1use crate::Nibbles;
2use alloc::{sync::Arc, vec::Vec};
3use alloy_primitives::map::{B256Map, B256Set};
4
5#[derive(Clone, Default, Debug)]
7pub struct TriePrefixSetsMut {
8 pub account_prefix_set: PrefixSetMut,
10 pub storage_prefix_sets: B256Map<PrefixSetMut>,
13 pub destroyed_accounts: B256Set,
15}
16
17impl TriePrefixSetsMut {
18 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 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 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 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#[derive(Default, Debug)]
59pub struct TriePrefixSets {
60 pub account_prefix_set: PrefixSet,
62 pub storage_prefix_sets: B256Map<PrefixSet>,
65 pub destroyed_accounts: B256Set,
67}
68
69#[derive(PartialEq, Eq, Clone, Default, Debug)]
97pub struct PrefixSetMut {
98 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 pub fn with_capacity(capacity: usize) -> Self {
116 Self { all: false, keys: Vec::with_capacity(capacity) }
117 }
118
119 pub const fn all() -> Self {
121 Self { all: true, keys: Vec::new() }
122 }
123
124 pub fn insert(&mut self, nibbles: Nibbles) {
126 self.keys.push(nibbles);
127 }
128
129 pub fn extend(&mut self, other: Self) {
131 self.all |= other.all;
132 self.keys.extend(other.keys);
133 }
134
135 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 pub const fn len(&self) -> usize {
145 self.keys.len()
146 }
147
148 pub const fn is_empty(&self) -> bool {
150 self.keys.is_empty()
151 }
152
153 pub fn clear(&mut self) {
155 self.all = false;
156 self.keys.clear();
157 }
158
159 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 self.keys.shrink_to_fit();
171 PrefixSet { index: 0, all: false, keys: Arc::new(self.keys) }
172 }
173 }
174}
175
176#[derive(Debug, Default, Clone)]
180pub struct PrefixSet {
181 all: bool,
183 index: usize,
184 keys: Arc<Vec<Nibbles>>,
185}
186
187impl PrefixSet {
188 #[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 pub fn iter(&self) -> core::slice::Iter<'_, Nibbles> {
229 self.keys.iter()
230 }
231
232 pub const fn all(&self) -> bool {
234 self.all
235 }
236
237 pub fn len(&self) -> usize {
239 self.keys.len()
240 }
241
242 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])); 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); }
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])); assert_eq!(prefix_set_mut.keys.len(), 4); assert_eq!(prefix_set_mut.keys.capacity(), 4); 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); assert_eq!(prefix_set.keys.capacity(), 3); }
293
294 #[test]
295 fn test_freeze_shrinks_existing_capacity() {
296 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])); assert_eq!(prefix_set_mut.keys.len(), 4); assert_eq!(prefix_set_mut.keys.capacity(), 101); 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); assert_eq!(prefix_set.keys.capacity(), 3); }
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}