reth_trie_common/
prefix_set.rs

1use crate::Nibbles;
2use alloc::{sync::Arc, vec::Vec};
3use alloy_primitives::map::{B256Map, B256Set};
4
5/// Collection of mutable prefix sets.
6#[derive(Clone, Default, Debug)]
7pub struct TriePrefixSetsMut {
8    /// A set of account prefixes that have changed.
9    pub account_prefix_set: PrefixSetMut,
10    /// A map containing storage changes with the hashed address as key and a set of storage key
11    /// prefixes as the value.
12    pub storage_prefix_sets: B256Map<PrefixSetMut>,
13    /// A set of hashed addresses of destroyed accounts.
14    pub destroyed_accounts: B256Set,
15}
16
17impl TriePrefixSetsMut {
18    /// Returns `true` if all prefix sets are empty.
19    pub fn is_empty(&self) -> bool {
20        self.account_prefix_set.is_empty() &&
21            self.storage_prefix_sets.is_empty() &&
22            self.destroyed_accounts.is_empty()
23    }
24
25    /// Extends prefix sets with contents of another prefix set.
26    pub fn extend(&mut self, other: Self) {
27        self.account_prefix_set.extend(other.account_prefix_set);
28        for (hashed_address, prefix_set) in other.storage_prefix_sets {
29            self.storage_prefix_sets.entry(hashed_address).or_default().extend(prefix_set);
30        }
31        self.destroyed_accounts.extend(other.destroyed_accounts);
32    }
33
34    /// Returns a `TriePrefixSets` with the same elements as these sets.
35    ///
36    /// If not yet sorted, the elements will be sorted and deduplicated.
37    pub fn freeze(self) -> TriePrefixSets {
38        TriePrefixSets {
39            account_prefix_set: self.account_prefix_set.freeze(),
40            storage_prefix_sets: self
41                .storage_prefix_sets
42                .into_iter()
43                .map(|(hashed_address, prefix_set)| (hashed_address, prefix_set.freeze()))
44                .collect(),
45            destroyed_accounts: self.destroyed_accounts,
46        }
47    }
48
49    /// Clears the prefix sets and destroyed accounts map.
50    pub fn clear(&mut self) {
51        self.destroyed_accounts.clear();
52        self.storage_prefix_sets.clear();
53        self.account_prefix_set.clear();
54    }
55}
56
57/// Collection of trie prefix sets.
58#[derive(Default, Debug)]
59pub struct TriePrefixSets {
60    /// A set of account prefixes that have changed.
61    pub account_prefix_set: PrefixSet,
62    /// A map containing storage changes with the hashed address as key and a set of storage key
63    /// prefixes as the value.
64    pub storage_prefix_sets: B256Map<PrefixSet>,
65    /// A set of hashed addresses of destroyed accounts.
66    pub destroyed_accounts: B256Set,
67}
68
69/// A container for efficiently storing and checking for the presence of key prefixes.
70///
71/// This data structure stores a set of `Nibbles` and provides methods to insert
72/// new elements and check whether any existing element has a given prefix.
73///
74/// Internally, this implementation uses a `Vec` and aims to act like a `BTreeSet` in being both
75/// sorted and deduplicated. It does this by keeping a `sorted` flag. The `sorted` flag represents
76/// whether or not the `Vec` is definitely sorted. When a new element is added, it is set to
77/// `false.`. The `Vec` is sorted and deduplicated when `sorted` is `true` and:
78///  * An element is being checked for inclusion (`contains`), or
79///  * The set is being converted into an immutable `PrefixSet` (`freeze`)
80///
81/// This means that a `PrefixSet` will always be sorted and deduplicated when constructed from a
82/// `PrefixSetMut`.
83///
84/// # Examples
85///
86/// ```
87/// use reth_trie_common::{prefix_set::PrefixSetMut, Nibbles};
88///
89/// let mut prefix_set_mut = PrefixSetMut::default();
90/// prefix_set_mut.insert(Nibbles::from_nibbles_unchecked(&[0xa, 0xb]));
91/// prefix_set_mut.insert(Nibbles::from_nibbles_unchecked(&[0xa, 0xb, 0xc]));
92/// let mut prefix_set = prefix_set_mut.freeze();
93/// assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([0xa, 0xb])));
94/// assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([0xa, 0xb, 0xc])));
95/// ```
96#[derive(PartialEq, Eq, Clone, Default, Debug)]
97pub struct PrefixSetMut {
98    /// Flag indicating that any entry should be considered changed.
99    /// If set, the keys will be discarded.
100    all: bool,
101    keys: Vec<Nibbles>,
102}
103
104impl<I> From<I> for PrefixSetMut
105where
106    I: IntoIterator<Item = Nibbles>,
107{
108    fn from(value: I) -> Self {
109        Self { all: false, keys: value.into_iter().collect() }
110    }
111}
112
113impl PrefixSetMut {
114    /// Create [`PrefixSetMut`] with pre-allocated capacity.
115    pub fn with_capacity(capacity: usize) -> Self {
116        Self { all: false, keys: Vec::with_capacity(capacity) }
117    }
118
119    /// Create [`PrefixSetMut`] that considers all key changed.
120    pub const fn all() -> Self {
121        Self { all: true, keys: Vec::new() }
122    }
123
124    /// Inserts the given `nibbles` into the set.
125    pub fn insert(&mut self, nibbles: Nibbles) {
126        self.keys.push(nibbles);
127    }
128
129    /// Extend prefix set with contents of another prefix set.
130    pub fn extend(&mut self, other: Self) {
131        self.all |= other.all;
132        self.keys.extend(other.keys);
133    }
134
135    /// Extend prefix set keys with contents of provided iterator.
136    pub fn extend_keys<I>(&mut self, keys: I)
137    where
138        I: IntoIterator<Item = Nibbles>,
139    {
140        self.keys.extend(keys);
141    }
142
143    /// Returns the number of elements in the set.
144    pub const fn len(&self) -> usize {
145        self.keys.len()
146    }
147
148    /// Returns `true` if the set is empty.
149    pub const fn is_empty(&self) -> bool {
150        self.keys.is_empty()
151    }
152
153    /// Clears the inner vec for reuse, setting `all` to `false`.
154    pub fn clear(&mut self) {
155        self.all = false;
156        self.keys.clear();
157    }
158
159    /// Returns a `PrefixSet` with the same elements as this set.
160    ///
161    /// If not yet sorted, the elements will be sorted and deduplicated.
162    pub fn freeze(mut self) -> PrefixSet {
163        if self.all {
164            PrefixSet { index: 0, all: true, keys: Arc::new(Vec::new()) }
165        } else {
166            self.keys.sort_unstable();
167            self.keys.dedup();
168            // We need to shrink in both the sorted and non-sorted cases because deduping may have
169            // occurred either on `freeze`, or during `contains`.
170            self.keys.shrink_to_fit();
171            PrefixSet { index: 0, all: false, keys: Arc::new(self.keys) }
172        }
173    }
174}
175
176/// A sorted prefix set that has an immutable _sorted_ list of unique keys.
177///
178/// See also [`PrefixSetMut::freeze`].
179#[derive(Debug, Default, Clone)]
180pub struct PrefixSet {
181    /// Flag indicating that any entry should be considered changed.
182    all: bool,
183    index: usize,
184    keys: Arc<Vec<Nibbles>>,
185}
186
187impl PrefixSet {
188    /// Returns `true` if any of the keys in the set has the given prefix
189    ///
190    /// # Note on Mutability
191    ///
192    /// This method requires `&mut self` (unlike typical `contains` methods) because it maintains an
193    /// internal position tracker (`self.index`) between calls. This enables significant performance
194    /// optimization for sequential lookups in sorted order, which is common during trie traversal.
195    ///
196    /// The `index` field allows subsequent searches to start where previous ones left off,
197    /// avoiding repeated full scans of the prefix array when keys are accessed in nearby ranges.
198    ///
199    /// This optimization was inspired by Silkworm's implementation and significantly improves
200    /// incremental state root calculation performance
201    /// ([see PR #2417](https://github.com/paradigmxyz/reth/pull/2417)).
202    #[inline]
203    pub fn contains(&mut self, prefix: &Nibbles) -> bool {
204        if self.all {
205            return true
206        }
207
208        while self.index > 0 && &self.keys[self.index] > prefix {
209            self.index -= 1;
210        }
211
212        for (idx, key) in self.keys[self.index..].iter().enumerate() {
213            if key.starts_with(prefix) {
214                self.index += idx;
215                return true
216            }
217
218            if key > prefix {
219                self.index += idx;
220                return false
221            }
222        }
223
224        false
225    }
226
227    /// Returns an iterator over reference to _all_ nibbles regardless of cursor position.
228    pub fn iter(&self) -> core::slice::Iter<'_, Nibbles> {
229        self.keys.iter()
230    }
231
232    /// Returns true if every entry should be considered changed.
233    pub const fn all(&self) -> bool {
234        self.all
235    }
236
237    /// Returns the number of elements in the set.
238    pub fn len(&self) -> usize {
239        self.keys.len()
240    }
241
242    /// Returns `true` if the set is empty.
243    pub fn is_empty(&self) -> bool {
244        self.keys.is_empty()
245    }
246}
247
248impl<'a> IntoIterator for &'a PrefixSet {
249    type Item = &'a Nibbles;
250    type IntoIter = core::slice::Iter<'a, Nibbles>;
251    fn into_iter(self) -> Self::IntoIter {
252        self.iter()
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259
260    #[test]
261    fn test_contains_with_multiple_inserts_and_duplicates() {
262        let mut prefix_set_mut = PrefixSetMut::default();
263        prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3]));
264        prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 4]));
265        prefix_set_mut.insert(Nibbles::from_nibbles([4, 5, 6]));
266        prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3])); // Duplicate
267
268        let mut prefix_set = prefix_set_mut.freeze();
269        assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([1, 2])));
270        assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([4, 5])));
271        assert!(!prefix_set.contains(&Nibbles::from_nibbles_unchecked([7, 8])));
272        assert_eq!(prefix_set.len(), 3); // Length should be 3 (excluding duplicate)
273    }
274
275    #[test]
276    fn test_freeze_shrinks_capacity() {
277        let mut prefix_set_mut = PrefixSetMut::default();
278        prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3]));
279        prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 4]));
280        prefix_set_mut.insert(Nibbles::from_nibbles([4, 5, 6]));
281        prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3])); // Duplicate
282
283        assert_eq!(prefix_set_mut.keys.len(), 4); // Length should be 3 (including duplicate)
284        assert_eq!(prefix_set_mut.keys.capacity(), 4); // Capacity should be 4 (including duplicate)
285
286        let mut prefix_set = prefix_set_mut.freeze();
287        assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([1, 2])));
288        assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([4, 5])));
289        assert!(!prefix_set.contains(&Nibbles::from_nibbles_unchecked([7, 8])));
290        assert_eq!(prefix_set.keys.len(), 3); // Length should be 3 (excluding duplicate)
291        assert_eq!(prefix_set.keys.capacity(), 3); // Capacity should be 3 after shrinking
292    }
293
294    #[test]
295    fn test_freeze_shrinks_existing_capacity() {
296        // do the above test but with preallocated capacity
297        let mut prefix_set_mut = PrefixSetMut::with_capacity(101);
298        prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3]));
299        prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 4]));
300        prefix_set_mut.insert(Nibbles::from_nibbles([4, 5, 6]));
301        prefix_set_mut.insert(Nibbles::from_nibbles([1, 2, 3])); // Duplicate
302
303        assert_eq!(prefix_set_mut.keys.len(), 4); // Length should be 3 (including duplicate)
304        assert_eq!(prefix_set_mut.keys.capacity(), 101); // Capacity should be 101 (including duplicate)
305
306        let mut prefix_set = prefix_set_mut.freeze();
307        assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([1, 2])));
308        assert!(prefix_set.contains(&Nibbles::from_nibbles_unchecked([4, 5])));
309        assert!(!prefix_set.contains(&Nibbles::from_nibbles_unchecked([7, 8])));
310        assert_eq!(prefix_set.keys.len(), 3); // Length should be 3 (excluding duplicate)
311        assert_eq!(prefix_set.keys.capacity(), 3); // Capacity should be 3 after shrinking
312    }
313
314    #[test]
315    fn test_prefix_set_all_extend() {
316        let mut prefix_set_mut = PrefixSetMut::default();
317        prefix_set_mut.extend(PrefixSetMut::all());
318        assert!(prefix_set_mut.all);
319    }
320}