reth_trie_common/
prefix_set.rs

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

reth_trie_common/
prefix_set.rs
