Skip to main content

reth_trie_sparse/
state.rs

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/// Holds data that should be dropped after any locks are released.
21///
22/// This is used to defer expensive deallocations (like proof node buffers) until after final state
23/// root is calculated
24#[derive(Debug, Default)]
25pub struct DeferredDrops {
26    /// Each nodes reveal operation creates a new buffer, uses it, and pushes it here.
27    pub proof_nodes_bufs: Vec<Vec<ProofTrieNodeV2>>,
28}
29
30#[derive(Debug)]
31/// Sparse state trie representing lazy-loaded Ethereum state trie.
32pub struct SparseStateTrie<
33    A = ArenaParallelSparseTrie, // Account trie implementation
34    S = ArenaParallelSparseTrie, // Storage trie implementation
35> {
36    /// Sparse account trie.
37    state: RevealableSparseTrie<A>,
38    /// State related to storage tries.
39    storage: StorageTries<S>,
40    /// Flag indicating whether trie updates should be retained.
41    retain_updates: bool,
42    /// Flag indicating whether changed node base paths should be retained.
43    retain_changed_paths: bool,
44    /// Holds data that should be dropped after final state root is calculated.
45    deferred_drops: DeferredDrops,
46    /// Global LFU tracker for hot `(address, slot)` storage entries.
47    hot_slots_lfu: BucketedLfu<HotSlotKey>,
48    /// Global LFU tracker for hot account entries.
49    hot_accounts_lfu: BucketedLfu<B256>,
50    /// Metrics for the sparse state trie.
51    #[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    /// Create state trie from state trie.
78    pub fn from_state(state: RevealableSparseTrie) -> Self {
79        Self { state, ..Default::default() }
80    }
81}
82
83impl<A, S> SparseStateTrie<A, S> {
84    /// Set the retention of branch node updates and deletions.
85    pub const fn set_updates(&mut self, retain_updates: bool) {
86        self.retain_updates = retain_updates;
87    }
88
89    /// Set the retention of branch node updates and deletions.
90    pub const fn with_updates(mut self, retain_updates: bool) -> Self {
91        self.set_updates(retain_updates);
92        self
93    }
94
95    /// Seeds the hot account/storage LFU caches with their configured capacities.
96    ///
97    /// This must happen before the first `record_*_touch` call, otherwise touches are ignored while
98    /// the LFUs still have zero capacity.
99    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    /// Seeds the hot account/storage LFU caches with their configured capacities.
105    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    /// Set the accounts trie to the given `RevealableSparseTrie`.
115    pub fn set_accounts_trie(&mut self, trie: RevealableSparseTrie<A>) {
116        self.state = trie;
117    }
118
119    /// Set the accounts trie to the given `RevealableSparseTrie`.
120    pub fn with_accounts_trie(mut self, trie: RevealableSparseTrie<A>) -> Self {
121        self.set_accounts_trie(trie);
122        self
123    }
124
125    /// Set the default trie which will be cloned when creating new storage
126    /// [`RevealableSparseTrie`]s.
127    pub fn set_default_storage_trie(&mut self, trie: RevealableSparseTrie<S>) {
128        self.storage.default_trie = trie;
129    }
130
131    /// Set the default trie which will be cloned when creating new storage
132    /// [`RevealableSparseTrie`]s.
133    pub fn with_default_storage_trie(mut self, trie: RevealableSparseTrie<S>) -> Self {
134        self.set_default_storage_trie(trie);
135        self
136    }
137
138    /// Takes the data structures for deferred dropping.
139    ///
140    /// This allows the caller to drop the buffers later, avoiding expensive deallocations while
141    /// calculating the state root.
142    pub fn take_deferred_drops(&mut self) -> DeferredDrops {
143        core::mem::take(&mut self.deferred_drops)
144    }
145}
146
147impl SparseStateTrie {
148    /// Create new [`SparseStateTrie`] with the default trie implementation.
149    pub fn new() -> Self {
150        Self::default()
151    }
152}
153
154impl<A: SparseTrieTrait, S: SparseTrieTrait> SparseStateTrie<A, S> {
155    /// Set the retention of changed node base paths.
156    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    /// Set the retention of changed node base paths.
169    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    /// Returns storage trie changed paths for tries that have been revealed.
175    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    /// Returns changed paths by taking them from the revealed sparse tries.
187    ///
188    /// Returns `None` if the accounts trie is not revealed.
189    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    /// Takes all debug recorders from the account trie and all revealed storage tries.
199    ///
200    /// Returns a vec of `(Option<B256>, TrieDebugRecorder)` where `None` is the account trie
201    /// key, and `Some(address)` are storage trie keys.
202    #[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    /// Returns mutable reference to account trie.
223    pub const fn trie_mut(&mut self) -> &mut RevealableSparseTrie<A> {
224        &mut self.state
225    }
226
227    /// Returns `true` if the account path has been revealed in the sparse trie.
228    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    /// Was the storage-slot witness for (`address`,`slot`) complete?
239    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    /// Records a storage slot access/update in the global LFU tracker.
250    #[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    /// Records an account access/update in the global LFU tracker.
256    #[inline]
257    pub fn record_account_touch(&mut self, account: B256) {
258        self.hot_accounts_lfu.touch(account);
259    }
260
261    /// Returns reference to bytes representing leaf value for the target account.
262    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    /// Returns reference to bytes representing leaf value for the target account and storage slot.
267    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    /// Returns reference to state trie if it was revealed.
272    pub const fn state_trie_ref(&self) -> Option<&A> {
273        self.state.as_revealed_ref()
274    }
275
276    /// Returns reference to storage trie if it was revealed.
277    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    /// Returns mutable reference to storage sparse trie if it was revealed.
282    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    /// Returns mutable reference to storage tries.
287    pub const fn storage_tries_mut(&mut self) -> &mut B256Map<RevealableSparseTrie<S>> {
288        &mut self.storage.tries
289    }
290
291    /// Takes the storage trie for the provided address.
292    pub fn take_storage_trie(&mut self, address: &B256) -> Option<RevealableSparseTrie<S>> {
293        self.storage.tries.remove(address)
294    }
295
296    /// Takes the storage trie for the provided address, creating a blind one if it doesn't exist.
297    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    /// Inserts storage trie for the provided address.
304    pub fn insert_storage_trie(&mut self, address: B256, storage_trie: RevealableSparseTrie<S>) {
305        self.storage.tries.insert(address, storage_trie);
306    }
307
308    /// Returns mutable reference to storage sparse trie, creating a blind one if it doesn't exist.
309    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    /// Reveal unknown trie paths from multiproof.
317    /// NOTE: This method does not extensively validate the proof.
318    pub fn reveal_multiproof(&mut self, multiproof: MultiProof) -> SparseStateTrieResult<()> {
319        // first decode the multiproof
320        let decoded_multiproof = multiproof.try_into()?;
321
322        // then reveal the decoded multiproof
323        self.reveal_decoded_multiproof(decoded_multiproof)
324    }
325
326    /// Reveal unknown trie paths from decoded multiproof.
327    /// NOTE: This method does not extensively validate the proof.
328    #[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    /// Reveals a V2 decoded multiproof.
337    ///
338    /// V2 multiproofs use a simpler format where proof nodes are stored as vectors rather than
339    /// hashmaps, with masks already included in the `ProofTrieNode` structure.
340    #[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        // Collect `(trie, proof_nodes)` pairs for both the account trie and every storage trie
349        // touched by this multiproof.
350        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        // Ensure a storage trie exists for every address whose proofs we're about to reveal
359        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        // Accumulate the first error and defer dropping the proof node buffers.
426        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    /// Wipe the storage trie at the provided address.
438    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    /// Calculates the hashes of subtries.
446    ///
447    /// If the trie has not been revealed, this function does nothing.
448    #[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    /// Returns storage sparse trie root if the trie has been revealed.
456    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    /// Returns mutable reference to the revealed account sparse trie.
461    fn revealed_trie_mut(&mut self) -> SparseStateTrieResult<&mut A> {
462        self.state.as_revealed_mut().ok_or_else(|| SparseTrieErrorKind::Blind.into())
463    }
464
465    /// Returns sparse trie root.
466    pub fn root(&mut self) -> SparseStateTrieResult<B256> {
467        // record revealed node metrics
468        #[cfg(feature = "metrics")]
469        self.metrics.record();
470
471        Ok(self.revealed_trie_mut()?.root())
472    }
473
474    /// Returns sparse trie root and trie updates.
475    ///
476    /// Returns an error if the account trie is still blind.
477    #[instrument(level = "debug", target = "trie::sparse", skip_all)]
478    pub fn root_with_updates(&mut self) -> SparseStateTrieResult<(B256, TrieUpdates)> {
479        // record revealed node metrics
480        #[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    /// Returns storage trie updates for tries that have been revealed.
496    ///
497    /// Panics if any of the storage tries are not revealed.
498    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    /// Returns [`TrieUpdates`] by taking the updates from the revealed sparse tries.
517    ///
518    /// Returns `None` if the accounts trie is not revealed.
519    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    /// Clears all trie data while preserving allocations for reuse.
538    ///
539    /// This resets the trie to an empty state but keeps the underlying memory allocations,
540    /// which can significantly reduce allocation overhead when the trie is reused.
541    pub fn clear(&mut self) {
542        self.state.clear();
543        self.storage.clear();
544    }
545
546    /// Returns a heuristic for the total in-memory size of this state trie in bytes.
547    ///
548    /// This aggregates the memory usage of the account trie, all revealed storage tries
549    /// (including cleared ones retained for allocation reuse), and auxiliary data structures.
550    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    /// Returns the number of storage tries currently retained (active + cleared).
581    pub fn retained_storage_tries_count(&self) -> usize {
582        self.storage.tries.len() + self.storage.cleared_tries.len()
583    }
584
585    /// Prunes account/storage tries according to global LFU retention and retained paths.
586    ///
587    /// - Top LFU `(address, slot)` entries are retained up to `max_hot_slots`.
588    ///
589    /// - Top LFU `(address, slot)` entries are retained in storage tries.
590    /// - Account trie retains only paths for accounts tracked by the account LFU.
591    /// - Storage tries retain only paths needed for retained slots.
592    /// - Additional retained paths are unioned with LFU-selected paths.
593    /// - All other revealed paths are pruned to hash stubs or fully evicted.
594    ///
595    /// # Preconditions
596    ///
597    /// All revealed account and storage tries must already have computed hashes via `root()`
598    /// / `storage_root()` for their current state. Pruning a dirty revealed trie is a hard
599    /// error and may panic.
600    #[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        // Prune account and storage tries in parallel using the same retained set.
633        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/// The fields of [`SparseStateTrie`] related to storage tries. This is kept separate from the rest
683/// of [`SparseStateTrie`] to help enforce allocation re-use.
684#[derive(Debug, Default)]
685struct StorageTries<S = ArenaParallelSparseTrie> {
686    /// Sparse storage tries.
687    tries: B256Map<RevealableSparseTrie<S>>,
688    /// Cleared storage tries, kept for re-use.
689    cleared_tries: Vec<RevealableSparseTrie<S>>,
690    /// A default cleared trie instance, which will be cloned when creating new tries.
691    default_trie: RevealableSparseTrie<S>,
692}
693
694#[cfg(feature = "std")]
695impl<S: SparseTrieTrait> StorageTries<S> {
696    /// Prunes storage tries using LFU-retained slots.
697    ///
698    /// Tries without retained slots are evicted entirely. Tries with retained slots are pruned to
699    /// those slots.
700    fn prune_by_retained_slots(
701        &mut self,
702        retained_slots: B256Map<PrefixSet>,
703        parent_span: &tracing::Span,
704    ) -> usize {
705        // Parallel pass: prune retained tries and clear evicted ones in place.
706        {
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        // Cheap sequential drain: move already-cleared tries into the reuse pool.
729        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    /// Returns all fields to a cleared state, equivalent to the default state, keeping cleared
750    /// collections for re-use later when possible.
751    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    // Returns mutable reference to storage sparse trie, creating a blind one if it doesn't exist.
761    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/// Key for identifying a storage slot in the global LFU cache.
769#[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    /// Create a leaf key (suffix) with given nibbles padded with zeros to reach `total_len`.
796    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        // Full 64-nibble paths
814        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        // Leaf key is 63 nibbles (suffix after 1-nibble node path)
819        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        // Reveal multiproof and check that the state trie contains the leaf node and value
847        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        // Remove the leaf node and check that the state trie does not contain the leaf node and
858        // value
859        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        // Full 64-nibble path
872        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        // Reveal multiproof and check that the storage trie contains the leaf node and value
910        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        // Remove the leaf node and check that the storage trie does not contain the leaf node and
921        // value
922        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        // Full 64-nibble path
1072        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        // Create V2 proof nodes with masks already included
1089        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        // Reveal V2 proof nodes
1103        sparse
1104            .reveal_decoded_multiproof_v2(reth_trie_common::DecodedMultiProofV2 {
1105                account_proofs: v2_proof_nodes,
1106                ..Default::default()
1107            })
1108            .unwrap();
1109
1110        // Check that the state trie contains the leaf node and value
1111        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        // Remove the leaf node
1121        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        // Full 64-nibble path
1130        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        // Reveal V2 storage proof nodes for account
1153        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        // Check that the storage trie contains the leaf node and value
1161        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        // Remove the leaf node
1171        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 = generators::rng();
1199        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        // TODO(alexey): assert against real state root calculation updates
1322        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}