reth_trie_sparse/
state.rs

1use crate::{
2    blinded::{BlindedProvider, BlindedProviderFactory, DefaultBlindedProviderFactory},
3    LeafLookup, RevealedSparseTrie, SparseTrie, TrieMasks,
4};
5use alloc::{collections::VecDeque, vec::Vec};
6use alloy_primitives::{
7    hex,
8    map::{B256Map, HashMap, HashSet},
9    Bytes, B256,
10};
11use alloy_rlp::{Decodable, Encodable};
12use alloy_trie::proof::DecodedProofNodes;
13use core::{fmt, iter::Peekable};
14use reth_execution_errors::{SparseStateTrieErrorKind, SparseStateTrieResult, SparseTrieErrorKind};
15use reth_primitives_traits::Account;
16use reth_trie_common::{
17    proof::ProofNodes,
18    updates::{StorageTrieUpdates, TrieUpdates},
19    DecodedMultiProof, DecodedStorageMultiProof, MultiProof, Nibbles, RlpNode, StorageMultiProof,
20    TrieAccount, TrieMask, TrieNode, EMPTY_ROOT_HASH, TRIE_ACCOUNT_RLP_MAX_SIZE,
21};
22use tracing::trace;
23
24/// Sparse state trie representing lazy-loaded Ethereum state trie.
25pub struct SparseStateTrie<F: BlindedProviderFactory = DefaultBlindedProviderFactory> {
26    /// Blinded node provider factory.
27    provider_factory: F,
28    /// Sparse account trie.
29    state: SparseTrie<F::AccountNodeProvider>,
30    /// Sparse storage tries.
31    storages: B256Map<SparseTrie<F::StorageNodeProvider>>,
32    /// Collection of revealed account trie paths.
33    revealed_account_paths: HashSet<Nibbles>,
34    /// Collection of revealed storage trie paths, per account.
35    revealed_storage_paths: B256Map<HashSet<Nibbles>>,
36    /// Flag indicating whether trie updates should be retained.
37    retain_updates: bool,
38    /// Reusable buffer for RLP encoding of trie accounts.
39    account_rlp_buf: Vec<u8>,
40    /// Metrics for the sparse state trie.
41    #[cfg(feature = "metrics")]
42    metrics: crate::metrics::SparseStateTrieMetrics,
43}
44
45#[cfg(test)]
46impl Default for SparseStateTrie {
47    fn default() -> Self {
48        Self {
49            provider_factory: DefaultBlindedProviderFactory,
50            state: Default::default(),
51            storages: Default::default(),
52            revealed_account_paths: Default::default(),
53            revealed_storage_paths: Default::default(),
54            retain_updates: false,
55            account_rlp_buf: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
56            #[cfg(feature = "metrics")]
57            metrics: Default::default(),
58        }
59    }
60}
61
62impl<P: BlindedProviderFactory> fmt::Debug for SparseStateTrie<P> {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        f.debug_struct("SparseStateTrie")
65            .field("state", &self.state)
66            .field("storages", &self.storages)
67            .field("revealed_account_paths", &self.revealed_account_paths)
68            .field("revealed_storage_paths", &self.revealed_storage_paths)
69            .field("retain_updates", &self.retain_updates)
70            .field("account_rlp_buf", &hex::encode(&self.account_rlp_buf))
71            .finish_non_exhaustive()
72    }
73}
74
75#[cfg(test)]
76impl SparseStateTrie {
77    /// Create state trie from state trie.
78    pub fn from_state(state: SparseTrie) -> Self {
79        Self { state, ..Default::default() }
80    }
81}
82
83impl<F: BlindedProviderFactory> SparseStateTrie<F> {
84    /// Create new [`SparseStateTrie`] with blinded node provider factory.
85    pub fn new(provider_factory: F) -> Self {
86        Self {
87            provider_factory,
88            state: Default::default(),
89            storages: Default::default(),
90            revealed_account_paths: Default::default(),
91            revealed_storage_paths: Default::default(),
92            retain_updates: false,
93            account_rlp_buf: Vec::with_capacity(TRIE_ACCOUNT_RLP_MAX_SIZE),
94            #[cfg(feature = "metrics")]
95            metrics: Default::default(),
96        }
97    }
98
99    /// Set the retention of branch node updates and deletions.
100    pub const fn with_updates(mut self, retain_updates: bool) -> Self {
101        self.retain_updates = retain_updates;
102        self
103    }
104
105    /// Returns `true` if account was already revealed.
106    pub fn is_account_revealed(&self, account: B256) -> bool {
107        self.revealed_account_paths.contains(&Nibbles::unpack(account))
108    }
109
110    /// Was the account witness for `address` complete?
111    pub fn check_valid_account_witness(&self, address: B256) -> bool {
112        let path = Nibbles::unpack(address);
113        let trie = match self.state_trie_ref() {
114            Some(t) => t,
115            None => return false,
116        };
117
118        matches!(
119            trie.find_leaf(&path, None),
120            Ok(LeafLookup::Exists | LeafLookup::NonExistent { .. })
121        )
122    }
123
124    /// Was the storage-slot witness for (`address`,`slot`) complete?
125    pub fn check_valid_storage_witness(&self, address: B256, slot: B256) -> bool {
126        let path = Nibbles::unpack(slot);
127        let trie = match self.storage_trie_ref(&address) {
128            Some(t) => t,
129            None => return false,
130        };
131
132        matches!(
133            trie.find_leaf(&path, None),
134            Ok(LeafLookup::Exists | LeafLookup::NonExistent { .. })
135        )
136    }
137
138    /// Returns `true` if storage slot for account was already revealed.
139    pub fn is_storage_slot_revealed(&self, account: B256, slot: B256) -> bool {
140        self.revealed_storage_paths
141            .get(&account)
142            .is_some_and(|slots| slots.contains(&Nibbles::unpack(slot)))
143    }
144
145    /// Returns reference to bytes representing leaf value for the target account.
146    pub fn get_account_value(&self, account: &B256) -> Option<&Vec<u8>> {
147        self.state.as_revealed_ref()?.get_leaf_value(&Nibbles::unpack(account))
148    }
149
150    /// Returns reference to bytes representing leaf value for the target account and storage slot.
151    pub fn get_storage_slot_value(&self, account: &B256, slot: &B256) -> Option<&Vec<u8>> {
152        self.storages.get(account)?.as_revealed_ref()?.get_leaf_value(&Nibbles::unpack(slot))
153    }
154
155    /// Returns reference to state trie if it was revealed.
156    pub const fn state_trie_ref(&self) -> Option<&RevealedSparseTrie<F::AccountNodeProvider>> {
157        self.state.as_revealed_ref()
158    }
159
160    /// Returns reference to storage trie if it was revealed.
161    pub fn storage_trie_ref(
162        &self,
163        address: &B256,
164    ) -> Option<&RevealedSparseTrie<F::StorageNodeProvider>> {
165        self.storages.get(address).and_then(|e| e.as_revealed_ref())
166    }
167
168    /// Returns mutable reference to storage sparse trie if it was revealed.
169    pub fn storage_trie_mut(
170        &mut self,
171        address: &B256,
172    ) -> Option<&mut RevealedSparseTrie<F::StorageNodeProvider>> {
173        self.storages.get_mut(address).and_then(|e| e.as_revealed_mut())
174    }
175
176    /// Takes the storage trie for the provided address.
177    pub fn take_storage_trie(
178        &mut self,
179        address: &B256,
180    ) -> Option<SparseTrie<F::StorageNodeProvider>> {
181        self.storages.remove(address)
182    }
183
184    /// Inserts storage trie for the provided address.
185    pub fn insert_storage_trie(
186        &mut self,
187        address: B256,
188        storage_trie: SparseTrie<F::StorageNodeProvider>,
189    ) {
190        self.storages.insert(address, storage_trie);
191    }
192
193    /// Reveal unknown trie paths from provided leaf path and its proof for the account.
194    ///
195    /// Panics if trie updates retention is enabled.
196    ///
197    /// NOTE: This method does not extensively validate the proof.
198    pub fn reveal_account(
199        &mut self,
200        account: B256,
201        proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
202    ) -> SparseStateTrieResult<()> {
203        assert!(!self.retain_updates);
204
205        if self.is_account_revealed(account) {
206            return Ok(());
207        }
208
209        let mut proof = proof.into_iter().peekable();
210
211        let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) };
212
213        // Reveal root node if it wasn't already.
214        let trie = self.state.reveal_root_with_provider(
215            self.provider_factory.account_node_provider(),
216            root_node,
217            TrieMasks::none(),
218            self.retain_updates,
219        )?;
220
221        // Reveal the remaining proof nodes.
222        for (path, bytes) in proof {
223            if self.revealed_account_paths.contains(&path) {
224                continue;
225            }
226            let node = TrieNode::decode(&mut &bytes[..])?;
227            trie.reveal_node(path.clone(), node, TrieMasks::none())?;
228
229            // Track the revealed path.
230            self.revealed_account_paths.insert(path);
231        }
232
233        Ok(())
234    }
235
236    /// Reveal unknown trie paths from provided leaf path and its proof for the storage slot.
237    ///
238    /// Panics if trie updates retention is enabled.
239    ///
240    /// NOTE: This method does not extensively validate the proof.
241    pub fn reveal_storage_slot(
242        &mut self,
243        account: B256,
244        slot: B256,
245        proof: impl IntoIterator<Item = (Nibbles, Bytes)>,
246    ) -> SparseStateTrieResult<()> {
247        assert!(!self.retain_updates);
248
249        if self.is_storage_slot_revealed(account, slot) {
250            return Ok(());
251        }
252
253        let mut proof = proof.into_iter().peekable();
254
255        let Some(root_node) = self.validate_root_node(&mut proof)? else { return Ok(()) };
256
257        // Reveal root node if it wasn't already.
258        let trie = self.storages.entry(account).or_default().reveal_root_with_provider(
259            self.provider_factory.storage_node_provider(account),
260            root_node,
261            TrieMasks::none(),
262            self.retain_updates,
263        )?;
264
265        let revealed_nodes = self.revealed_storage_paths.entry(account).or_default();
266
267        // Reveal the remaining proof nodes.
268        for (path, bytes) in proof {
269            // If the node is already revealed, skip it.
270            if revealed_nodes.contains(&path) {
271                continue;
272            }
273            let node = TrieNode::decode(&mut &bytes[..])?;
274            trie.reveal_node(path.clone(), node, TrieMasks::none())?;
275
276            // Track the revealed path.
277            revealed_nodes.insert(path);
278        }
279
280        Ok(())
281    }
282
283    /// Reveal unknown trie paths from multiproof.
284    /// NOTE: This method does not extensively validate the proof.
285    pub fn reveal_multiproof(&mut self, multiproof: MultiProof) -> SparseStateTrieResult<()> {
286        // first decode the multiproof
287        let decoded_multiproof = multiproof.try_into()?;
288
289        // then reveal the decoded multiproof
290        self.reveal_decoded_multiproof(decoded_multiproof)
291    }
292
293    /// Reveal unknown trie paths from decoded multiproof.
294    /// NOTE: This method does not extensively validate the proof.
295    pub fn reveal_decoded_multiproof(
296        &mut self,
297        multiproof: DecodedMultiProof,
298    ) -> SparseStateTrieResult<()> {
299        let DecodedMultiProof {
300            account_subtree,
301            storages,
302            branch_node_hash_masks,
303            branch_node_tree_masks,
304        } = multiproof;
305
306        // first reveal the account proof nodes
307        self.reveal_decoded_account_multiproof(
308            account_subtree,
309            branch_node_hash_masks,
310            branch_node_tree_masks,
311        )?;
312
313        // then reveal storage proof nodes for each storage trie
314        for (account, storage_subtree) in storages {
315            self.reveal_decoded_storage_multiproof(account, storage_subtree)?;
316        }
317
318        Ok(())
319    }
320
321    /// Reveals an account multiproof.
322    pub fn reveal_account_multiproof(
323        &mut self,
324        account_subtree: ProofNodes,
325        branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
326        branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
327    ) -> SparseStateTrieResult<()> {
328        // decode the multiproof first
329        let decoded_multiproof = account_subtree.try_into()?;
330        self.reveal_decoded_account_multiproof(
331            decoded_multiproof,
332            branch_node_hash_masks,
333            branch_node_tree_masks,
334        )
335    }
336
337    /// Reveals a decoded account multiproof.
338    pub fn reveal_decoded_account_multiproof(
339        &mut self,
340        account_subtree: DecodedProofNodes,
341        branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
342        branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
343    ) -> SparseStateTrieResult<()> {
344        let FilteredProofNodes {
345            nodes,
346            new_nodes,
347            total_nodes: _total_nodes,
348            skipped_nodes: _skipped_nodes,
349        } = filter_revealed_nodes(account_subtree, &self.revealed_account_paths)?;
350        #[cfg(feature = "metrics")]
351        {
352            self.metrics.increment_total_account_nodes(_total_nodes as u64);
353            self.metrics.increment_skipped_account_nodes(_skipped_nodes as u64);
354        }
355        let mut account_nodes = nodes.into_iter().peekable();
356
357        if let Some(root_node) = Self::validate_root_node_decoded(&mut account_nodes)? {
358            // Reveal root node if it wasn't already.
359            let trie = self.state.reveal_root_with_provider(
360                self.provider_factory.account_node_provider(),
361                root_node,
362                TrieMasks {
363                    hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(),
364                    tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(),
365                },
366                self.retain_updates,
367            )?;
368
369            // Reserve the capacity for new nodes ahead of time.
370            trie.reserve_nodes(new_nodes);
371
372            // Reveal the remaining proof nodes.
373            for (path, node) in account_nodes {
374                // // seismic upstream merge: unsure if this section is necessary. kept during
375                // self.metrics.increment_total_account_nodes();
376                // // If the node is already revealed, skip it.
377                // if self.revealed_account_paths.contains(&path) {
378                //     self.metrics.increment_skipped_account_nodes();
379                //     continue;
380                // }
381
382                let (hash_mask, tree_mask) = if let TrieNode::Branch(_) = node {
383                    (
384                        branch_node_hash_masks.get(&path).copied(),
385                        branch_node_tree_masks.get(&path).copied(),
386                    )
387                } else {
388                    (None, None)
389                };
390
391                trace!(target: "trie::sparse", ?path, ?node, ?hash_mask, ?tree_mask, "Revealing account node");
392                trie.reveal_node(path.clone(), node, TrieMasks { hash_mask, tree_mask })?;
393
394                // Track the revealed path.
395                self.revealed_account_paths.insert(path);
396            }
397        }
398
399        Ok(())
400    }
401
402    /// Reveals a storage multiproof for the given address.
403    pub fn reveal_storage_multiproof(
404        &mut self,
405        account: B256,
406        storage_subtree: StorageMultiProof,
407    ) -> SparseStateTrieResult<()> {
408        // decode the multiproof first
409        let decoded_multiproof = storage_subtree.try_into()?;
410        self.reveal_decoded_storage_multiproof(account, decoded_multiproof)
411    }
412
413    /// Reveals a decoded storage multiproof for the given address.
414    pub fn reveal_decoded_storage_multiproof(
415        &mut self,
416        account: B256,
417        storage_subtree: DecodedStorageMultiProof,
418    ) -> SparseStateTrieResult<()> {
419        let revealed_nodes = self.revealed_storage_paths.entry(account).or_default();
420
421        let FilteredProofNodes {
422            nodes,
423            new_nodes,
424            total_nodes: _total_nodes,
425            skipped_nodes: _skipped_nodes,
426        } = filter_revealed_nodes(storage_subtree.subtree, revealed_nodes)?;
427        #[cfg(feature = "metrics")]
428        {
429            self.metrics.increment_total_storage_nodes(_total_nodes as u64);
430            self.metrics.increment_skipped_storage_nodes(_skipped_nodes as u64);
431        }
432        let mut nodes = nodes.into_iter().peekable();
433
434        if let Some(root_node) = Self::validate_root_node_decoded(&mut nodes)? {
435            // Reveal root node if it wasn't already.
436            let trie = self.storages.entry(account).or_default().reveal_root_with_provider(
437                self.provider_factory.storage_node_provider(account),
438                root_node,
439                TrieMasks {
440                    hash_mask: storage_subtree
441                        .branch_node_hash_masks
442                        .get(&Nibbles::default())
443                        .copied(),
444                    tree_mask: storage_subtree
445                        .branch_node_tree_masks
446                        .get(&Nibbles::default())
447                        .copied(),
448                },
449                self.retain_updates,
450            )?;
451
452            // Reserve the capacity for new nodes ahead of time.
453            trie.reserve_nodes(new_nodes);
454
455            // Reveal the remaining proof nodes.
456            for (path, node) in nodes {
457                // // seismic upstream merge: unsure if this section is necessary. kept during
458                // self.metrics.increment_total_storage_nodes();
459                // // If the node is already revealed, skip it.
460                // if revealed_nodes.contains(&path) {
461                //     self.metrics.increment_skipped_storage_nodes();
462                //     continue;
463                // }
464
465                let (hash_mask, tree_mask) = if let TrieNode::Branch(_) = node {
466                    (
467                        storage_subtree.branch_node_hash_masks.get(&path).copied(),
468                        storage_subtree.branch_node_tree_masks.get(&path).copied(),
469                    )
470                } else {
471                    (None, None)
472                };
473
474                trace!(target: "trie::sparse", ?account, ?path, ?node, ?hash_mask, ?tree_mask, "Revealing storage node");
475                trie.reveal_node(path.clone(), node, TrieMasks { hash_mask, tree_mask })?;
476
477                // Track the revealed path.
478                revealed_nodes.insert(path);
479            }
480        }
481
482        Ok(())
483    }
484
485    /// Reveal state witness with the given state root.
486    /// The state witness is expected to be a map of `keccak(rlp(node)): rlp(node).`
487    /// NOTE: This method does not extensively validate the witness.
488    pub fn reveal_witness(
489        &mut self,
490        state_root: B256,
491        witness: &B256Map<Bytes>,
492    ) -> SparseStateTrieResult<()> {
493        // Create a `(hash, path, maybe_account)` queue for traversing witness trie nodes
494        // starting from the root node.
495        let mut queue = VecDeque::from([(state_root, Nibbles::default(), None)]);
496
497        while let Some((hash, path, maybe_account)) = queue.pop_front() {
498            // Retrieve the trie node and decode it.
499            let Some(trie_node_bytes) = witness.get(&hash) else { continue };
500            let trie_node = TrieNode::decode(&mut &trie_node_bytes[..])?;
501
502            // Push children nodes into the queue.
503            match &trie_node {
504                TrieNode::Branch(branch) => {
505                    for (idx, maybe_child) in branch.as_ref().children() {
506                        if let Some(child_hash) = maybe_child.and_then(RlpNode::as_hash) {
507                            let mut child_path = path.clone();
508                            child_path.push_unchecked(idx);
509                            queue.push_back((child_hash, child_path, maybe_account));
510                        }
511                    }
512                }
513                TrieNode::Extension(ext) => {
514                    if let Some(child_hash) = ext.child.as_hash() {
515                        let mut child_path = path.clone();
516                        child_path.extend_from_slice_unchecked(&ext.key);
517                        queue.push_back((child_hash, child_path, maybe_account));
518                    }
519                }
520                TrieNode::Leaf(leaf) => {
521                    let mut full_path = path.clone();
522                    full_path.extend_from_slice_unchecked(&leaf.key);
523                    if maybe_account.is_none() {
524                        let hashed_address = B256::from_slice(&full_path.pack());
525                        let account = TrieAccount::decode(&mut &leaf.value[..])?;
526                        if account.storage_root != EMPTY_ROOT_HASH {
527                            queue.push_back((
528                                account.storage_root,
529                                Nibbles::default(),
530                                Some(hashed_address),
531                            ));
532                        }
533                    }
534                }
535                TrieNode::EmptyRoot => {} // nothing to do here
536            };
537
538            // Reveal the node itself.
539            if let Some(account) = maybe_account {
540                // Check that the path was not already revealed.
541                if self
542                    .revealed_storage_paths
543                    .get(&account)
544                    .is_none_or(|paths| !paths.contains(&path))
545                {
546                    let storage_trie_entry = self.storages.entry(account).or_default();
547                    if path.is_empty() {
548                        // Handle special storage state root node case.
549                        storage_trie_entry.reveal_root_with_provider(
550                            self.provider_factory.storage_node_provider(account),
551                            trie_node,
552                            TrieMasks::none(),
553                            self.retain_updates,
554                        )?;
555                    } else {
556                        // Reveal non-root storage trie node.
557                        storage_trie_entry
558                            .as_revealed_mut()
559                            .ok_or(SparseTrieErrorKind::Blind)?
560                            .reveal_node(path.clone(), trie_node, TrieMasks::none())?;
561                    }
562
563                    // Track the revealed path.
564                    self.revealed_storage_paths.entry(account).or_default().insert(path);
565                }
566            }
567            // Check that the path was not already revealed.
568            else if !self.revealed_account_paths.contains(&path) {
569                if path.is_empty() {
570                    // Handle special state root node case.
571                    self.state.reveal_root_with_provider(
572                        self.provider_factory.account_node_provider(),
573                        trie_node,
574                        TrieMasks::none(),
575                        self.retain_updates,
576                    )?;
577                } else {
578                    // Reveal non-root state trie node.
579                    self.state.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?.reveal_node(
580                        path.clone(),
581                        trie_node,
582                        TrieMasks::none(),
583                    )?;
584                }
585
586                // Track the revealed path.
587                self.revealed_account_paths.insert(path);
588            }
589        }
590
591        Ok(())
592    }
593
594    /// Validates the root node of the proof and returns it if it exists and is valid.
595    fn validate_root_node<I: Iterator<Item = (Nibbles, Bytes)>>(
596        &self,
597        proof: &mut Peekable<I>,
598    ) -> SparseStateTrieResult<Option<TrieNode>> {
599        // Validate root node.
600        let Some((path, node)) = proof.next() else { return Ok(None) };
601        if !path.is_empty() {
602            return Err(SparseStateTrieErrorKind::InvalidRootNode { path, node }.into());
603        }
604
605        // Decode root node and perform sanity check.
606        let root_node = TrieNode::decode(&mut &node[..])?;
607        if matches!(root_node, TrieNode::EmptyRoot) && proof.peek().is_some() {
608            return Err(SparseStateTrieErrorKind::InvalidRootNode { path, node }.into());
609        }
610
611        Ok(Some(root_node))
612    }
613
614    /// Validates the decoded root node of the proof and returns it if it exists and is valid.
615    fn validate_root_node_decoded<I: Iterator<Item = (Nibbles, TrieNode)>>(
616        proof: &mut Peekable<I>,
617    ) -> SparseStateTrieResult<Option<TrieNode>> {
618        // Validate root node.
619        let Some((path, root_node)) = proof.next() else { return Ok(None) };
620        if !path.is_empty() {
621            return Err(SparseStateTrieErrorKind::InvalidRootNode {
622                path,
623                node: alloy_rlp::encode(&root_node).into(),
624            }
625            .into())
626        }
627
628        // Perform sanity check.
629        if matches!(root_node, TrieNode::EmptyRoot) && proof.peek().is_some() {
630            return Err(SparseStateTrieErrorKind::InvalidRootNode {
631                path,
632                node: alloy_rlp::encode(&root_node).into(),
633            }
634            .into())
635        }
636
637        Ok(Some(root_node))
638    }
639
640    /// Wipe the storage trie at the provided address.
641    pub fn wipe_storage(&mut self, address: B256) -> SparseStateTrieResult<()> {
642        if let Some(trie) = self.storages.get_mut(&address) {
643            trie.wipe()?;
644        }
645        Ok(())
646    }
647
648    /// Calculates the hashes of the nodes below the provided level.
649    ///
650    /// If the trie has not been revealed, this function does nothing.
651    pub fn calculate_below_level(&mut self, level: usize) {
652        if let SparseTrie::Revealed(trie) = &mut self.state {
653            trie.update_rlp_node_level(level);
654        }
655    }
656
657    /// Returns storage sparse trie root if the trie has been revealed.
658    pub fn storage_root(&mut self, account: B256) -> Option<B256> {
659        self.storages.get_mut(&account).and_then(|trie| trie.root())
660    }
661
662    /// Returns mutable reference to the revealed sparse trie.
663    ///
664    /// If the trie is not revealed yet, its root will be revealed using the blinded node provider.
665    fn revealed_trie_mut(
666        &mut self,
667    ) -> SparseStateTrieResult<&mut RevealedSparseTrie<F::AccountNodeProvider>> {
668        match self.state {
669            SparseTrie::Blind => {
670                let (root_node, hash_mask, tree_mask) = self
671                    .provider_factory
672                    .account_node_provider()
673                    .blinded_node(&Nibbles::default())?
674                    .map(|node| {
675                        TrieNode::decode(&mut &node.node[..])
676                            .map(|decoded| (decoded, node.hash_mask, node.tree_mask))
677                    })
678                    .transpose()?
679                    .unwrap_or((TrieNode::EmptyRoot, None, None));
680                self.state
681                    .reveal_root_with_provider(
682                        self.provider_factory.account_node_provider(),
683                        root_node,
684                        TrieMasks { hash_mask, tree_mask },
685                        self.retain_updates,
686                    )
687                    .map_err(Into::into)
688            }
689            SparseTrie::Revealed(ref mut trie) => Ok(trie),
690        }
691    }
692
693    /// Returns sparse trie root.
694    ///
695    /// If the trie has not been revealed, this function reveals the root node and returns its hash.
696    pub fn root(&mut self) -> SparseStateTrieResult<B256> {
697        // record revealed node metrics
698        #[cfg(feature = "metrics")]
699        self.metrics.record();
700
701        Ok(self.revealed_trie_mut()?.root())
702    }
703
704    /// Returns sparse trie root and trie updates if the trie has been revealed.
705    pub fn root_with_updates(&mut self) -> SparseStateTrieResult<(B256, TrieUpdates)> {
706        // record revealed node metrics
707        #[cfg(feature = "metrics")]
708        self.metrics.record();
709
710        let storage_tries = self.storage_trie_updates();
711        let revealed = self.revealed_trie_mut()?;
712
713        let (root, updates) = (revealed.root(), revealed.take_updates());
714        let updates = TrieUpdates {
715            account_nodes: updates.updated_nodes,
716            removed_nodes: updates.removed_nodes,
717            storage_tries,
718        };
719        Ok((root, updates))
720    }
721
722    /// Returns storage trie updates for tries that have been revealed.
723    ///
724    /// Panics if any of the storage tries are not revealed.
725    pub fn storage_trie_updates(&mut self) -> B256Map<StorageTrieUpdates> {
726        self.storages
727            .iter_mut()
728            .map(|(address, trie)| {
729                let trie = trie.as_revealed_mut().unwrap();
730                let updates = trie.take_updates();
731                let updates = StorageTrieUpdates {
732                    is_deleted: updates.wiped,
733                    storage_nodes: updates.updated_nodes,
734                    removed_nodes: updates.removed_nodes,
735                };
736                (*address, updates)
737            })
738            .filter(|(_, updates)| !updates.is_empty())
739            .collect()
740    }
741
742    /// Returns [`TrieUpdates`] by taking the updates from the revealed sparse tries.
743    ///
744    /// Returns `None` if the accounts trie is not revealed.
745    pub fn take_trie_updates(&mut self) -> Option<TrieUpdates> {
746        let storage_tries = self.storage_trie_updates();
747        self.state.as_revealed_mut().map(|state| {
748            let updates = state.take_updates();
749            TrieUpdates {
750                account_nodes: updates.updated_nodes,
751                removed_nodes: updates.removed_nodes,
752                storage_tries,
753            }
754        })
755    }
756
757    /// Update the account leaf node.
758    pub fn update_account_leaf(
759        &mut self,
760        path: Nibbles,
761        value: Vec<u8>,
762    ) -> SparseStateTrieResult<()> {
763        if !self.revealed_account_paths.contains(&path) {
764            self.revealed_account_paths.insert(path.clone());
765        }
766        let is_private = false; // account leaves are always public. Their storage leaves can be private.
767
768        self.state.update_leaf(path, value, is_private)?;
769        Ok(())
770    }
771
772    /// Update the leaf node of a storage trie at the provided address.
773    pub fn update_storage_leaf(
774        &mut self,
775        address: B256,
776        slot: Nibbles,
777        value: Vec<u8>,
778        is_private: bool,
779    ) -> SparseStateTrieResult<()> {
780        if !self.revealed_storage_paths.get(&address).is_some_and(|slots| slots.contains(&slot)) {
781            self.revealed_storage_paths.entry(address).or_default().insert(slot.clone());
782        }
783
784        let storage_trie = self.storages.get_mut(&address).ok_or(SparseTrieErrorKind::Blind)?;
785        storage_trie.update_leaf(slot, value, is_private)?;
786        Ok(())
787    }
788
789    /// Update or remove trie account based on new account info. This method will either recompute
790    /// the storage root based on update storage trie or look it up from existing leaf value.
791    ///
792    /// If the new account info and storage trie are empty, the account leaf will be removed.
793    pub fn update_account(&mut self, address: B256, account: Account) -> SparseStateTrieResult<()> {
794        let nibbles = Nibbles::unpack(address);
795
796        let storage_root = if let Some(storage_trie) = self.storages.get_mut(&address) {
797            trace!(target: "trie::sparse", ?address, "Calculating storage root to update account");
798            storage_trie.root().ok_or(SparseTrieErrorKind::Blind)?
799        } else if self.is_account_revealed(address) {
800            trace!(target: "trie::sparse", ?address, "Retrieving storage root from account leaf to update account");
801            // The account was revealed, either...
802            if let Some(value) = self.get_account_value(&address) {
803                // ..it exists and we should take it's current storage root or...
804                TrieAccount::decode(&mut &value[..])?.storage_root
805            } else {
806                // ...the account is newly created and the storage trie is empty.
807                EMPTY_ROOT_HASH
808            }
809        } else {
810            return Err(SparseTrieErrorKind::Blind.into());
811        };
812
813        if account.is_empty() && storage_root == EMPTY_ROOT_HASH {
814            trace!(target: "trie::sparse", ?address, "Removing account");
815            self.remove_account_leaf(&nibbles)
816        } else {
817            trace!(target: "trie::sparse", ?address, "Updating account");
818            self.account_rlp_buf.clear();
819            account.into_trie_account(storage_root).encode(&mut self.account_rlp_buf);
820            self.update_account_leaf(nibbles, self.account_rlp_buf.clone())
821        }
822    }
823
824    /// Update the storage root of a revealed account.
825    ///
826    /// If the account doesn't exist in the trie, the function is a no-op.
827    ///
828    /// If the new storage root is empty, and the account info was already empty, the account leaf
829    /// will be removed.
830    pub fn update_account_storage_root(&mut self, address: B256) -> SparseStateTrieResult<()> {
831        if !self.is_account_revealed(address) {
832            return Err(SparseTrieErrorKind::Blind.into());
833        }
834
835        // Nothing to update if the account doesn't exist in the trie.
836        let Some(mut trie_account) = self
837            .get_account_value(&address)
838            .map(|v| TrieAccount::decode(&mut &v[..]))
839            .transpose()?
840        else {
841            trace!(target: "trie::sparse", ?address, "Account not found in trie, skipping storage root update");
842            return Ok(());
843        };
844
845        // Calculate the new storage root. If the storage trie doesn't exist, the storage root will
846        // be empty.
847        let storage_root = if let Some(storage_trie) = self.storages.get_mut(&address) {
848            trace!(target: "trie::sparse", ?address, "Calculating storage root to update account");
849            storage_trie.root().ok_or(SparseTrieErrorKind::Blind)?
850        } else {
851            EMPTY_ROOT_HASH
852        };
853
854        // Update the account with the new storage root.
855        trie_account.storage_root = storage_root;
856
857        let nibbles = Nibbles::unpack(address);
858        if trie_account == TrieAccount::default() {
859            // If the account is empty, remove it.
860            trace!(target: "trie::sparse", ?address, "Removing account because the storage root is empty");
861            self.remove_account_leaf(&nibbles)?;
862        } else {
863            // Otherwise, update the account leaf.
864            trace!(target: "trie::sparse", ?address, "Updating account with the new storage root");
865            self.account_rlp_buf.clear();
866            trie_account.encode(&mut self.account_rlp_buf);
867            self.update_account_leaf(nibbles, self.account_rlp_buf.clone())?;
868        }
869
870        Ok(())
871    }
872
873    /// Remove the account leaf node.
874    pub fn remove_account_leaf(&mut self, path: &Nibbles) -> SparseStateTrieResult<()> {
875        self.state.remove_leaf(path)?;
876        Ok(())
877    }
878
879    /// Update the leaf node of a storage trie at the provided address.
880    pub fn remove_storage_leaf(
881        &mut self,
882        address: B256,
883        slot: &Nibbles,
884    ) -> SparseStateTrieResult<()> {
885        let storage_trie = self.storages.get_mut(&address).ok_or(SparseTrieErrorKind::Blind)?;
886        storage_trie.remove_leaf(slot)?;
887        Ok(())
888    }
889}
890
891/// Result of [`filter_revealed_nodes`].
892#[derive(Debug, PartialEq, Eq)]
893struct FilteredProofNodes {
894    /// Filtered, decoded and sorted proof nodes.
895    nodes: Vec<(Nibbles, TrieNode)>,
896    /// Number of nodes in the proof.
897    total_nodes: usize,
898    /// Number of nodes that were skipped because they were already revealed.
899    skipped_nodes: usize,
900    /// Number of new nodes that will be revealed. This includes all children of branch nodes, even
901    /// if they are not in the proof.
902    new_nodes: usize,
903}
904
905/// Filters the decoded nodes that are already revealed and returns additional information about the
906/// number of total, skipped, and new nodes.
907fn filter_revealed_nodes(
908    proof_nodes: DecodedProofNodes,
909    revealed_nodes: &HashSet<Nibbles>,
910) -> alloy_rlp::Result<FilteredProofNodes> {
911    let mut result = FilteredProofNodes {
912        nodes: Vec::with_capacity(proof_nodes.len()),
913        total_nodes: 0,
914        skipped_nodes: 0,
915        new_nodes: 0,
916    };
917
918    for (path, node) in proof_nodes.into_inner() {
919        result.total_nodes += 1;
920        // If the node is already revealed, skip it.
921        if revealed_nodes.contains(&path) {
922            result.skipped_nodes += 1;
923            continue
924        }
925
926        result.new_nodes += 1;
927        // If it's a branch node, increase the number of new nodes by the number of children
928        // according to the state mask.
929        if let TrieNode::Branch(branch) = &node {
930            result.new_nodes += branch.state_mask.count_ones() as usize;
931        }
932
933        result.nodes.push((path, node));
934    }
935
936    result.nodes.sort_unstable_by(|a, b| a.0.cmp(&b.0));
937    Ok(result)
938}
939
940#[cfg(test)]
941mod tests {
942    use super::*;
943    use alloy_primitives::{
944        b256,
945        map::{HashMap, HashSet},
946        Bytes, U256,
947    };
948    use alloy_rlp::EMPTY_STRING_CODE;
949    use arbitrary::Arbitrary;
950    use assert_matches::assert_matches;
951    use rand::{rngs::StdRng, Rng, SeedableRng};
952    use reth_primitives_traits::Account;
953    use reth_trie::{updates::StorageTrieUpdates, HashBuilder, MultiProof, EMPTY_ROOT_HASH};
954    use reth_trie_common::{
955        proof::{ProofNodes, ProofRetainer},
956        BranchNode, LeafNode, StorageMultiProof, TrieMask,
957    };
958
959    #[test]
960    fn validate_root_node_first_node_not_root() {
961        let sparse = SparseStateTrie::default();
962        let proof = [(Nibbles::from_nibbles([0x1]), Bytes::from([EMPTY_STRING_CODE]))];
963        assert_matches!(
964            sparse.validate_root_node(&mut proof.into_iter().peekable()).map_err(|e| e.into_kind()),
965            Err(SparseStateTrieErrorKind::InvalidRootNode { .. })
966        );
967    }
968
969    #[test]
970    fn validate_root_node_invalid_proof_with_empty_root() {
971        let sparse = SparseStateTrie::default();
972        let proof = [
973            (Nibbles::default(), Bytes::from([EMPTY_STRING_CODE])),
974            (Nibbles::from_nibbles([0x1]), Bytes::new()),
975        ];
976        assert_matches!(
977            sparse.validate_root_node(&mut proof.into_iter().peekable()).map_err(|e| e.into_kind()),
978            Err(SparseStateTrieErrorKind::InvalidRootNode { .. })
979        );
980    }
981
982    #[test]
983    fn reveal_account_empty() {
984        let retainer = ProofRetainer::from_iter([Nibbles::default()]);
985        let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
986        hash_builder.root();
987        let proofs = hash_builder.take_proof_nodes();
988        assert_eq!(proofs.len(), 1);
989
990        let mut sparse = SparseStateTrie::default();
991        assert_eq!(sparse.state, SparseTrie::Blind);
992
993        sparse.reveal_account(Default::default(), proofs.into_inner()).unwrap();
994        assert_eq!(sparse.state, SparseTrie::revealed_empty());
995    }
996
997    #[test]
998    fn reveal_storage_slot_empty() {
999        let retainer = ProofRetainer::from_iter([Nibbles::default()]);
1000        let mut hash_builder = HashBuilder::default().with_proof_retainer(retainer);
1001        hash_builder.root();
1002        let proofs = hash_builder.take_proof_nodes();
1003        assert_eq!(proofs.len(), 1);
1004
1005        let mut sparse = SparseStateTrie::default();
1006        assert!(sparse.storages.is_empty());
1007
1008        sparse
1009            .reveal_storage_slot(Default::default(), Default::default(), proofs.into_inner())
1010            .unwrap();
1011        assert_eq!(
1012            sparse.storages,
1013            HashMap::from_iter([(Default::default(), SparseTrie::revealed_empty())])
1014        );
1015    }
1016
1017    #[test]
1018    fn reveal_account_path_twice() {
1019        let is_private = false; // hardcode to false for legacy test, TODO: make a private equivalent
1020        let mut sparse = SparseStateTrie::default();
1021
1022        let leaf_value = alloy_rlp::encode(TrieAccount::default());
1023        let leaf_1 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
1024            Nibbles::default(),
1025            leaf_value.clone(),
1026            is_private,
1027        )));
1028        let leaf_2 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
1029            Nibbles::default(),
1030            leaf_value.clone(),
1031            is_private,
1032        )));
1033
1034        let multiproof = MultiProof {
1035            account_subtree: ProofNodes::from_iter([
1036                (
1037                    Nibbles::default(),
1038                    alloy_rlp::encode(TrieNode::Branch(BranchNode {
1039                        stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)],
1040                        state_mask: TrieMask::new(0b11),
1041                    }))
1042                    .into(),
1043                ),
1044                (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()),
1045                (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
1046            ]),
1047            ..Default::default()
1048        };
1049
1050        // Reveal multiproof and check that the state trie contains the leaf node and value
1051        sparse.reveal_decoded_multiproof(multiproof.clone().try_into().unwrap()).unwrap();
1052        assert!(sparse
1053            .state_trie_ref()
1054            .unwrap()
1055            .nodes_ref()
1056            .contains_key(&Nibbles::from_nibbles([0x0])),);
1057        assert_eq!(
1058            sparse.state_trie_ref().unwrap().get_leaf_value(&Nibbles::from_nibbles([0x0])),
1059            Some(&leaf_value)
1060        );
1061
1062        // Remove the leaf node and check that the state trie does not contain the leaf node and
1063        // value
1064        sparse.remove_account_leaf(&Nibbles::from_nibbles([0x0])).unwrap();
1065        assert!(!sparse
1066            .state_trie_ref()
1067            .unwrap()
1068            .nodes_ref()
1069            .contains_key(&Nibbles::from_nibbles([0x0])),);
1070        assert!(sparse
1071            .state_trie_ref()
1072            .unwrap()
1073            .get_leaf_value(&Nibbles::from_nibbles([0x0]))
1074            .is_none());
1075
1076        // Reveal multiproof again and check that the state trie still does not contain the leaf
1077        // node and value, because they were already revealed before
1078        sparse.reveal_decoded_multiproof(multiproof.try_into().unwrap()).unwrap();
1079        assert!(!sparse
1080            .state_trie_ref()
1081            .unwrap()
1082            .nodes_ref()
1083            .contains_key(&Nibbles::from_nibbles([0x0])));
1084        assert!(sparse
1085            .state_trie_ref()
1086            .unwrap()
1087            .get_leaf_value(&Nibbles::from_nibbles([0x0]))
1088            .is_none());
1089    }
1090
1091    #[test]
1092    fn reveal_storage_path_twice() {
1093        let is_private = false; // hardcode to false for legacy test, TODO: make a private equivalent
1094        let mut sparse = SparseStateTrie::default();
1095
1096        let leaf_value = alloy_rlp::encode(TrieAccount::default());
1097        let leaf_1 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
1098            Nibbles::default(),
1099            leaf_value.clone(),
1100            is_private,
1101        )));
1102        let leaf_2 = alloy_rlp::encode(TrieNode::Leaf(LeafNode::new(
1103            Nibbles::default(),
1104            leaf_value.clone(),
1105            is_private,
1106        )));
1107
1108        let multiproof = MultiProof {
1109            storages: HashMap::from_iter([(
1110                B256::ZERO,
1111                StorageMultiProof {
1112                    root: B256::ZERO,
1113                    subtree: ProofNodes::from_iter([
1114                        (
1115                            Nibbles::default(),
1116                            alloy_rlp::encode(TrieNode::Branch(BranchNode {
1117                                stack: vec![RlpNode::from_rlp(&leaf_1), RlpNode::from_rlp(&leaf_2)],
1118                                state_mask: TrieMask::new(0b11),
1119                            }))
1120                            .into(),
1121                        ),
1122                        (Nibbles::from_nibbles([0x0]), leaf_1.clone().into()),
1123                        (Nibbles::from_nibbles([0x1]), leaf_1.clone().into()),
1124                    ]),
1125                    branch_node_hash_masks: Default::default(),
1126                    branch_node_tree_masks: Default::default(),
1127                },
1128            )]),
1129            ..Default::default()
1130        };
1131
1132        // Reveal multiproof and check that the storage trie contains the leaf node and value
1133        sparse.reveal_decoded_multiproof(multiproof.clone().try_into().unwrap()).unwrap();
1134        assert!(sparse
1135            .storage_trie_ref(&B256::ZERO)
1136            .unwrap()
1137            .nodes_ref()
1138            .contains_key(&Nibbles::from_nibbles([0x0])),);
1139        assert_eq!(
1140            sparse
1141                .storage_trie_ref(&B256::ZERO)
1142                .unwrap()
1143                .get_leaf_value(&Nibbles::from_nibbles([0x0])),
1144            Some(&leaf_value)
1145        );
1146
1147        // Remove the leaf node and check that the storage trie does not contain the leaf node and
1148        // value
1149        sparse.remove_storage_leaf(B256::ZERO, &Nibbles::from_nibbles([0x0])).unwrap();
1150        assert!(!sparse
1151            .storage_trie_ref(&B256::ZERO)
1152            .unwrap()
1153            .nodes_ref()
1154            .contains_key(&Nibbles::from_nibbles([0x0])),);
1155        assert!(sparse
1156            .storage_trie_ref(&B256::ZERO)
1157            .unwrap()
1158            .get_leaf_value(&Nibbles::from_nibbles([0x0]))
1159            .is_none());
1160
1161        // Reveal multiproof again and check that the storage trie still does not contain the leaf
1162        // node and value, because they were already revealed before
1163        sparse.reveal_decoded_multiproof(multiproof.try_into().unwrap()).unwrap();
1164        assert!(!sparse
1165            .storage_trie_ref(&B256::ZERO)
1166            .unwrap()
1167            .nodes_ref()
1168            .contains_key(&Nibbles::from_nibbles([0x0])));
1169        assert!(sparse
1170            .storage_trie_ref(&B256::ZERO)
1171            .unwrap()
1172            .get_leaf_value(&Nibbles::from_nibbles([0x0]))
1173            .is_none());
1174    }
1175
1176    #[test]
1177    fn take_trie_updates() {
1178        reth_tracing::init_test_tracing();
1179
1180        // let mut rng = generators::rng();
1181        let mut rng = StdRng::seed_from_u64(1);
1182
1183        let mut bytes = [0u8; 1024];
1184        rng.fill(bytes.as_mut_slice());
1185
1186        let slot_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
1187        let slot_path_1 = Nibbles::unpack(slot_1);
1188        let value_1 = U256::from(rng.random::<u64>());
1189        let slot_2 = b256!("0x1100000000000000000000000000000000000000000000000000000000000000");
1190        let slot_path_2 = Nibbles::unpack(slot_2);
1191        let value_2 = U256::from(rng.random::<u64>());
1192        let slot_3 = b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
1193        let slot_path_3 = Nibbles::unpack(slot_3);
1194        let value_3 = U256::from(rng.random::<u64>());
1195
1196        let mut storage_hash_builder =
1197            HashBuilder::default().with_proof_retainer(ProofRetainer::from_iter([
1198                slot_path_1.clone(),
1199                slot_path_2.clone(),
1200            ]));
1201        let is_private = false; // hardcode to false for legacy test, TODO: make a private equivalent
1202        storage_hash_builder.add_leaf(
1203            slot_path_1,
1204            &alloy_rlp::encode_fixed_size(&value_1),
1205            is_private,
1206        );
1207        storage_hash_builder.add_leaf(
1208            slot_path_2,
1209            &alloy_rlp::encode_fixed_size(&value_2),
1210            is_private,
1211        );
1212
1213        let storage_root = storage_hash_builder.root();
1214        let storage_proof_nodes = storage_hash_builder.take_proof_nodes();
1215        let storage_branch_node_hash_masks = HashMap::from_iter([
1216            (Nibbles::default(), TrieMask::new(0b010)),
1217            (Nibbles::from_nibbles([0x1]), TrieMask::new(0b11)),
1218        ]);
1219
1220        let address_1 = b256!("0x1000000000000000000000000000000000000000000000000000000000000000");
1221        let address_path_1 = Nibbles::unpack(address_1);
1222        let account_1 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
1223        let mut trie_account_1 = account_1.into_trie_account(storage_root);
1224        let address_2 = b256!("0x1100000000000000000000000000000000000000000000000000000000000000");
1225        let address_path_2 = Nibbles::unpack(address_2);
1226        let account_2 = Account::arbitrary(&mut arbitrary::Unstructured::new(&bytes)).unwrap();
1227        let mut trie_account_2 = account_2.into_trie_account(EMPTY_ROOT_HASH);
1228
1229        let mut hash_builder =
1230            HashBuilder::default().with_proof_retainer(ProofRetainer::from_iter([
1231                address_path_1.clone(),
1232                address_path_2.clone(),
1233            ]));
1234        let is_private = false; // account leaves are always public. Their storage leaves can be private.
1235        hash_builder.add_leaf(
1236            address_path_1.clone(),
1237            &alloy_rlp::encode(trie_account_1),
1238            is_private,
1239        );
1240        hash_builder.add_leaf(
1241            address_path_2.clone(),
1242            &alloy_rlp::encode(trie_account_2),
1243            is_private,
1244        );
1245
1246        let root = hash_builder.root();
1247        let proof_nodes = hash_builder.take_proof_nodes();
1248
1249        let mut sparse = SparseStateTrie::default().with_updates(true);
1250        sparse
1251            .reveal_decoded_multiproof(
1252                MultiProof {
1253                    account_subtree: proof_nodes,
1254                    branch_node_hash_masks: HashMap::from_iter([(
1255                        Nibbles::from_nibbles([0x1]),
1256                        TrieMask::new(0b00),
1257                    )]),
1258                    branch_node_tree_masks: HashMap::default(),
1259                    storages: HashMap::from_iter([
1260                        (
1261                            address_1,
1262                            StorageMultiProof {
1263                                root,
1264                                subtree: storage_proof_nodes.clone(),
1265                                branch_node_hash_masks: storage_branch_node_hash_masks.clone(),
1266                                branch_node_tree_masks: HashMap::default(),
1267                            },
1268                        ),
1269                        (
1270                            address_2,
1271                            StorageMultiProof {
1272                                root,
1273                                subtree: storage_proof_nodes,
1274                                branch_node_hash_masks: storage_branch_node_hash_masks,
1275                                branch_node_tree_masks: HashMap::default(),
1276                            },
1277                        ),
1278                    ]),
1279                }
1280                .try_into()
1281                .unwrap(),
1282            )
1283            .unwrap();
1284
1285        assert_eq!(sparse.root().unwrap(), root);
1286
1287        let address_3 = b256!("0x2000000000000000000000000000000000000000000000000000000000000000");
1288        let address_path_3 = Nibbles::unpack(address_3);
1289        let account_3 = Account { nonce: account_1.nonce + 1, ..account_1 };
1290        let trie_account_3 = account_3.into_trie_account(EMPTY_ROOT_HASH);
1291
1292        sparse.update_account_leaf(address_path_3, alloy_rlp::encode(trie_account_3)).unwrap();
1293
1294        let is_private = false; // legacy test does not use private storage
1295        sparse
1296            .update_storage_leaf(address_1, slot_path_3, alloy_rlp::encode(value_3), is_private)
1297            .unwrap();
1298        trie_account_1.storage_root = sparse.storage_root(address_1).unwrap();
1299        sparse.update_account_leaf(address_path_1, alloy_rlp::encode(trie_account_1)).unwrap();
1300
1301        sparse.wipe_storage(address_2).unwrap();
1302        trie_account_2.storage_root = sparse.storage_root(address_2).unwrap();
1303        sparse.update_account_leaf(address_path_2, alloy_rlp::encode(trie_account_2)).unwrap();
1304
1305        sparse.root().unwrap();
1306
1307        let sparse_updates = sparse.take_trie_updates().unwrap();
1308        // TODO(alexey): assert against real state root calculation updates
1309        pretty_assertions::assert_eq!(
1310            sparse_updates,
1311            TrieUpdates {
1312                account_nodes: HashMap::default(),
1313                storage_tries: HashMap::from_iter([(
1314                    b256!("0x1100000000000000000000000000000000000000000000000000000000000000"),
1315                    StorageTrieUpdates {
1316                        is_deleted: true,
1317                        storage_nodes: HashMap::default(),
1318                        removed_nodes: HashSet::default()
1319                    }
1320                )]),
1321                removed_nodes: HashSet::default()
1322            }
1323        );
1324    }
1325
1326    #[test]
1327    fn test_filter_revealed_nodes() {
1328        let is_private = false; // hardcode to false for legacy test
1329        let revealed_nodes = HashSet::from_iter([Nibbles::from_nibbles([0x0])]);
1330        let leaf =
1331            TrieNode::Leaf(LeafNode::new(Nibbles::default(), alloy_rlp::encode([]), is_private));
1332        let leaf_encoded = alloy_rlp::encode(&leaf);
1333        let branch = TrieNode::Branch(BranchNode::new(
1334            vec![RlpNode::from_rlp(&leaf_encoded), RlpNode::from_rlp(&leaf_encoded)],
1335            TrieMask::new(0b11),
1336        ));
1337        let proof_nodes = alloy_trie::proof::DecodedProofNodes::from_iter([
1338            (Nibbles::default(), branch.clone()),
1339            (Nibbles::from_nibbles([0x0]), leaf.clone()),
1340            (Nibbles::from_nibbles([0x1]), leaf.clone()),
1341        ]);
1342
1343        let decoded = filter_revealed_nodes(proof_nodes, &revealed_nodes).unwrap();
1344
1345        assert_eq!(
1346            decoded,
1347            FilteredProofNodes {
1348                nodes: vec![(Nibbles::default(), branch), (Nibbles::from_nibbles([0x1]), leaf)],
1349                // Branch, leaf, leaf
1350                total_nodes: 3,
1351                // Revealed leaf node with path 0x1
1352                skipped_nodes: 1,
1353                // Branch, two of its children, one leaf
1354                new_nodes: 4
1355            }
1356        );
1357    }
1358}