reth_provider/bundle_state/
state_reverts.rs

1use alloy_primitives::B256;
2use revm::{db::states::RevertToSlot, primitives::FlaggedStorage};
3use std::iter::Peekable;
4
5/// Iterator over storage reverts.
6/// See [`StorageRevertsIter::next`] for more details.
7#[allow(missing_debug_implementations)]
8pub struct StorageRevertsIter<R: Iterator, W: Iterator> {
9    reverts: Peekable<R>,
10    wiped: Peekable<W>,
11}
12
13impl<R, W> StorageRevertsIter<R, W>
14where
15    R: Iterator<Item = (B256, RevertToSlot)>,
16    W: Iterator<Item = (B256, FlaggedStorage)>,
17{
18    /// Create a new iterator over storage reverts.
19    pub fn new(
20        reverts: impl IntoIterator<IntoIter = R>,
21        wiped: impl IntoIterator<IntoIter = W>,
22    ) -> Self {
23        Self { reverts: reverts.into_iter().peekable(), wiped: wiped.into_iter().peekable() }
24    }
25
26    /// Consume next revert and return it.
27    fn next_revert(&mut self) -> Option<(B256, FlaggedStorage)> {
28        self.reverts.next().map(|(key, revert)| (key, revert.to_previous_value()))
29    }
30
31    /// Consume next wiped storage and return it.
32    fn next_wiped(&mut self) -> Option<(B256, FlaggedStorage)> {
33        self.wiped.next()
34    }
35}
36
37impl<R, W> Iterator for StorageRevertsIter<R, W>
38where
39    R: Iterator<Item = (B256, RevertToSlot)>,
40    W: Iterator<Item = (B256, FlaggedStorage)>,
41{
42    type Item = (B256, FlaggedStorage);
43
44    /// Iterate over storage reverts and wiped entries and return items in the sorted order.
45    /// NOTE: The implementation assumes that inner iterators are already sorted.
46    fn next(&mut self) -> Option<Self::Item> {
47        match (self.reverts.peek(), self.wiped.peek()) {
48            (Some(revert), Some(wiped)) => {
49                // Compare the keys and return the lesser.
50                use std::cmp::Ordering;
51                match revert.0.cmp(&wiped.0) {
52                    Ordering::Less => self.next_revert(),
53                    Ordering::Greater => self.next_wiped(),
54                    Ordering::Equal => {
55                        // Keys are the same, decide which one to return.
56                        let (key, revert_to) = *revert;
57
58                        let value = match revert_to {
59                            // If the slot is some, prefer the revert value.
60                            RevertToSlot::Some(value) => value,
61                            // If the slot was destroyed, prefer the database value.
62                            RevertToSlot::Destroyed => wiped.1,
63                        };
64
65                        // Consume both values from inner iterators.
66                        self.next_revert();
67                        self.next_wiped();
68
69                        Some((key, value))
70                    }
71                }
72            }
73            (Some(_revert), None) => self.next_revert(),
74            (None, Some(_wiped)) => self.next_wiped(),
75            (None, None) => None,
76        }
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    #[test]
85    fn test_storage_reverts_iter_empty() {
86        // Create empty sample data for reverts and wiped entries.
87        let reverts: Vec<(B256, RevertToSlot)> = vec![];
88        let wiped: Vec<(B256, FlaggedStorage)> = vec![];
89
90        // Create the iterator with the empty data.
91        let iter = StorageRevertsIter::new(reverts, wiped);
92
93        // Iterate and collect results into a vector for verification.
94        let results: Vec<_> = iter.collect();
95
96        // Verify that the results are empty.
97        assert_eq!(results, vec![]);
98    }
99
100    #[test]
101    fn test_storage_reverts_iter_reverts_only() {
102        // Create sample data for only reverts.
103        let reverts = vec![
104            (B256::from_slice(&[4; 32]), RevertToSlot::Destroyed),
105            (B256::from_slice(&[5; 32]), RevertToSlot::Some(FlaggedStorage::new_from_value(40))),
106        ];
107
108        // Create the iterator with only reverts and no wiped entries.
109        let iter = StorageRevertsIter::new(reverts, vec![]);
110
111        // Iterate and collect results into a vector for verification.
112        let results: Vec<_> = iter.collect();
113
114        // Verify the output order and values.
115        assert_eq!(
116            results,
117            vec![
118                (B256::from_slice(&[4; 32]), FlaggedStorage::ZERO), // Revert slot previous value
119                (B256::from_slice(&[5; 32]), FlaggedStorage::new_from_value(40)), /* Only revert
120                                                                     * present. */
121            ]
122        );
123    }
124
125    #[test]
126    fn test_storage_reverts_iter_wiped_only() {
127        // Create sample data for only wiped entries.
128        let wiped = vec![
129            (B256::from_slice(&[6; 32]), FlaggedStorage::new_from_value(50)),
130            (B256::from_slice(&[7; 32]), FlaggedStorage::new_from_value(60)),
131        ];
132
133        // Create the iterator with only wiped entries and no reverts.
134        let iter = StorageRevertsIter::new(vec![], wiped);
135
136        // Iterate and collect results into a vector for verification.
137        let results: Vec<_> = iter.collect();
138
139        // Verify the output order and values.
140        assert_eq!(
141            results,
142            vec![
143                (B256::from_slice(&[6; 32]), FlaggedStorage::new_from_value(50)), /* Only wiped
144                                                                                   * present. */
145                (B256::from_slice(&[7; 32]), FlaggedStorage::new_from_value(60)), /* Only wiped
146                                                                                   * present. */
147            ]
148        );
149    }
150
151    #[test]
152    fn test_storage_reverts_iter_interleaved() {
153        // Create sample data for interleaved reverts and wiped entries.
154        let reverts = vec![
155            (B256::from_slice(&[8; 32]), RevertToSlot::Some(FlaggedStorage::new_from_value(70))),
156            (B256::from_slice(&[9; 32]), RevertToSlot::Some(FlaggedStorage::new_from_value(80))),
157            // Some higher key than wiped
158            (B256::from_slice(&[15; 32]), RevertToSlot::Some(FlaggedStorage::new_from_value(90))),
159        ];
160
161        let wiped = vec![
162            (B256::from_slice(&[8; 32]), FlaggedStorage::new_from_value(75)), // Same key as revert
163            (B256::from_slice(&[10; 32]), FlaggedStorage::new_from_value(85)), // Wiped with new key
164        ];
165
166        // Create the iterator with the sample data.
167        let iter = StorageRevertsIter::new(reverts, wiped);
168
169        // Iterate and collect results into a vector for verification.
170        let results: Vec<_> = iter.collect();
171
172        // Verify the output order and values.
173        assert_eq!(
174            results,
175            vec![
176                (B256::from_slice(&[8; 32]), FlaggedStorage::new_from_value(70)), /* Revert takes priority. */
177                (B256::from_slice(&[9; 32]), FlaggedStorage::new_from_value(80)), /* Only revert
178                                                                                   * present. */
179                (B256::from_slice(&[10; 32]), FlaggedStorage::new_from_value(85)), // Wiped entry.
180                (B256::from_slice(&[15; 32]), FlaggedStorage::new_from_value(90)), /* WGreater revert entry */
181            ]
182        );
183    }
184}