1#[cfg(feature = "trie-debug")]
2use crate::debug_recorder::TrieDebugRecorder;
3use crate::{
4 lfu::BucketedLfu, traits::SparseTrie as SparseTrieTrait, ArenaParallelSparseTrie,
5 RevealableSparseTrie,
6};
7use alloc::vec::Vec;
8use alloy_primitives::{map::B256Map, B256};
9use either::Either;
10use reth_execution_errors::{SparseStateTrieResult, SparseTrieErrorKind};
11use reth_trie_common::{
12 prefix_set::{PrefixSet, PrefixSetMut, TriePrefixSets, TriePrefixSetsMut},
13 updates::{StorageTrieUpdates, TrieUpdates},
14 DecodedMultiProof, MultiProof, Nibbles, ProofTrieNodeV2,
15};
16#[cfg(feature = "std")]
17use tracing::debug;
18use tracing::instrument;
19
20#[derive(Debug, Default)]
25pub struct DeferredDrops {
26 pub proof_nodes_bufs: Vec<Vec<ProofTrieNodeV2>>,
28}
29
30#[derive(Debug)]
31pub struct SparseStateTrie<
33 A = ArenaParallelSparseTrie, S = ArenaParallelSparseTrie, > {
36 state: RevealableSparseTrie<A>,
38 storage: StorageTries<S>,
40 retain_updates: bool,
42 retain_changed_paths: bool,
44 deferred_drops: DeferredDrops,
46 hot_slots_lfu: BucketedLfu<HotSlotKey>,
48 hot_accounts_lfu: BucketedLfu<B256>,
50 #[cfg(feature = "metrics")]
52 metrics: crate::metrics::SparseStateTrieMetrics,
53}
54
55impl<A, S> Default for SparseStateTrie<A, S>
56where
57 A: Default,
58 S: Default,
59{
60 fn default() -> Self {
61 Self {
62 state: Default::default(),
63 storage: Default::default(),
64 retain_updates: false,
65 retain_changed_paths: false,
66 deferred_drops: DeferredDrops::default(),
67 hot_slots_lfu: BucketedLfu::default(),
68 hot_accounts_lfu: BucketedLfu::default(),
69 #[cfg(feature = "metrics")]
70 metrics: Default::default(),
71 }
72 }
73}
74
75#[cfg(test)]
76impl SparseStateTrie {
77 pub fn from_state(state: RevealableSparseTrie) -> Self {
79 Self { state, ..Default::default() }
80 }
81}
82
83impl<A, S> SparseStateTrie<A, S> {
84 pub const fn set_updates(&mut self, retain_updates: bool) {
86 self.retain_updates = retain_updates;
87 }
88
89 pub const fn with_updates(mut self, retain_updates: bool) -> Self {
91 self.set_updates(retain_updates);
92 self
93 }
94
95 pub fn set_hot_cache_capacities(&mut self, max_hot_slots: usize, max_hot_accounts: usize) {
100 self.hot_slots_lfu.decay_and_evict(max_hot_slots);
101 self.hot_accounts_lfu.decay_and_evict(max_hot_accounts);
102 }
103
104 pub fn with_hot_cache_capacities(
106 mut self,
107 max_hot_slots: usize,
108 max_hot_accounts: usize,
109 ) -> Self {
110 self.set_hot_cache_capacities(max_hot_slots, max_hot_accounts);
111 self
112 }
113
114 pub fn set_accounts_trie(&mut self, trie: RevealableSparseTrie<A>) {
116 self.state = trie;
117 }
118
119 pub fn with_accounts_trie(mut self, trie: RevealableSparseTrie<A>) -> Self {
121 self.set_accounts_trie(trie);
122 self
123 }
124
125 pub fn set_default_storage_trie(&mut self, trie: RevealableSparseTrie<S>) {
128 self.storage.default_trie = trie;
129 }
130
131 pub fn with_default_storage_trie(mut self, trie: RevealableSparseTrie<S>) -> Self {
134 self.set_default_storage_trie(trie);
135 self
136 }
137
138 pub fn take_deferred_drops(&mut self) -> DeferredDrops {
143 core::mem::take(&mut self.deferred_drops)
144 }
145}
146
147impl SparseStateTrie {
148 pub fn new() -> Self {
150 Self::default()
151 }
152}
153
154impl<A: SparseTrieTrait, S: SparseTrieTrait> SparseStateTrie<A, S> {
155 pub fn set_changed_paths(&mut self, retain_changed_paths: bool) {
157 self.retain_changed_paths = retain_changed_paths;
158 self.state.set_changed_paths(retain_changed_paths);
159 for trie in self.storage.tries.values_mut() {
160 trie.set_changed_paths(retain_changed_paths);
161 }
162 for trie in &mut self.storage.cleared_tries {
163 trie.set_changed_paths(retain_changed_paths);
164 }
165 self.storage.default_trie.set_changed_paths(retain_changed_paths);
166 }
167
168 pub fn with_changed_paths(mut self, retain_changed_paths: bool) -> Self {
170 self.set_changed_paths(retain_changed_paths);
171 self
172 }
173
174 fn storage_trie_changed_paths(&mut self) -> B256Map<PrefixSetMut> {
176 self.storage
177 .tries
178 .iter_mut()
179 .filter_map(|(address, trie)| {
180 let changed_paths = trie.take_changed_paths()?;
181 (!changed_paths.is_empty()).then_some((*address, changed_paths))
182 })
183 .collect()
184 }
185
186 pub fn take_changed_paths(&mut self) -> Option<TriePrefixSetsMut> {
190 let storage_prefix_sets = self.storage_trie_changed_paths();
191 self.state.take_changed_paths().map(|account_prefix_set| TriePrefixSetsMut {
192 account_prefix_set,
193 storage_prefix_sets,
194 destroyed_accounts: Default::default(),
195 })
196 }
197
198 #[cfg(feature = "trie-debug")]
203 pub fn take_debug_recorders(&mut self) -> alloc::vec::Vec<(Option<B256>, TrieDebugRecorder)> {
204 let mut recorders = alloc::vec::Vec::new();
205 if let Some(trie) = self.state.as_revealed_mut() {
206 recorders.push((None, trie.take_debug_recorder()));
207 }
208 for (address, trie) in &mut self.storage.tries {
209 if let Some(trie) = trie.as_revealed_mut() {
210 recorders.push((Some(*address), trie.take_debug_recorder()));
211 }
212 }
213 recorders
214 }
215}
216
217impl<A, S> SparseStateTrie<A, S>
218where
219 A: SparseTrieTrait + Default,
220 S: SparseTrieTrait + Default + Clone,
221{
222 pub const fn trie_mut(&mut self) -> &mut RevealableSparseTrie<A> {
224 &mut self.state
225 }
226
227 pub fn is_account_revealed(&self, account: B256) -> bool {
229 let path = Nibbles::unpack(account);
230 let trie = match self.state_trie_ref() {
231 Some(t) => t,
232 None => return false,
233 };
234
235 trie.find_leaf(&path, None).is_ok()
236 }
237
238 pub fn check_valid_storage_witness(&self, address: B256, slot: B256) -> bool {
240 let path = Nibbles::unpack(slot);
241 let trie = match self.storage_trie_ref(&address) {
242 Some(t) => t,
243 None => return false,
244 };
245
246 trie.find_leaf(&path, None).is_ok()
247 }
248
249 #[inline]
251 pub fn record_slot_touch(&mut self, account: B256, slot: B256) {
252 self.hot_slots_lfu.touch(HotSlotKey { address: account, slot });
253 }
254
255 #[inline]
257 pub fn record_account_touch(&mut self, account: B256) {
258 self.hot_accounts_lfu.touch(account);
259 }
260
261 pub fn get_account_value(&self, account: &B256) -> Option<&Vec<u8>> {
263 self.state.as_revealed_ref()?.get_leaf_value(&Nibbles::unpack(account))
264 }
265
266 pub fn get_storage_slot_value(&self, account: &B256, slot: &B256) -> Option<&Vec<u8>> {
268 self.storage.tries.get(account)?.as_revealed_ref()?.get_leaf_value(&Nibbles::unpack(slot))
269 }
270
271 pub const fn state_trie_ref(&self) -> Option<&A> {
273 self.state.as_revealed_ref()
274 }
275
276 pub fn storage_trie_ref(&self, address: &B256) -> Option<&S> {
278 self.storage.tries.get(address).and_then(|e| e.as_revealed_ref())
279 }
280
281 pub fn storage_trie_mut(&mut self, address: &B256) -> Option<&mut S> {
283 self.storage.tries.get_mut(address).and_then(|e| e.as_revealed_mut())
284 }
285
286 pub const fn storage_tries_mut(&mut self) -> &mut B256Map<RevealableSparseTrie<S>> {
288 &mut self.storage.tries
289 }
290
291 pub fn take_storage_trie(&mut self, address: &B256) -> Option<RevealableSparseTrie<S>> {
293 self.storage.tries.remove(address)
294 }
295
296 pub fn take_or_create_storage_trie(&mut self, address: &B256) -> RevealableSparseTrie<S> {
298 self.storage.tries.remove(address).unwrap_or_else(|| {
299 self.storage.cleared_tries.pop().unwrap_or_else(|| self.storage.default_trie.clone())
300 })
301 }
302
303 pub fn insert_storage_trie(&mut self, address: B256, storage_trie: RevealableSparseTrie<S>) {
305 self.storage.tries.insert(address, storage_trie);
306 }
307
308 pub fn get_or_create_storage_trie_mut(
310 &mut self,
311 address: B256,
312 ) -> &mut RevealableSparseTrie<S> {
313 self.storage.get_or_create_trie_mut(address)
314 }
315
316 pub fn reveal_multiproof(&mut self, multiproof: MultiProof) -> SparseStateTrieResult<()> {
319 let decoded_multiproof = multiproof.try_into()?;
321
322 self.reveal_decoded_multiproof(decoded_multiproof)
324 }
325
326 #[instrument(level = "debug", target = "trie::sparse", skip_all)]
329 pub fn reveal_decoded_multiproof(
330 &mut self,
331 multiproof: DecodedMultiProof,
332 ) -> SparseStateTrieResult<()> {
333 self.reveal_decoded_multiproof_v2(multiproof.into())
334 }
335
336 #[instrument(level = "debug", target = "trie::sparse", skip_all)]
341 pub fn reveal_decoded_multiproof_v2(
342 &mut self,
343 multiproof: reth_trie_common::DecodedMultiProofV2,
344 ) -> SparseStateTrieResult<()> {
345 let reth_trie_common::DecodedMultiProofV2 { account_proofs, mut storage_proofs, .. } =
346 multiproof;
347
348 let mut targets = Vec::with_capacity(storage_proofs.len() + 1);
351
352 if !account_proofs.is_empty() {
353 #[cfg(feature = "metrics")]
354 self.metrics.increment_total_account_nodes(account_proofs.len() as u64);
355 targets.push((None, Either::Left(&mut self.state), account_proofs));
356 }
357
358 for &account in storage_proofs.keys() {
360 let _ = self.storage.get_or_create_trie_mut(account);
361 }
362
363 for (account, trie) in &mut self.storage.tries {
364 if let Some(nodes) = storage_proofs.remove(account) {
365 #[cfg(feature = "metrics")]
366 self.metrics.increment_total_storage_nodes(nodes.len() as u64);
367 targets.push((Some(*account), Either::Right(trie), nodes));
368 }
369 }
370
371 let retain_updates = self.retain_updates;
372 let retain_changed_paths = self.retain_changed_paths;
373
374 #[cfg(not(feature = "std"))]
375 let results: Vec<_> = targets
376 .into_iter()
377 .map(|(_, target, mut nodes)| {
378 let result = match target {
379 Either::Left(trie) => {
380 trie.reveal_v2_proof_nodes(&mut nodes, retain_updates, retain_changed_paths)
381 }
382 Either::Right(trie) => {
383 trie.reveal_v2_proof_nodes(&mut nodes, retain_updates, retain_changed_paths)
384 }
385 };
386 (result, nodes)
387 })
388 .collect();
389
390 #[cfg(feature = "std")]
391 let results: Vec<_> = {
392 use rayon::iter::ParallelIterator;
393 use reth_primitives_traits::ParallelBridgeBuffered;
394
395 let parent_span = tracing::Span::current();
396 targets
397 .into_iter()
398 .par_bridge_buffered()
399 .map(|(hashed_address, target, mut nodes)| {
400 let _span = tracing::trace_span!(
401 target: "trie::sparse",
402 parent: &parent_span,
403 "reveal_v2_proof_nodes",
404 ?hashed_address,
405 )
406 .entered();
407
408 let result = match target {
409 Either::Left(trie) => trie.reveal_v2_proof_nodes(
410 &mut nodes,
411 retain_updates,
412 retain_changed_paths,
413 ),
414 Either::Right(trie) => trie.reveal_v2_proof_nodes(
415 &mut nodes,
416 retain_updates,
417 retain_changed_paths,
418 ),
419 };
420 (result, nodes)
421 })
422 .collect()
423 };
424
425 let mut any_err = Ok(());
427 for (result, nodes) in results {
428 if result.is_err() && any_err.is_ok() {
429 any_err = result.map_err(Into::into);
430 }
431 self.deferred_drops.proof_nodes_bufs.push(nodes);
432 }
433
434 any_err
435 }
436
437 pub fn wipe_storage(&mut self, address: B256) -> SparseStateTrieResult<()> {
439 if let Some(trie) = self.storage.tries.get_mut(&address) {
440 trie.wipe()?;
441 }
442 Ok(())
443 }
444
445 #[instrument(level = "debug", target = "trie::sparse", skip_all)]
449 pub fn calculate_subtries(&mut self) {
450 if let RevealableSparseTrie::Revealed(trie) = &mut self.state {
451 trie.update_subtrie_hashes();
452 }
453 }
454
455 pub fn storage_root(&mut self, account: &B256) -> Option<B256> {
457 self.storage.tries.get_mut(account).and_then(|trie| trie.root())
458 }
459
460 fn revealed_trie_mut(&mut self) -> SparseStateTrieResult<&mut A> {
462 self.state.as_revealed_mut().ok_or_else(|| SparseTrieErrorKind::Blind.into())
463 }
464
465 pub fn root(&mut self) -> SparseStateTrieResult<B256> {
467 #[cfg(feature = "metrics")]
469 self.metrics.record();
470
471 Ok(self.revealed_trie_mut()?.root())
472 }
473
474 #[instrument(level = "debug", target = "trie::sparse", skip_all)]
478 pub fn root_with_updates(&mut self) -> SparseStateTrieResult<(B256, TrieUpdates)> {
479 #[cfg(feature = "metrics")]
481 self.metrics.record();
482
483 let storage_tries = self.storage_trie_updates();
484 let revealed = self.revealed_trie_mut()?;
485
486 let (root, updates) = (revealed.root(), revealed.take_updates());
487 let updates = TrieUpdates {
488 account_nodes: updates.updated_nodes,
489 removed_nodes: updates.removed_nodes,
490 storage_tries,
491 };
492 Ok((root, updates))
493 }
494
495 pub fn storage_trie_updates(&mut self) -> B256Map<StorageTrieUpdates> {
499 self.storage
500 .tries
501 .iter_mut()
502 .map(|(address, trie)| {
503 let trie = trie.as_revealed_mut().unwrap();
504 let updates = trie.take_updates();
505 let updates = StorageTrieUpdates {
506 is_deleted: updates.wiped,
507 storage_nodes: updates.updated_nodes,
508 removed_nodes: updates.removed_nodes,
509 };
510 (*address, updates)
511 })
512 .filter(|(_, updates)| !updates.is_empty())
513 .collect()
514 }
515
516 pub fn take_trie_updates(&mut self) -> Option<TrieUpdates> {
520 let storage_tries = self.storage_trie_updates();
521 self.state.as_revealed_mut().map(|state| {
522 let updates = state.take_updates();
523 TrieUpdates {
524 account_nodes: updates.updated_nodes,
525 removed_nodes: updates.removed_nodes,
526 storage_tries,
527 }
528 })
529 }
530}
531
532impl<A, S> SparseStateTrie<A, S>
533where
534 A: SparseTrieTrait + Default,
535 S: SparseTrieTrait + Default + Clone,
536{
537 pub fn clear(&mut self) {
542 self.state.clear();
543 self.storage.clear();
544 }
545
546 pub fn memory_size(&self) -> usize {
551 let mut size = core::mem::size_of::<Self>();
552
553 size += match &self.state {
554 RevealableSparseTrie::Revealed(t) | RevealableSparseTrie::Blind(Some(t)) => {
555 t.memory_size()
556 }
557 RevealableSparseTrie::Blind(None) => 0,
558 };
559
560 for trie in self.storage.tries.values() {
561 size += match trie {
562 RevealableSparseTrie::Revealed(t) | RevealableSparseTrie::Blind(Some(t)) => {
563 t.memory_size()
564 }
565 RevealableSparseTrie::Blind(None) => 0,
566 };
567 }
568 for trie in &self.storage.cleared_tries {
569 size += match trie {
570 RevealableSparseTrie::Revealed(t) | RevealableSparseTrie::Blind(Some(t)) => {
571 t.memory_size()
572 }
573 RevealableSparseTrie::Blind(None) => 0,
574 };
575 }
576
577 size
578 }
579
580 pub fn retained_storage_tries_count(&self) -> usize {
582 self.storage.tries.len() + self.storage.cleared_tries.len()
583 }
584
585 #[cfg(feature = "std")]
601 #[instrument(
602 level = "debug",
603 name = "SparseStateTrie::prune",
604 target = "trie::sparse",
605 skip_all,
606 fields(%max_hot_slots, %max_hot_accounts)
607 )]
608 pub fn prune(
609 &mut self,
610 max_hot_slots: usize,
611 max_hot_accounts: usize,
612 mut retained_paths: TriePrefixSetsMut,
613 ) {
614 self.hot_slots_lfu.decay_and_evict(max_hot_slots);
615 self.hot_accounts_lfu.decay_and_evict(max_hot_accounts);
616 extend_retained_paths_from_lfus(
617 &mut retained_paths,
618 &self.hot_accounts_lfu,
619 &self.hot_slots_lfu,
620 );
621 let retained_paths = retained_paths.freeze();
622
623 let retained_accounts = retained_paths.account_prefix_set.len();
624 let retained_storage_tries = retained_paths.storage_prefix_sets.len();
625 let total_storage_tries_before = self.storage.tries.len();
626
627 let TriePrefixSets { account_prefix_set, storage_prefix_sets, .. } = retained_paths;
628
629 let parent_span = tracing::Span::current();
630 let account_parent_span = parent_span.clone();
631
632 let (account_nodes_pruned, storage_tries_evicted) = rayon::join(
634 || {
635 let hashed_address = Option::<B256>::None;
636 let _span = tracing::trace_span!(
637 target: "trie::sparse",
638 parent: &account_parent_span,
639 "prune_trie",
640 ?hashed_address,
641 )
642 .entered();
643
644 self.state
645 .as_revealed_mut()
646 .map(|trie| trie.prune(account_prefix_set.slice()))
647 .unwrap_or(0)
648 },
649 || self.storage.prune_by_retained_slots(storage_prefix_sets, &parent_span),
650 );
651
652 debug!(
653 target: "trie::sparse",
654 retained_accounts,
655 retained_storage_tries,
656 account_nodes_pruned,
657 storage_tries_evicted,
658 storage_tries_after = total_storage_tries_before - storage_tries_evicted,
659 "SparseStateTrie::prune completed"
660 );
661 }
662}
663
664#[cfg(feature = "std")]
665fn extend_retained_paths_from_lfus(
666 retained_paths: &mut TriePrefixSetsMut,
667 hot_accounts: &BucketedLfu<B256>,
668 hot_slots: &BucketedLfu<HotSlotKey>,
669) {
670 retained_paths
671 .account_prefix_set
672 .extend_keys(hot_accounts.keys().map(|key| Nibbles::unpack(*key)));
673 for key in hot_slots.keys() {
674 retained_paths
675 .storage_prefix_sets
676 .entry(key.address)
677 .or_default()
678 .insert(Nibbles::unpack(key.slot));
679 }
680}
681
682#[derive(Debug, Default)]
685struct StorageTries<S = ArenaParallelSparseTrie> {
686 tries: B256Map<RevealableSparseTrie<S>>,
688 cleared_tries: Vec<RevealableSparseTrie<S>>,
690 default_trie: RevealableSparseTrie<S>,
692}
693
694#[cfg(feature = "std")]
695impl<S: SparseTrieTrait> StorageTries<S> {
696 fn prune_by_retained_slots(
701 &mut self,
702 retained_slots: B256Map<PrefixSet>,
703 parent_span: &tracing::Span,
704 ) -> usize {
705 {
707 use rayon::iter::{IntoParallelRefMutIterator, ParallelIterator};
708 self.tries.par_iter_mut().for_each(|(address, trie)| {
709 let hashed_address = Some(*address);
710 let _span = tracing::trace_span!(
711 target: "trie::sparse",
712 parent: parent_span,
713 "prune_trie",
714 ?hashed_address,
715 )
716 .entered();
717
718 if let Some(slots) = retained_slots.get(address) {
719 if let Some(t) = trie.as_revealed_mut() {
720 t.prune(slots.slice());
721 }
722 } else {
723 trie.clear();
724 }
725 });
726 }
727
728 let addresses_to_evict: Vec<B256> = self
730 .tries
731 .keys()
732 .filter(|address| !retained_slots.contains_key(*address))
733 .copied()
734 .collect();
735
736 let evicted = addresses_to_evict.len();
737 self.cleared_tries.reserve(evicted);
738 for address in &addresses_to_evict {
739 if let Some(trie) = self.tries.remove(address) {
740 self.cleared_tries.push(trie);
741 }
742 }
743
744 evicted
745 }
746}
747
748impl<S: SparseTrieTrait> StorageTries<S> {
749 fn clear(&mut self) {
752 self.cleared_tries.extend(self.tries.drain().map(|(_, mut trie)| {
753 trie.clear();
754 trie
755 }));
756 }
757}
758
759impl<S: SparseTrieTrait + Clone> StorageTries<S> {
760 fn get_or_create_trie_mut(&mut self, address: B256) -> &mut RevealableSparseTrie<S> {
762 self.tries.entry(address).or_insert_with(|| {
763 self.cleared_tries.pop().unwrap_or_else(|| self.default_trie.clone())
764 })
765 }
766}
767
768#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
770struct HotSlotKey {
771 address: B256,
772 slot: B256,
773}
774
775#[cfg(test)]
776mod tests {
777 use super::*;
778 use crate::{ArenaParallelSparseTrie, LeafLookup, LeafUpdate};
779 use alloy_primitives::{
780 b256,
781 map::{HashMap, HashSet},
782 U256,
783 };
784 use arbitrary::Arbitrary;
785 use rand::{rngs::StdRng, Rng, SeedableRng};
786 use reth_execution_errors::{SparseStateTrieErrorKind, SparseTrieErrorKind};
787 use reth_primitives_traits::Account;
788 use reth_trie::{updates::StorageTrieUpdates, HashBuilder, MultiProof, EMPTY_ROOT_HASH};
789 use reth_trie_common::{
790 proof::{ProofNodes, ProofRetainer},
791 BranchNodeMasks, BranchNodeMasksMap, BranchNodeV2, LeafNode, RlpNode, StorageMultiProof,
792 TrieAccount, TrieMask, TrieNodeV2,
793 };
794
795 fn leaf_key(suffix: impl AsRef<[u8]>, total_len: usize) -> Nibbles {
797 let suffix = suffix.as_ref();
798 let mut nibbles = Nibbles::from_nibbles(suffix);
799 nibbles.extend(&Nibbles::from_nibbles_unchecked(vec![0; total_len - suffix.len()]));
800 nibbles
801 }
802
803 fn apply_account_update(sparse: &mut SparseStateTrie, address: B256, update: LeafUpdate) {
804 let mut updates = B256Map::from_iter([(address, update)]);
805 sparse.trie_mut().update_leaves(&mut updates, |_, _| {}).unwrap();
806 assert!(updates.is_empty());
807 }
808
809 #[test]
810 fn reveal_account_path_twice() {
811 let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
812
813 let full_path_0 = leaf_key([0x0], 64);
815 let _full_path_1 = leaf_key([0x1], 64);
816
817 let leaf_value = alloy_rlp::encode(TrieAccount::default());
818 let leaf_1 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(
820 leaf_key([], 63),
821 leaf_value.clone(),
822 )));
823 let leaf_2 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(
824 leaf_key([], 63),
825 leaf_value.clone(),
826 )));
827
828 let multiproof = MultiProof {
829 account_subtree: ProofNodes::from_iter([
830 (
831 Nibbles::default(),
832 alloy_rlp::encode(TrieNodeV2::Branch(BranchNodeV2 {
833 key: Nibbles::default(),
834 stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)],
835 state_mask: TrieMask::new(0b11),
836 branch_rlp_node: None,
837 }))
838 .into(),
839 ),
840 (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()),
841 (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
842 ]),
843 ..Default::default()
844 };
845
846 sparse.reveal_decoded_multiproof(multiproof.try_into().unwrap()).unwrap();
848 assert!(matches!(
849 sparse.state_trie_ref().unwrap().find_leaf(&full_path_0, None),
850 Ok(LeafLookup::Exists)
851 ));
852 assert_eq!(
853 sparse.state_trie_ref().unwrap().get_leaf_value(&full_path_0),
854 Some(&leaf_value)
855 );
856
857 apply_account_update(&mut sparse, B256::ZERO, LeafUpdate::Changed(Vec::new()));
860 assert!(matches!(
861 sparse.state_trie_ref().unwrap().find_leaf(&full_path_0, None),
862 Ok(LeafLookup::NonExistent)
863 ));
864 assert!(sparse.state_trie_ref().unwrap().get_leaf_value(&full_path_0).is_none());
865 }
866
867 #[test]
868 fn reveal_storage_path_twice() {
869 let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
870
871 let full_path_0 = leaf_key([0x0], 64);
873
874 let leaf_value = alloy_rlp::encode(TrieAccount::default());
875 let leaf_1 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(
876 leaf_key([], 63),
877 leaf_value.clone(),
878 )));
879 let leaf_2 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(
880 leaf_key([], 63),
881 leaf_value.clone(),
882 )));
883
884 let multiproof = MultiProof {
885 storages: HashMap::from_iter([(
886 B256::ZERO,
887 StorageMultiProof {
888 root: B256::ZERO,
889 subtree: ProofNodes::from_iter([
890 (
891 Nibbles::default(),
892 alloy_rlp::encode(TrieNodeV2::Branch(BranchNodeV2 {
893 key: Nibbles::default(),
894 stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)],
895 state_mask: TrieMask::new(0b11),
896 branch_rlp_node: None,
897 }))
898 .into(),
899 ),
900 (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()),
901 (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
902 ]),
903 branch_node_masks: Default::default(),
904 },
905 )]),
906 ..Default::default()
907 };
908
909 sparse.reveal_decoded_multiproof(multiproof.try_into().unwrap()).unwrap();
911 assert!(matches!(
912 sparse.storage_trie_ref(&B256::ZERO).unwrap().find_leaf(&full_path_0, None),
913 Ok(LeafLookup::Exists)
914 ));
915 assert_eq!(
916 sparse.storage_trie_ref(&B256::ZERO).unwrap().get_leaf_value(&full_path_0),
917 Some(&leaf_value)
918 );
919
920 let mut updates = B256Map::from_iter([(B256::ZERO, LeafUpdate::Changed(Vec::new()))]);
923 sparse
924 .storage_trie_mut(&B256::ZERO)
925 .unwrap()
926 .update_leaves(&mut updates, |_, _| {})
927 .unwrap();
928 assert!(updates.is_empty());
929 assert!(matches!(
930 sparse.storage_trie_ref(&B256::ZERO).unwrap().find_leaf(&full_path_0, None),
931 Ok(LeafLookup::NonExistent)
932 ));
933 assert!(sparse
934 .storage_trie_ref(&B256::ZERO)
935 .unwrap()
936 .get_leaf_value(&full_path_0)
937 .is_none());
938 }
939
940 #[test]
941 fn seeded_hot_cache_capacities_preserve_first_cycle_touches() {
942 let account = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
943 let slot = b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
944 let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
945
946 sparse.set_hot_cache_capacities(1, 1);
947 sparse.record_account_touch(account);
948 sparse.record_slot_touch(account, slot);
949 sparse.prune(1, 1, TriePrefixSetsMut::default());
950
951 assert_eq!(sparse.hot_accounts_lfu.keys().copied().collect::<Vec<_>>(), vec![account]);
952 assert_eq!(
953 sparse.hot_slots_lfu.keys().copied().collect::<Vec<_>>(),
954 vec![HotSlotKey { address: account, slot }]
955 );
956 }
957
958 #[test]
959 fn take_changed_paths_from_sparse_state_trie() {
960 let account = B256::with_last_byte(0x01);
961 let slot = B256::with_last_byte(0x02);
962 let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
963 sparse.set_accounts_trie(RevealableSparseTrie::revealed_empty());
964 sparse.insert_storage_trie(account, RevealableSparseTrie::revealed_empty());
965 sparse.set_changed_paths(true);
966
967 let mut account_updates =
968 B256Map::from_iter([(account, LeafUpdate::Changed(vec![0x01; 32]))]);
969 sparse.trie_mut().update_leaves(&mut account_updates, |_, _| {}).unwrap();
970 assert!(account_updates.is_empty());
971 let _ = sparse.root().unwrap();
972
973 let mut storage_updates = B256Map::from_iter([(slot, LeafUpdate::Changed(vec![0x02; 32]))]);
974 sparse
975 .storage_trie_mut(&account)
976 .unwrap()
977 .update_leaves(&mut storage_updates, |_, _| {})
978 .unwrap();
979 assert!(storage_updates.is_empty());
980 let _ = sparse.storage_root(&account).unwrap();
981
982 let changed_paths = sparse.take_changed_paths().unwrap();
983 assert!(changed_paths.account_prefix_set.iter().any(|path| *path == Nibbles::default()));
984 assert!(changed_paths.storage_prefix_sets[&account]
985 .iter()
986 .any(|path| *path == Nibbles::default()));
987 }
988
989 #[test]
990 fn prune_keeps_retained_paths_overlay_account_and_storage() {
991 let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
992
993 let account = B256::ZERO;
994 let slot = B256::ZERO;
995 let account_path = leaf_key([0x0], 64);
996 let storage_path = leaf_key([0x0], 64);
997
998 let leaf_value = alloy_rlp::encode(TrieAccount::default());
999 let leaf_0 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(
1000 leaf_key([], 63),
1001 leaf_value.clone(),
1002 )));
1003 let leaf_1 =
1004 alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(leaf_key([], 63), leaf_value)));
1005
1006 let subtree = || {
1007 ProofNodes::from_iter([
1008 (
1009 Nibbles::default(),
1010 alloy_rlp::encode(TrieNodeV2::Branch(BranchNodeV2 {
1011 key: Nibbles::default(),
1012 stack: vec![RlpNode::from_rlp(&leaf_0), RlpNode::from_rlp(&leaf_1)],
1013 state_mask: TrieMask::new(0b11),
1014 branch_rlp_node: None,
1015 }))
1016 .into(),
1017 ),
1018 (Nibbles::from_nibbles([0x0]), leaf_0.clone().into()),
1019 (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
1020 ])
1021 };
1022
1023 let multiproof = MultiProof {
1024 account_subtree: subtree(),
1025 storages: HashMap::from_iter([(
1026 account,
1027 StorageMultiProof {
1028 root: B256::ZERO,
1029 subtree: subtree(),
1030 branch_node_masks: Default::default(),
1031 },
1032 )]),
1033 ..Default::default()
1034 };
1035
1036 sparse.reveal_decoded_multiproof(multiproof.try_into().unwrap()).unwrap();
1037 let trie_account = TrieAccount {
1038 storage_root: sparse.storage_root(&account).unwrap(),
1039 ..Default::default()
1040 };
1041 apply_account_update(
1042 &mut sparse,
1043 account,
1044 LeafUpdate::Changed(alloy_rlp::encode(trie_account)),
1045 );
1046 sparse.root().unwrap();
1047
1048 let mut retained_paths = TriePrefixSetsMut::default();
1049 retained_paths.account_prefix_set.insert(Nibbles::unpack(account));
1050 retained_paths
1051 .storage_prefix_sets
1052 .entry(account)
1053 .or_default()
1054 .insert(Nibbles::unpack(slot));
1055 sparse.prune(0, 0, retained_paths);
1056
1057 assert!(matches!(
1058 sparse.state_trie_ref().unwrap().find_leaf(&account_path, None),
1059 Ok(LeafLookup::Exists)
1060 ));
1061 assert!(matches!(
1062 sparse.storage_trie_ref(&account).unwrap().find_leaf(&storage_path, None),
1063 Ok(LeafLookup::Exists)
1064 ));
1065 }
1066
1067 #[test]
1068 fn reveal_v2_proof_nodes() {
1069 let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
1070
1071 let full_path_0 = leaf_key([0x0], 64);
1073
1074 let leaf_value = alloy_rlp::encode(TrieAccount::default());
1075 let leaf_1_node = TrieNodeV2::Leaf(LeafNode::new(leaf_key([], 63), leaf_value.clone()));
1076 let leaf_2_node = TrieNodeV2::Leaf(LeafNode::new(leaf_key([], 63), leaf_value.clone()));
1077
1078 let branch_node = TrieNodeV2::Branch(BranchNodeV2 {
1079 key: Nibbles::default(),
1080 stack: vec![
1081 RlpNode::from_rlp(&alloy_rlp::encode(&leaf_1_node)),
1082 RlpNode::from_rlp(&alloy_rlp::encode(&leaf_2_node)),
1083 ],
1084 state_mask: TrieMask::new(0b11),
1085 branch_rlp_node: None,
1086 });
1087
1088 let v2_proof_nodes = vec![
1090 ProofTrieNodeV2 {
1091 path: Nibbles::default(),
1092 node: branch_node,
1093 masks: Some(BranchNodeMasks {
1094 hash_mask: TrieMask::default(),
1095 tree_mask: TrieMask::default(),
1096 }),
1097 },
1098 ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x0]), node: leaf_1_node, masks: None },
1099 ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x1]), node: leaf_2_node, masks: None },
1100 ];
1101
1102 sparse
1104 .reveal_decoded_multiproof_v2(reth_trie_common::DecodedMultiProofV2 {
1105 account_proofs: v2_proof_nodes,
1106 ..Default::default()
1107 })
1108 .unwrap();
1109
1110 assert!(matches!(
1112 sparse.state_trie_ref().unwrap().find_leaf(&full_path_0, None),
1113 Ok(LeafLookup::Exists)
1114 ));
1115 assert_eq!(
1116 sparse.state_trie_ref().unwrap().get_leaf_value(&full_path_0),
1117 Some(&leaf_value)
1118 );
1119
1120 apply_account_update(&mut sparse, B256::ZERO, LeafUpdate::Changed(Vec::new()));
1122 assert!(sparse.state_trie_ref().unwrap().get_leaf_value(&full_path_0).is_none());
1123 }
1124
1125 #[test]
1126 fn reveal_storage_v2_proof_nodes() {
1127 let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
1128
1129 let full_path_0 = leaf_key([0x0], 64);
1131
1132 let storage_value: Vec<u8> = alloy_rlp::encode_fixed_size(&U256::from(42)).to_vec();
1133 let leaf_1_node = TrieNodeV2::Leaf(LeafNode::new(leaf_key([], 63), storage_value.clone()));
1134 let leaf_2_node = TrieNodeV2::Leaf(LeafNode::new(leaf_key([], 63), storage_value.clone()));
1135
1136 let branch_node = TrieNodeV2::Branch(BranchNodeV2 {
1137 key: Nibbles::default(),
1138 stack: vec![
1139 RlpNode::from_rlp(&alloy_rlp::encode(&leaf_1_node)),
1140 RlpNode::from_rlp(&alloy_rlp::encode(&leaf_2_node)),
1141 ],
1142 state_mask: TrieMask::new(0b11),
1143 branch_rlp_node: None,
1144 });
1145
1146 let v2_proof_nodes = vec![
1147 ProofTrieNodeV2 { path: Nibbles::default(), node: branch_node, masks: None },
1148 ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x0]), node: leaf_1_node, masks: None },
1149 ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x1]), node: leaf_2_node, masks: None },
1150 ];
1151
1152 sparse
1154 .reveal_decoded_multiproof_v2(reth_trie_common::DecodedMultiProofV2 {
1155 storage_proofs: B256Map::from_iter([(B256::ZERO, v2_proof_nodes)]),
1156 ..Default::default()
1157 })
1158 .unwrap();
1159
1160 assert!(matches!(
1162 sparse.storage_trie_ref(&B256::ZERO).unwrap().find_leaf(&full_path_0, None),
1163 Ok(LeafLookup::Exists)
1164 ));
1165 assert_eq!(
1166 sparse.storage_trie_ref(&B256::ZERO).unwrap().get_leaf_value(&full_path_0),
1167 Some(&storage_value)
1168 );
1169
1170 let mut updates = B256Map::from_iter([(B256::ZERO, LeafUpdate::Changed(Vec::new()))]);
1172 sparse
1173 .storage_trie_mut(&B256::ZERO)
1174 .unwrap()
1175 .update_leaves(&mut updates, |_, _| {})
1176 .unwrap();
1177 assert!(updates.is_empty());
1178 assert!(sparse
1179 .storage_trie_ref(&B256::ZERO)
1180 .unwrap()
1181 .get_leaf_value(&full_path_0)
1182 .is_none());
1183 }
1184
1185 #[test]
1186 fn root_on_blind_trie_returns_blind_error() {
1187 let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
1188
1189 let err = sparse.root().unwrap_err();
1190
1191 assert!(matches!(err.kind(), SparseStateTrieErrorKind::Sparse(SparseTrieErrorKind::Blind)));
1192 }
1193
1194 #[test]
1195 fn take_trie_updates() {
1196 reth_tracing::init_test_tracing();
1197
1198 let mut rng = StdRng::seed_from_u64(1);
1200
1201 let mut bytes = [0u8; 1024];
1202 rng.fill(bytes.as_mut_slice());
1203
1204 let slot_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
1205 let slot_path_1 = Nibbles::unpack(slot_1);
1206 let value_1 = U256::from(rng.random::<u64>());
1207 let slot_2 = b256!("0x1100000000000000000000000000000000000000000000000000000000000000");
1208 let slot_path_2 = Nibbles::unpack(slot_2);
1209 let value_2 = U256::from(rng.random::<u64>());
1210 let slot_3 = b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
1211 let value_3 = U256::from(rng.random::<u64>());
1212
1213 let mut storage_hash_builder = HashBuilder::default()
1214 .with_proof_retainer(ProofRetainer::from_iter([slot_path_1, slot_path_2]));
1215 storage_hash_builder.add_leaf(slot_path_1, &alloy_rlp::encode_fixed_size(&value_1));
1216 storage_hash_builder.add_leaf(slot_path_2, &alloy_rlp::encode_fixed_size(&value_2));
1217
1218 let storage_root = storage_hash_builder.root();
1219 let storage_proof_nodes = storage_hash_builder.take_proof_nodes();
1220 let storage_branch_node_masks = BranchNodeMasksMap::from_iter([
1221 (
1222 Nibbles::default(),
1223 BranchNodeMasks { hash_mask: TrieMask::new(0b010), tree_mask: TrieMask::default() },
1224 ),
1225 (
1226 Nibbles::from_nibbles([0x1]),
1227 BranchNodeMasks { hash_mask: TrieMask::new(0b11), tree_mask: TrieMask::default() },
1228 ),
1229 ]);
1230
1231 let address_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
1232 let address_path_1 = Nibbles::unpack(address_1);
1233 let account_1 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
1234 let mut trie_account_1 = account_1.into_trie_account(storage_root);
1235 let address_2 = b256!("0x1100000000000000000000000000000000000000000000000000000000000000");
1236 let address_path_2 = Nibbles::unpack(address_2);
1237 let account_2 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
1238 let mut trie_account_2 = account_2.into_trie_account(EMPTY_ROOT_HASH);
1239
1240 let mut hash_builder = HashBuilder::default()
1241 .with_proof_retainer(ProofRetainer::from_iter([address_path_1, address_path_2]));
1242 hash_builder.add_leaf(address_path_1, &alloy_rlp::encode(trie_account_1));
1243 hash_builder.add_leaf(address_path_2, &alloy_rlp::encode(trie_account_2));
1244
1245 let root = hash_builder.root();
1246 let proof_nodes = hash_builder.take_proof_nodes();
1247 let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default().with_updates(true);
1248 sparse
1249 .reveal_decoded_multiproof(
1250 MultiProof {
1251 account_subtree: proof_nodes,
1252 branch_node_masks: BranchNodeMasksMap::from_iter([(
1253 Nibbles::from_nibbles([0x1]),
1254 BranchNodeMasks {
1255 hash_mask: TrieMask::new(0b00),
1256 tree_mask: TrieMask::default(),
1257 },
1258 )]),
1259 storages: HashMap::from_iter([
1260 (
1261 address_1,
1262 StorageMultiProof {
1263 root,
1264 subtree: storage_proof_nodes.clone(),
1265 branch_node_masks: storage_branch_node_masks.clone(),
1266 },
1267 ),
1268 (
1269 address_2,
1270 StorageMultiProof {
1271 root,
1272 subtree: storage_proof_nodes,
1273 branch_node_masks: storage_branch_node_masks,
1274 },
1275 ),
1276 ]),
1277 }
1278 .try_into()
1279 .unwrap(),
1280 )
1281 .unwrap();
1282
1283 assert_eq!(sparse.root().unwrap(), root);
1284
1285 let address_3 = b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
1286 let account_3 = Account { nonce: account_1.nonce + 1, ..account_1 };
1287 let trie_account_3 = account_3.into_trie_account(EMPTY_ROOT_HASH);
1288
1289 apply_account_update(
1290 &mut sparse,
1291 address_3,
1292 LeafUpdate::Changed(alloy_rlp::encode(trie_account_3)),
1293 );
1294
1295 let mut updates =
1296 B256Map::from_iter([(slot_3, LeafUpdate::Changed(alloy_rlp::encode(value_3)))]);
1297 sparse
1298 .storage_trie_mut(&address_1)
1299 .unwrap()
1300 .update_leaves(&mut updates, |_, _| {})
1301 .unwrap();
1302 assert!(updates.is_empty());
1303 trie_account_1.storage_root = sparse.storage_root(&address_1).unwrap();
1304 apply_account_update(
1305 &mut sparse,
1306 address_1,
1307 LeafUpdate::Changed(alloy_rlp::encode(trie_account_1)),
1308 );
1309
1310 sparse.wipe_storage(address_2).unwrap();
1311 trie_account_2.storage_root = sparse.storage_root(&address_2).unwrap();
1312 apply_account_update(
1313 &mut sparse,
1314 address_2,
1315 LeafUpdate::Changed(alloy_rlp::encode(trie_account_2)),
1316 );
1317
1318 sparse.root().unwrap();
1319
1320 let sparse_updates = sparse.take_trie_updates().unwrap();
1321 pretty_assertions::assert_eq!(
1323 sparse_updates,
1324 TrieUpdates {
1325 account_nodes: HashMap::default(),
1326 storage_tries: HashMap::from_iter([
1327 (
1328 b256!("0x1000000000000000000000000000000000000000000000000000000000000000"),
1329 StorageTrieUpdates {
1330 is_deleted: false,
1331 storage_nodes: HashMap::default(),
1332 removed_nodes: HashSet::from_iter([Nibbles::from_nibbles([0x1])])
1333 }
1334 ),
1335 (
1336 b256!("0x1100000000000000000000000000000000000000000000000000000000000000"),
1337 StorageTrieUpdates {
1338 is_deleted: true,
1339 storage_nodes: HashMap::default(),
1340 removed_nodes: HashSet::default()
1341 }
1342 )
1343 ]),
1344 removed_nodes: HashSet::default()
1345 }
1346 );
1347 }
1348}