reth_trie_sparse/
state.rs

1use crate::{
2    traits::SparseTrie as SparseTrieTrait, ArenaParallelSparseTrie, RevealableSparseTrie,
3    TrieNodeEpoch,
4};
5use alloc::vec::Vec;
6use alloy_primitives::{map::B256Map, B256};
7use either::Either;
8use reth_execution_errors::{SparseStateTrieResult, SparseTrieErrorKind};
9use reth_trie_common::{
10    updates::{StorageTrieUpdates, TrieUpdates},
11    DecodedMultiProof, MultiProof, Nibbles, ProofTrieNodeV2,
12};
13#[cfg(feature = "std")]
14use tracing::debug;
15use tracing::instrument;
16
17/// Holds data that should be dropped after any locks are released.
18///
19/// This is used to defer expensive deallocations (like proof node buffers) until after final state
20/// root is calculated
21#[derive(Debug, Default)]
22pub struct DeferredDrops {
23    /// Each nodes reveal operation creates a new buffer, uses it, and pushes it here.
24    pub proof_nodes_bufs: Vec<Vec<ProofTrieNodeV2>>,
25}
26
27#[derive(Debug)]
28/// Sparse state trie representing lazy-loaded Ethereum state trie.
29pub struct SparseStateTrie<
30    A = ArenaParallelSparseTrie, // Account trie implementation
31    S = ArenaParallelSparseTrie, // Storage trie implementation
32> {
33    /// Sparse account trie.
34    state: RevealableSparseTrie<A>,
35    /// State related to storage tries.
36    storage: StorageTries<S>,
37    /// Flag indicating whether trie updates should be retained.
38    retain_updates: bool,
39    /// Holds data that should be dropped after final state root is calculated.
40    deferred_drops: DeferredDrops,
41    /// Metrics for the sparse state trie.
42    #[cfg(feature = "metrics")]
43    metrics: crate::metrics::SparseStateTrieMetrics,
44}
45
46impl<A, S> Default for SparseStateTrie<A, S>
47where
48    A: Default,
49    S: Default,
50{
51    fn default() -> Self {
52        Self {
53            state: Default::default(),
54            storage: Default::default(),
55            retain_updates: false,
56            deferred_drops: DeferredDrops::default(),
57            #[cfg(feature = "metrics")]
58            metrics: Default::default(),
59        }
60    }
61}
62
63#[cfg(test)]
64impl SparseStateTrie {
65    /// Create state trie from state trie.
66    pub fn from_state(state: RevealableSparseTrie) -> Self {
67        Self { state, ..Default::default() }
68    }
69}
70
71impl<A, S> SparseStateTrie<A, S> {
72    /// Set the retention of branch node updates and deletions.
73    pub const fn set_updates(&mut self, retain_updates: bool) {
74        self.retain_updates = retain_updates;
75    }
76
77    /// Set the retention of branch node updates and deletions.
78    pub const fn with_updates(mut self, retain_updates: bool) -> Self {
79        self.set_updates(retain_updates);
80        self
81    }
82
83    /// Returns whether branch node updates and deletions are retained.
84    pub const fn retains_updates(&self) -> bool {
85        self.retain_updates
86    }
87
88    /// Set the accounts trie to the given `RevealableSparseTrie`.
89    pub fn set_accounts_trie(&mut self, trie: RevealableSparseTrie<A>) {
90        self.state = trie;
91    }
92
93    /// Set the accounts trie to the given `RevealableSparseTrie`.
94    pub fn with_accounts_trie(mut self, trie: RevealableSparseTrie<A>) -> Self {
95        self.set_accounts_trie(trie);
96        self
97    }
98
99    /// Set the default trie which will be cloned when creating new storage
100    /// [`RevealableSparseTrie`]s.
101    pub fn set_default_storage_trie(&mut self, trie: RevealableSparseTrie<S>) {
102        self.storage.default_trie = trie;
103    }
104
105    /// Set the default trie which will be cloned when creating new storage
106    /// [`RevealableSparseTrie`]s.
107    pub fn with_default_storage_trie(mut self, trie: RevealableSparseTrie<S>) -> Self {
108        self.set_default_storage_trie(trie);
109        self
110    }
111
112    /// Takes the data structures for deferred dropping.
113    ///
114    /// This allows the caller to drop the buffers later, avoiding expensive deallocations while
115    /// calculating the state root.
116    pub fn take_deferred_drops(&mut self) -> DeferredDrops {
117        core::mem::take(&mut self.deferred_drops)
118    }
119}
120
121impl SparseStateTrie {
122    /// Create new [`SparseStateTrie`] with the default trie implementation.
123    pub fn new() -> Self {
124        Self::default()
125    }
126}
127
128impl<A, S> SparseStateTrie<A, S>
129where
130    A: SparseTrieTrait + Default,
131    S: SparseTrieTrait + Default + Clone,
132{
133    /// Returns mutable reference to account trie.
134    pub const fn trie_mut(&mut self) -> &mut RevealableSparseTrie<A> {
135        &mut self.state
136    }
137
138    /// Returns `true` if the account path has been revealed in the sparse trie.
139    pub fn is_account_revealed(&self, account: B256) -> bool {
140        let path = Nibbles::unpack(account);
141        let trie = match self.state_trie_ref() {
142            Some(t) => t,
143            None => return false,
144        };
145
146        trie.find_leaf(&path, None).is_ok()
147    }
148
149    /// Was the storage-slot witness for (`address`,`slot`) complete?
150    pub fn check_valid_storage_witness(&self, address: B256, slot: B256) -> bool {
151        let path = Nibbles::unpack(slot);
152        let trie = match self.storage_trie_ref(&address) {
153            Some(t) => t,
154            None => return false,
155        };
156
157        trie.find_leaf(&path, None).is_ok()
158    }
159
160    /// Returns reference to bytes representing leaf value for the target account.
161    pub fn get_account_value(&self, account: &B256) -> Option<&Vec<u8>> {
162        self.state.as_revealed_ref()?.get_leaf_value(&Nibbles::unpack(account))
163    }
164
165    /// Returns reference to bytes representing leaf value for the target account and storage slot.
166    pub fn get_storage_slot_value(&self, account: &B256, slot: &B256) -> Option<&Vec<u8>> {
167        self.storage.tries.get(account)?.as_revealed_ref()?.get_leaf_value(&Nibbles::unpack(slot))
168    }
169
170    /// Returns reference to state trie if it was revealed.
171    pub const fn state_trie_ref(&self) -> Option<&A> {
172        self.state.as_revealed_ref()
173    }
174
175    /// Returns reference to storage trie if it was revealed.
176    pub fn storage_trie_ref(&self, address: &B256) -> Option<&S> {
177        self.storage.tries.get(address).and_then(|e| e.as_revealed_ref())
178    }
179
180    /// Returns mutable reference to storage sparse trie if it was revealed.
181    pub fn storage_trie_mut(&mut self, address: &B256) -> Option<&mut S> {
182        self.storage.tries.get_mut(address).and_then(|e| e.as_revealed_mut())
183    }
184
185    /// Returns mutable reference to storage tries.
186    pub const fn storage_tries_mut(&mut self) -> &mut B256Map<RevealableSparseTrie<S>> {
187        &mut self.storage.tries
188    }
189
190    /// Takes the storage trie for the provided address.
191    pub fn take_storage_trie(&mut self, address: &B256) -> Option<RevealableSparseTrie<S>> {
192        self.storage.tries.remove(address)
193    }
194
195    /// Takes the storage trie for the provided address, creating a blind one if it doesn't exist.
196    pub fn take_or_create_storage_trie(&mut self, address: &B256) -> RevealableSparseTrie<S> {
197        self.storage.tries.remove(address).unwrap_or_else(|| {
198            self.storage.cleared_tries.pop().unwrap_or_else(|| self.storage.default_trie.clone())
199        })
200    }
201
202    /// Inserts storage trie for the provided address.
203    pub fn insert_storage_trie(&mut self, address: B256, storage_trie: RevealableSparseTrie<S>) {
204        self.storage.tries.insert(address, storage_trie);
205    }
206
207    /// Returns mutable reference to storage sparse trie, creating a blind one if it doesn't exist.
208    pub fn get_or_create_storage_trie_mut(
209        &mut self,
210        address: B256,
211    ) -> &mut RevealableSparseTrie<S> {
212        self.storage.get_or_create_trie_mut(address)
213    }
214
215    /// Reveal unknown trie paths from multiproof.
216    /// NOTE: This method does not extensively validate the proof.
217    pub fn reveal_multiproof(&mut self, multiproof: MultiProof) -> SparseStateTrieResult<()> {
218        // first decode the multiproof
219        let decoded_multiproof = multiproof.try_into()?;
220
221        // then reveal the decoded multiproof
222        self.reveal_decoded_multiproof(decoded_multiproof)
223    }
224
225    /// Reveal unknown trie paths from decoded multiproof.
226    /// NOTE: This method does not extensively validate the proof.
227    #[instrument(level = "debug", target = "trie::sparse", skip_all)]
228    pub fn reveal_decoded_multiproof(
229        &mut self,
230        multiproof: DecodedMultiProof,
231    ) -> SparseStateTrieResult<()> {
232        self.reveal_decoded_multiproof_v2(multiproof.into())
233    }
234
235    /// Reveals a V2 decoded multiproof.
236    ///
237    /// V2 multiproofs use a simpler format where proof nodes are stored as vectors rather than
238    /// hashmaps, with masks already included in the `ProofTrieNode` structure.
239    #[instrument(level = "debug", target = "trie::sparse", skip_all)]
240    pub fn reveal_decoded_multiproof_v2(
241        &mut self,
242        multiproof: reth_trie_common::DecodedMultiProofV2,
243    ) -> SparseStateTrieResult<()> {
244        let reth_trie_common::DecodedMultiProofV2 { account_proofs, mut storage_proofs, .. } =
245            multiproof;
246
247        // Collect `(trie, proof_nodes)` pairs for both the account trie and every storage trie
248        // touched by this multiproof.
249        let mut targets = Vec::with_capacity(storage_proofs.len() + 1);
250
251        if !account_proofs.is_empty() {
252            #[cfg(feature = "metrics")]
253            self.metrics.increment_total_account_nodes(account_proofs.len() as u64);
254            targets.push((None, Either::Left(&mut self.state), account_proofs));
255        }
256
257        // Ensure a storage trie exists for every address whose proofs we're about to reveal
258        for &account in storage_proofs.keys() {
259            let _ = self.storage.get_or_create_trie_mut(account);
260        }
261
262        for (account, trie) in &mut self.storage.tries {
263            if let Some(nodes) = storage_proofs.remove(account) {
264                #[cfg(feature = "metrics")]
265                self.metrics.increment_total_storage_nodes(nodes.len() as u64);
266                targets.push((Some(*account), Either::Right(trie), nodes));
267            }
268        }
269
270        let retain_updates = self.retain_updates;
271
272        #[cfg(not(feature = "std"))]
273        let results: Vec<_> = targets
274            .into_iter()
275            .map(|(_, target, mut nodes)| {
276                let result = match target {
277                    Either::Left(trie) => trie.reveal_v2_proof_nodes(&mut nodes, retain_updates),
278                    Either::Right(trie) => trie.reveal_v2_proof_nodes(&mut nodes, retain_updates),
279                };
280                (result, nodes)
281            })
282            .collect();
283
284        #[cfg(feature = "std")]
285        let results: Vec<_> = {
286            use rayon::iter::ParallelIterator;
287            use reth_primitives_traits::ParallelBridgeBuffered;
288
289            let parent_span = tracing::Span::current();
290            targets
291                .into_iter()
292                .par_bridge_buffered()
293                .map(|(hashed_address, target, mut nodes)| {
294                    let _span = tracing::trace_span!(
295                        target: "trie::sparse",
296                        parent: &parent_span,
297                        "reveal_v2_proof_nodes",
298                        ?hashed_address,
299                    )
300                    .entered();
301
302                    let result = match target {
303                        Either::Left(trie) => {
304                            trie.reveal_v2_proof_nodes(&mut nodes, retain_updates)
305                        }
306                        Either::Right(trie) => {
307                            trie.reveal_v2_proof_nodes(&mut nodes, retain_updates)
308                        }
309                    };
310                    (result, nodes)
311                })
312                .collect()
313        };
314
315        // Accumulate the first error and defer dropping the proof node buffers.
316        let mut any_err = Ok(());
317        for (result, nodes) in results {
318            if result.is_err() && any_err.is_ok() {
319                any_err = result.map_err(Into::into);
320            }
321            self.deferred_drops.proof_nodes_bufs.push(nodes);
322        }
323
324        any_err
325    }
326
327    /// Reveals account trie proof nodes on the calling thread.
328    ///
329    /// Unlike [`Self::reveal_decoded_multiproof_v2`] this touches only the account trie, so a
330    /// caller that owns some of the storage tries itself can reveal the two halves separately.
331    pub fn reveal_account_proof_nodes(
332        &mut self,
333        mut nodes: Vec<ProofTrieNodeV2>,
334    ) -> SparseStateTrieResult<()> {
335        if nodes.is_empty() {
336            return Ok(())
337        }
338
339        #[cfg(feature = "metrics")]
340        self.metrics.increment_total_account_nodes(nodes.len() as u64);
341
342        let result = self.state.reveal_v2_proof_nodes(&mut nodes, self.retain_updates);
343        self.deferred_drops.proof_nodes_bufs.push(nodes);
344
345        Ok(result?)
346    }
347
348    /// Records storage trie nodes that were revealed into a trie taken out of this state trie, so
349    /// the reveal metrics stay complete.
350    pub const fn record_revealed_storage_nodes(&mut self, nodes: usize) {
351        #[cfg(feature = "metrics")]
352        self.metrics.increment_total_storage_nodes(nodes as u64);
353        #[cfg(not(feature = "metrics"))]
354        let _ = nodes;
355    }
356
357    /// Calculates the hashes of subtries.
358    ///
359    /// If the trie has not been revealed, this function does nothing.
360    #[instrument(level = "debug", target = "trie::sparse", skip_all)]
361    pub fn calculate_subtries(&mut self, new_epoch: TrieNodeEpoch) {
362        if let RevealableSparseTrie::Revealed(trie) = &mut self.state {
363            trie.update_subtrie_hashes(new_epoch);
364        }
365    }
366
367    /// Returns storage sparse trie root if the trie has been revealed.
368    pub fn storage_root(&mut self, account: &B256, new_epoch: TrieNodeEpoch) -> Option<B256> {
369        self.storage.tries.get_mut(account).and_then(|trie| trie.root(new_epoch))
370    }
371
372    /// Returns the roots of all revealed storage tries, keyed by hashed account address.
373    ///
374    /// Dirty roots are computed at `new_epoch`; blind storage tries are omitted.
375    pub fn storage_roots(&mut self, new_epoch: TrieNodeEpoch) -> B256Map<B256> {
376        self.storage
377            .tries
378            .iter_mut()
379            .filter_map(|(address, trie)| trie.root(new_epoch).map(|root| (*address, root)))
380            .collect()
381    }
382
383    /// Returns mutable reference to the revealed account sparse trie.
384    fn revealed_trie_mut(&mut self) -> SparseStateTrieResult<&mut A> {
385        self.state.as_revealed_mut().ok_or_else(|| SparseTrieErrorKind::Blind.into())
386    }
387
388    /// Returns sparse trie root.
389    pub fn root(&mut self, new_epoch: TrieNodeEpoch) -> SparseStateTrieResult<B256> {
390        // record revealed node metrics
391        #[cfg(feature = "metrics")]
392        self.metrics.record();
393
394        Ok(self.revealed_trie_mut()?.root(new_epoch))
395    }
396
397    /// Returns sparse trie root and trie updates.
398    ///
399    /// Returns an error if the account trie is still blind.
400    #[instrument(level = "debug", target = "trie::sparse", skip_all)]
401    pub fn root_with_updates(
402        &mut self,
403        new_epoch: TrieNodeEpoch,
404    ) -> SparseStateTrieResult<(B256, TrieUpdates)> {
405        // record revealed node metrics
406        #[cfg(feature = "metrics")]
407        self.metrics.record();
408
409        let storage_tries = self.storage_trie_updates();
410        let revealed = self.revealed_trie_mut()?;
411
412        let (root, updates) = (revealed.root(new_epoch), revealed.take_updates());
413        let updates = TrieUpdates {
414            account_nodes: updates.updated_nodes,
415            removed_nodes: updates.removed_nodes,
416            storage_tries,
417        };
418        Ok((root, updates))
419    }
420
421    /// Returns storage trie updates for tries that have been revealed.
422    ///
423    /// Panics if any of the storage tries are not revealed.
424    pub fn storage_trie_updates(&mut self) -> B256Map<StorageTrieUpdates> {
425        self.storage
426            .tries
427            .iter_mut()
428            .map(|(address, trie)| {
429                let trie = trie.as_revealed_mut().unwrap();
430                let updates = trie.take_updates();
431                let updates = StorageTrieUpdates {
432                    storage_nodes: updates.updated_nodes,
433                    removed_nodes: updates.removed_nodes,
434                };
435                (*address, updates)
436            })
437            .filter(|(_, updates)| !updates.is_empty())
438            .collect()
439    }
440
441    /// Returns [`TrieUpdates`] by taking the updates from the revealed sparse tries.
442    ///
443    /// Returns `None` if the accounts trie is not revealed.
444    pub fn take_trie_updates(&mut self) -> Option<TrieUpdates> {
445        let storage_tries = self.storage_trie_updates();
446        self.state.as_revealed_mut().map(|state| {
447            let updates = state.take_updates();
448            TrieUpdates {
449                account_nodes: updates.updated_nodes,
450                removed_nodes: updates.removed_nodes,
451                storage_tries,
452            }
453        })
454    }
455}
456
457impl<A, S> SparseStateTrie<A, S>
458where
459    A: SparseTrieTrait + Default,
460    S: SparseTrieTrait + Default + Clone,
461{
462    /// Clears all trie data while preserving allocations for reuse.
463    ///
464    /// This resets the trie to an empty state but keeps the underlying memory allocations,
465    /// which can significantly reduce allocation overhead when the trie is reused.
466    pub fn clear(&mut self) {
467        self.state.clear();
468        self.storage.clear();
469    }
470
471    /// Returns the number of storage tries currently retained (active + cleared).
472    pub fn retained_storage_tries_count(&self) -> usize {
473        self.storage.tries.len() + self.storage.cleared_tries.len()
474    }
475
476    /// Prunes account and storage trie nodes last modified before `prune_before`.
477    ///
478    /// Storage tries whose root epochs predate the cutoff are fully evicted.
479    ///
480    /// # Preconditions
481    ///
482    /// Modified account and storage tries must already have computed hashes via `root()` /
483    /// `storage_root()` for their current state. Unmodified storage roots revealed only by
484    /// prewarming are treated as epoch zero.
485    #[cfg(feature = "std")]
486    #[instrument(
487        level = "debug",
488        name = "SparseStateTrie::prune",
489        target = "trie::sparse",
490        skip_all
491    )]
492    pub fn prune(&mut self, prune_before: TrieNodeEpoch) {
493        let total_storage_tries_before = self.storage.tries.len();
494
495        let parent_span = tracing::Span::current();
496        let account_parent_span = parent_span.clone();
497
498        // Prune account and storage tries in parallel using the same epoch cutoff.
499        let (account_nodes_pruned, storage_tries_evicted) = rayon::join(
500            || {
501                let hashed_address = Option::<B256>::None;
502                let _span = tracing::trace_span!(
503                    target: "trie::sparse",
504                    parent: &account_parent_span,
505                    "prune_trie",
506                    ?hashed_address,
507                )
508                .entered();
509
510                self.state.as_revealed_mut().map(|trie| trie.prune(prune_before)).unwrap_or(0)
511            },
512            || self.storage.prune(prune_before, &parent_span),
513        );
514
515        debug!(
516            target: "trie::sparse",
517            prune_before = prune_before.get(),
518            account_nodes_pruned,
519            storage_tries_evicted,
520            storage_tries_after = total_storage_tries_before - storage_tries_evicted,
521            "SparseStateTrie::prune completed"
522        );
523    }
524}
525
526/// The fields of [`SparseStateTrie`] related to storage tries. This is kept separate from the rest
527/// of [`SparseStateTrie`] to help enforce allocation re-use.
528#[derive(Debug, Default)]
529struct StorageTries<S = ArenaParallelSparseTrie> {
530    /// Sparse storage tries.
531    tries: B256Map<RevealableSparseTrie<S>>,
532    /// Cleared storage tries, kept for re-use.
533    cleared_tries: Vec<RevealableSparseTrie<S>>,
534    /// A default cleared trie instance, which will be cloned when creating new tries.
535    default_trie: RevealableSparseTrie<S>,
536}
537
538#[cfg(feature = "std")]
539impl<S: SparseTrieTrait> StorageTries<S> {
540    /// Prunes storage tries by epoch, returning fully old tries to the reuse pool.
541    fn prune(&mut self, prune_before: TrieNodeEpoch, parent_span: &tracing::Span) -> usize {
542        use rayon::iter::{IntoParallelRefMutIterator, ParallelIterator};
543
544        let addresses_to_evict: Vec<B256> = self
545            .tries
546            .par_iter_mut()
547            .filter_map(|(address, trie)| {
548                let hashed_address = Some(*address);
549                let _span = tracing::trace_span!(
550                    target: "trie::sparse",
551                    parent: parent_span,
552                    "prune_trie",
553                    ?hashed_address,
554                )
555                .entered();
556
557                let evict = trie.as_revealed_mut().is_none_or(|revealed| {
558                    // Avoid traversing and compacting a trie that will be immediately evicted.
559                    let root_epoch =
560                        revealed.root_epoch().expect("storage trie root must not be dirty");
561                    if !root_epoch.should_prune(prune_before) {
562                        revealed.prune(prune_before);
563                    }
564                    root_epoch.should_prune(prune_before)
565                });
566
567                evict.then(|| {
568                    trie.clear();
569                    *address
570                })
571            })
572            .collect();
573
574        let evicted = addresses_to_evict.len();
575        self.cleared_tries.reserve(evicted);
576        for address in addresses_to_evict {
577            if let Some(trie) = self.tries.remove(&address) {
578                self.cleared_tries.push(trie);
579            }
580        }
581
582        evicted
583    }
584}
585
586impl<S: SparseTrieTrait> StorageTries<S> {
587    /// Returns all fields to a cleared state, equivalent to the default state, keeping cleared
588    /// collections for re-use later when possible.
589    fn clear(&mut self) {
590        self.cleared_tries.extend(self.tries.drain().map(|(_, mut trie)| {
591            trie.clear();
592            trie
593        }));
594    }
595}
596
597impl<S: SparseTrieTrait + Clone> StorageTries<S> {
598    // Returns mutable reference to storage sparse trie, creating a blind one if it doesn't exist.
599    fn get_or_create_trie_mut(&mut self, address: B256) -> &mut RevealableSparseTrie<S> {
600        self.tries.entry(address).or_insert_with(|| {
601            self.cleared_tries.pop().unwrap_or_else(|| self.default_trie.clone())
602        })
603    }
604}
605
606#[cfg(test)]
607mod tests {
608    use super::*;
609    use crate::{ArenaParallelSparseTrie, LeafLookup, LeafUpdate};
610    use alloy_primitives::{
611        b256,
612        map::{HashMap, HashSet},
613        U256,
614    };
615    use arbitrary::Arbitrary;
616    use rand::{rngs::StdRng, Rng, SeedableRng};
617    use reth_execution_errors::{SparseStateTrieErrorKind, SparseTrieErrorKind};
618    use reth_primitives_traits::Account;
619    use reth_trie::{updates::StorageTrieUpdates, HashBuilder, MultiProof, EMPTY_ROOT_HASH};
620    use reth_trie_common::{
621        proof::{ProofNodes, ProofRetainer},
622        BranchNodeMasks, BranchNodeMasksMap, BranchNodeV2, LeafNode, RlpNode, StorageMultiProof,
623        TrieAccount, TrieMask, TrieNodeV2,
624    };
625
626    const fn epoch(value: u64) -> TrieNodeEpoch {
627        TrieNodeEpoch::new(value)
628    }
629
630    /// Create a leaf key (suffix) with given nibbles padded with zeros to reach `total_len`.
631    fn leaf_key(suffix: impl AsRef<[u8]>, total_len: usize) -> Nibbles {
632        let suffix = suffix.as_ref();
633        let mut nibbles = Nibbles::from_nibbles(suffix);
634        nibbles.extend(&Nibbles::from_nibbles_unchecked(vec![0; total_len - suffix.len()]));
635        nibbles
636    }
637
638    fn apply_account_update(sparse: &mut SparseStateTrie, address: B256, update: LeafUpdate) {
639        let mut updates = B256Map::from_iter([(address, update)]);
640        sparse.trie_mut().update_leaves(&mut updates, |_, _| {}).unwrap();
641        assert!(updates.is_empty());
642    }
643
644    #[test]
645    fn reveal_account_path_twice() {
646        let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
647
648        // Full 64-nibble paths
649        let full_path_0 = leaf_key([0x0], 64);
650        let _full_path_1 = leaf_key([0x1], 64);
651
652        let leaf_value = alloy_rlp::encode(TrieAccount::default());
653        // Leaf key is 63 nibbles (suffix after 1-nibble node path)
654        let leaf_1 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(
655            leaf_key([], 63),
656            leaf_value.clone(),
657        )));
658        let leaf_2 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(
659            leaf_key([], 63),
660            leaf_value.clone(),
661        )));
662
663        let multiproof = MultiProof {
664            account_subtree: ProofNodes::from_iter([
665                (
666                    Nibbles::default(),
667                    alloy_rlp::encode(TrieNodeV2::Branch(BranchNodeV2 {
668                        key: Nibbles::default(),
669                        stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)],
670                        state_mask: TrieMask::new(0b11),
671                        branch_rlp_node: None,
672                    }))
673                    .into(),
674                ),
675                (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()),
676                (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
677            ]),
678            ..Default::default()
679        };
680
681        // Reveal multiproof and check that the state trie contains the leaf node and value
682        sparse.reveal_decoded_multiproof(multiproof.try_into().unwrap()).unwrap();
683        assert!(matches!(
684            sparse.state_trie_ref().unwrap().find_leaf(&full_path_0, None),
685            Ok(LeafLookup::Exists)
686        ));
687        assert_eq!(
688            sparse.state_trie_ref().unwrap().get_leaf_value(&full_path_0),
689            Some(&leaf_value)
690        );
691
692        // Remove the leaf node and check that the state trie does not contain the leaf node and
693        // value
694        apply_account_update(&mut sparse, B256::ZERO, LeafUpdate::Changed(Vec::new()));
695        assert!(matches!(
696            sparse.state_trie_ref().unwrap().find_leaf(&full_path_0, None),
697            Ok(LeafLookup::NonExistent)
698        ));
699        assert!(sparse.state_trie_ref().unwrap().get_leaf_value(&full_path_0).is_none());
700    }
701
702    #[test]
703    fn reveal_storage_path_twice() {
704        let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
705
706        // Full 64-nibble path
707        let full_path_0 = leaf_key([0x0], 64);
708
709        let leaf_value = alloy_rlp::encode(TrieAccount::default());
710        let leaf_1 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(
711            leaf_key([], 63),
712            leaf_value.clone(),
713        )));
714        let leaf_2 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(
715            leaf_key([], 63),
716            leaf_value.clone(),
717        )));
718
719        let multiproof = MultiProof {
720            storages: HashMap::from_iter([(
721                B256::ZERO,
722                StorageMultiProof {
723                    root: B256::ZERO,
724                    subtree: ProofNodes::from_iter([
725                        (
726                            Nibbles::default(),
727                            alloy_rlp::encode(TrieNodeV2::Branch(BranchNodeV2 {
728                                key: Nibbles::default(),
729                                stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)],
730                                state_mask: TrieMask::new(0b11),
731                                branch_rlp_node: None,
732                            }))
733                            .into(),
734                        ),
735                        (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()),
736                        (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
737                    ]),
738                    branch_node_masks: Default::default(),
739                },
740            )]),
741            ..Default::default()
742        };
743
744        // Reveal multiproof and check that the storage trie contains the leaf node and value
745        sparse.reveal_decoded_multiproof(multiproof.try_into().unwrap()).unwrap();
746        assert!(matches!(
747            sparse.storage_trie_ref(&B256::ZERO).unwrap().find_leaf(&full_path_0, None),
748            Ok(LeafLookup::Exists)
749        ));
750        assert_eq!(
751            sparse.storage_trie_ref(&B256::ZERO).unwrap().get_leaf_value(&full_path_0),
752            Some(&leaf_value)
753        );
754
755        // Remove the leaf node and check that the storage trie does not contain the leaf node and
756        // value
757        let mut updates = B256Map::from_iter([(B256::ZERO, LeafUpdate::Changed(Vec::new()))]);
758        sparse
759            .storage_trie_mut(&B256::ZERO)
760            .unwrap()
761            .update_leaves(&mut updates, |_, _| {})
762            .unwrap();
763        assert!(updates.is_empty());
764        assert!(matches!(
765            sparse.storage_trie_ref(&B256::ZERO).unwrap().find_leaf(&full_path_0, None),
766            Ok(LeafLookup::NonExistent)
767        ));
768        assert!(sparse
769            .storage_trie_ref(&B256::ZERO)
770            .unwrap()
771            .get_leaf_value(&full_path_0)
772            .is_none());
773    }
774
775    #[test]
776    fn prune_uses_epochs_for_account_and_storage_tries() {
777        let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
778
779        let account = B256::ZERO;
780        let old_account =
781            b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
782        let prewarmed_account =
783            b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
784        let slot = B256::ZERO;
785        let account_path = leaf_key([0x0], 64);
786        let old_account_path = leaf_key([0x1], 64);
787        let storage_path = leaf_key([0x0], 64);
788
789        let leaf_value = alloy_rlp::encode(TrieAccount::default());
790        let leaf_0 = alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(
791            leaf_key([], 63),
792            leaf_value.clone(),
793        )));
794        let leaf_1 =
795            alloy_rlp::encode(TrieNodeV2::Leaf(LeafNode::new(leaf_key([], 63), leaf_value)));
796
797        let subtree = || {
798            ProofNodes::from_iter([
799                (
800                    Nibbles::default(),
801                    alloy_rlp::encode(TrieNodeV2::Branch(BranchNodeV2 {
802                        key: Nibbles::default(),
803                        stack: vec![RlpNode::from_rlp(&leaf_0), RlpNode::from_rlp(&leaf_1)],
804                        state_mask: TrieMask::new(0b11),
805                        branch_rlp_node: None,
806                    }))
807                    .into(),
808                ),
809                (Nibbles::from_nibbles([0x0]), leaf_0.clone().into()),
810                (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
811            ])
812        };
813
814        let multiproof = MultiProof {
815            account_subtree: subtree(),
816            storages: HashMap::from_iter([
817                (
818                    account,
819                    StorageMultiProof {
820                        root: B256::ZERO,
821                        subtree: subtree(),
822                        branch_node_masks: Default::default(),
823                    },
824                ),
825                (
826                    old_account,
827                    StorageMultiProof {
828                        root: B256::ZERO,
829                        subtree: subtree(),
830                        branch_node_masks: Default::default(),
831                    },
832                ),
833                (
834                    prewarmed_account,
835                    StorageMultiProof {
836                        root: B256::ZERO,
837                        subtree: subtree(),
838                        branch_node_masks: Default::default(),
839                    },
840                ),
841            ]),
842            ..Default::default()
843        };
844
845        sparse.reveal_decoded_multiproof(multiproof.try_into().unwrap()).unwrap();
846
847        sparse.storage_root(&account, epoch(0)).unwrap();
848        sparse.storage_root(&old_account, epoch(0)).unwrap();
849        sparse.root(epoch(0)).unwrap();
850        assert!(!sparse.storage_trie_ref(&prewarmed_account).unwrap().is_root_cached());
851
852        let mut storage_updates = B256Map::from_iter([(
853            slot,
854            LeafUpdate::Changed(alloy_rlp::encode_fixed_size(&U256::from(2)).to_vec()),
855        )]);
856        sparse
857            .storage_trie_mut(&account)
858            .unwrap()
859            .update_leaves(&mut storage_updates, |_, _| {
860                panic!("fully revealed storage trie must not request proofs")
861            })
862            .unwrap();
863        assert!(storage_updates.is_empty());
864
865        let trie_account = TrieAccount {
866            storage_root: sparse.storage_root(&account, epoch(10)).unwrap(),
867            ..Default::default()
868        };
869        apply_account_update(
870            &mut sparse,
871            account,
872            LeafUpdate::Changed(alloy_rlp::encode(trie_account)),
873        );
874        let root_before = sparse.root(epoch(10)).unwrap();
875        sparse.prune(epoch(10));
876
877        assert!(matches!(
878            sparse.state_trie_ref().unwrap().find_leaf(&account_path, None),
879            Ok(LeafLookup::Exists)
880        ));
881        assert!(matches!(
882            sparse.storage_trie_ref(&account).unwrap().find_leaf(&storage_path, None),
883            Ok(LeafLookup::Exists)
884        ));
885        assert!(matches!(
886            sparse.state_trie_ref().unwrap().find_leaf(&old_account_path, None),
887            Err(crate::LeafLookupError::BlindedNode { .. })
888        ));
889        assert!(sparse.storage_trie_ref(&old_account).is_none());
890        assert!(sparse.storage_trie_ref(&prewarmed_account).is_none());
891        assert_eq!(sparse.root(epoch(10)).unwrap(), root_before);
892
893        // An intermediate storage root does not advance without another change, so the trie
894        // becomes eligible once the cutoff passes its last modification epoch.
895        sparse.prune(epoch(11));
896        assert!(sparse.storage_trie_ref(&account).is_none());
897        assert_eq!(sparse.root(epoch(11)).unwrap(), root_before);
898    }
899
900    #[test]
901    fn reveal_v2_proof_nodes() {
902        let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
903
904        // Full 64-nibble path
905        let full_path_0 = leaf_key([0x0], 64);
906
907        let leaf_value = alloy_rlp::encode(TrieAccount::default());
908        let leaf_1_node = TrieNodeV2::Leaf(LeafNode::new(leaf_key([], 63), leaf_value.clone()));
909        let leaf_2_node = TrieNodeV2::Leaf(LeafNode::new(leaf_key([], 63), leaf_value.clone()));
910
911        let branch_node = TrieNodeV2::Branch(BranchNodeV2 {
912            key: Nibbles::default(),
913            stack: vec![
914                RlpNode::from_rlp(&alloy_rlp::encode(&leaf_1_node)),
915                RlpNode::from_rlp(&alloy_rlp::encode(&leaf_2_node)),
916            ],
917            state_mask: TrieMask::new(0b11),
918            branch_rlp_node: None,
919        });
920
921        // Create V2 proof nodes with masks already included
922        let v2_proof_nodes = vec![
923            ProofTrieNodeV2 {
924                path: Nibbles::default(),
925                node: branch_node,
926                masks: Some(BranchNodeMasks {
927                    hash_mask: TrieMask::default(),
928                    tree_mask: TrieMask::default(),
929                }),
930            },
931            ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x0]), node: leaf_1_node, masks: None },
932            ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x1]), node: leaf_2_node, masks: None },
933        ];
934
935        // Reveal V2 proof nodes
936        sparse
937            .reveal_decoded_multiproof_v2(reth_trie_common::DecodedMultiProofV2 {
938                account_proofs: v2_proof_nodes,
939                ..Default::default()
940            })
941            .unwrap();
942
943        // Check that the state trie contains the leaf node and value
944        assert!(matches!(
945            sparse.state_trie_ref().unwrap().find_leaf(&full_path_0, None),
946            Ok(LeafLookup::Exists)
947        ));
948        assert_eq!(
949            sparse.state_trie_ref().unwrap().get_leaf_value(&full_path_0),
950            Some(&leaf_value)
951        );
952
953        // Remove the leaf node
954        apply_account_update(&mut sparse, B256::ZERO, LeafUpdate::Changed(Vec::new()));
955        assert!(sparse.state_trie_ref().unwrap().get_leaf_value(&full_path_0).is_none());
956    }
957
958    #[test]
959    fn reveal_storage_v2_proof_nodes() {
960        let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
961
962        // Full 64-nibble path
963        let full_path_0 = leaf_key([0x0], 64);
964
965        let storage_value: Vec<u8> = alloy_rlp::encode_fixed_size(&U256::from(42)).to_vec();
966        let leaf_1_node = TrieNodeV2::Leaf(LeafNode::new(leaf_key([], 63), storage_value.clone()));
967        let leaf_2_node = TrieNodeV2::Leaf(LeafNode::new(leaf_key([], 63), storage_value.clone()));
968
969        let branch_node = TrieNodeV2::Branch(BranchNodeV2 {
970            key: Nibbles::default(),
971            stack: vec![
972                RlpNode::from_rlp(&alloy_rlp::encode(&leaf_1_node)),
973                RlpNode::from_rlp(&alloy_rlp::encode(&leaf_2_node)),
974            ],
975            state_mask: TrieMask::new(0b11),
976            branch_rlp_node: None,
977        });
978
979        let v2_proof_nodes = vec![
980            ProofTrieNodeV2 { path: Nibbles::default(), node: branch_node, masks: None },
981            ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x0]), node: leaf_1_node, masks: None },
982            ProofTrieNodeV2 { path: Nibbles::from_nibbles([0x1]), node: leaf_2_node, masks: None },
983        ];
984
985        // Reveal V2 storage proof nodes for account
986        sparse
987            .reveal_decoded_multiproof_v2(reth_trie_common::DecodedMultiProofV2 {
988                storage_proofs: B256Map::from_iter([(B256::ZERO, v2_proof_nodes)]),
989                ..Default::default()
990            })
991            .unwrap();
992
993        // Check that the storage trie contains the leaf node and value
994        assert!(matches!(
995            sparse.storage_trie_ref(&B256::ZERO).unwrap().find_leaf(&full_path_0, None),
996            Ok(LeafLookup::Exists)
997        ));
998        assert_eq!(
999            sparse.storage_trie_ref(&B256::ZERO).unwrap().get_leaf_value(&full_path_0),
1000            Some(&storage_value)
1001        );
1002
1003        // Remove the leaf node
1004        let mut updates = B256Map::from_iter([(B256::ZERO, LeafUpdate::Changed(Vec::new()))]);
1005        sparse
1006            .storage_trie_mut(&B256::ZERO)
1007            .unwrap()
1008            .update_leaves(&mut updates, |_, _| {})
1009            .unwrap();
1010        assert!(updates.is_empty());
1011        assert!(sparse
1012            .storage_trie_ref(&B256::ZERO)
1013            .unwrap()
1014            .get_leaf_value(&full_path_0)
1015            .is_none());
1016    }
1017
1018    #[test]
1019    fn root_on_blind_trie_returns_blind_error() {
1020        let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default();
1021
1022        let err = sparse.root(epoch(0)).unwrap_err();
1023
1024        assert!(matches!(err.kind(), SparseStateTrieErrorKind::Sparse(SparseTrieErrorKind::Blind)));
1025    }
1026
1027    #[test]
1028    #[allow(clippy::clone_on_copy)]
1029    fn take_trie_updates() {
1030        reth_tracing::init_test_tracing();
1031
1032        // let mut rng = generators::rng();
1033        let mut rng = StdRng::seed_from_u64(1);
1034
1035        let mut bytes = [0u8; 1024];
1036        rng.fill(bytes.as_mut_slice());
1037
1038        let slot_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
1039        let slot_path_1 = Nibbles::unpack(slot_1);
1040        let value_1 = U256::from(rng.random::<u64>());
1041        let slot_2 = b256!("0x1100000000000000000000000000000000000000000000000000000000000000");
1042        let slot_path_2 = Nibbles::unpack(slot_2);
1043        let value_2 = U256::from(rng.random::<u64>());
1044        let slot_3 = b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
1045        let value_3 = U256::from(rng.random::<u64>());
1046
1047        let mut storage_hash_builder = HashBuilder::default()
1048            .with_proof_retainer(ProofRetainer::from_iter([slot_path_1, slot_path_2]));
1049        storage_hash_builder.add_leaf(slot_path_1, &alloy_rlp::encode_fixed_size(&value_1));
1050        storage_hash_builder.add_leaf(slot_path_2, &alloy_rlp::encode_fixed_size(&value_2));
1051
1052        let storage_root = storage_hash_builder.root();
1053        let storage_proof_nodes = storage_hash_builder.take_proof_nodes();
1054        let storage_branch_node_masks = BranchNodeMasksMap::from_iter([
1055            (
1056                Nibbles::default(),
1057                BranchNodeMasks { hash_mask: TrieMask::new(0b010), tree_mask: TrieMask::default() },
1058            ),
1059            (
1060                Nibbles::from_nibbles([0x1]),
1061                BranchNodeMasks { hash_mask: TrieMask::new(0b11), tree_mask: TrieMask::default() },
1062            ),
1063        ]);
1064
1065        let address_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
1066        let address_path_1 = Nibbles::unpack(address_1);
1067        let account_1 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
1068        let mut trie_account_1 = account_1.clone().into_trie_account(storage_root);
1069        let address_2 = b256!("0x1100000000000000000000000000000000000000000000000000000000000000");
1070        let address_path_2 = Nibbles::unpack(address_2);
1071        let account_2 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
1072        let trie_account_2 = account_2.into_trie_account(EMPTY_ROOT_HASH);
1073
1074        let mut hash_builder = HashBuilder::default()
1075            .with_proof_retainer(ProofRetainer::from_iter([address_path_1, address_path_2]));
1076        hash_builder.add_leaf(address_path_1, &alloy_rlp::encode(trie_account_1.clone()));
1077        hash_builder.add_leaf(address_path_2, &alloy_rlp::encode(trie_account_2));
1078
1079        let root = hash_builder.root();
1080        let proof_nodes = hash_builder.take_proof_nodes();
1081        let mut sparse = SparseStateTrie::<ArenaParallelSparseTrie>::default().with_updates(true);
1082        sparse
1083            .reveal_decoded_multiproof(
1084                MultiProof {
1085                    account_subtree: proof_nodes,
1086                    branch_node_masks: BranchNodeMasksMap::from_iter([(
1087                        Nibbles::from_nibbles([0x1]),
1088                        BranchNodeMasks {
1089                            hash_mask: TrieMask::new(0b00),
1090                            tree_mask: TrieMask::default(),
1091                        },
1092                    )]),
1093                    storages: HashMap::from_iter([
1094                        (
1095                            address_1,
1096                            StorageMultiProof {
1097                                root,
1098                                subtree: storage_proof_nodes.clone(),
1099                                branch_node_masks: storage_branch_node_masks.clone(),
1100                            },
1101                        ),
1102                        (
1103                            address_2,
1104                            StorageMultiProof {
1105                                root,
1106                                subtree: storage_proof_nodes,
1107                                branch_node_masks: storage_branch_node_masks,
1108                            },
1109                        ),
1110                    ]),
1111                }
1112                .try_into()
1113                .unwrap(),
1114            )
1115            .unwrap();
1116
1117        assert_eq!(sparse.root(epoch(0)).unwrap(), root);
1118
1119        let address_3 = b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
1120        let account_3 = Account { nonce: account_1.nonce + 1, ..account_1 };
1121        let trie_account_3 = account_3.into_trie_account(EMPTY_ROOT_HASH);
1122
1123        apply_account_update(
1124            &mut sparse,
1125            address_3,
1126            LeafUpdate::Changed(alloy_rlp::encode(trie_account_3)),
1127        );
1128
1129        let mut updates =
1130            B256Map::from_iter([(slot_3, LeafUpdate::Changed(alloy_rlp::encode(value_3)))]);
1131        sparse
1132            .storage_trie_mut(&address_1)
1133            .unwrap()
1134            .update_leaves(&mut updates, |_, _| {})
1135            .unwrap();
1136        assert!(updates.is_empty());
1137        trie_account_1.storage_root = sparse.storage_root(&address_1, epoch(0)).unwrap();
1138        apply_account_update(
1139            &mut sparse,
1140            address_1,
1141            LeafUpdate::Changed(alloy_rlp::encode(trie_account_1)),
1142        );
1143
1144        sparse.root(epoch(0)).unwrap();
1145
1146        let sparse_updates = sparse.take_trie_updates().unwrap();
1147        // TODO(alexey): assert against real state root calculation updates
1148        pretty_assertions::assert_eq!(
1149            sparse_updates,
1150            TrieUpdates {
1151                account_nodes: HashMap::default(),
1152                storage_tries: HashMap::from_iter([(
1153                    b256!("0x1000000000000000000000000000000000000000000000000000000000000000"),
1154                    StorageTrieUpdates {
1155                        storage_nodes: HashMap::default(),
1156                        removed_nodes: HashSet::from_iter([Nibbles::from_nibbles([0x1])])
1157                    }
1158                )]),
1159                removed_nodes: HashSet::default()
1160            }
1161        );
1162    }
1163}
reth_trie_sparse/state.rs

reth_trie_sparse/
state.rs
