reth_trie_sparse/
trie.rs

1use crate::blinded::{BlindedProvider, DefaultBlindedProvider, RevealedNode};
2use alloc::{
3    borrow::Cow,
4    boxed::Box,
5    fmt,
6    string::{String, ToString},
7    vec,
8    vec::Vec,
9};
10use alloy_primitives::{
11    hex, keccak256,
12    map::{Entry, HashMap, HashSet},
13    B256,
14};
15use alloy_rlp::Decodable;
16use reth_execution_errors::{SparseTrieErrorKind, SparseTrieResult};
17use reth_trie_common::{
18    prefix_set::{PrefixSet, PrefixSetMut},
19    BranchNodeCompact, BranchNodeRef, ExtensionNodeRef, LeafNodeRef, Nibbles, RlpNode, TrieMask,
20    TrieNode, CHILD_INDEX_RANGE, EMPTY_ROOT_HASH,
21};
22use smallvec::SmallVec;
23use tracing::trace;
24
25/// Struct for passing around branch node mask information.
26///
27/// Branch nodes can have up to 16 children (one for each nibble).
28/// The masks represent which children are stored in different ways:
29/// - `hash_mask`: Indicates which children are stored as hashes in the database
30/// - `tree_mask`: Indicates which children are complete subtrees stored in the database
31///
32/// These masks are essential for efficient trie traversal and serialization, as they
33/// determine how nodes should be encoded and stored on disk.
34#[derive(Debug)]
35pub struct TrieMasks {
36    /// Branch node hash mask, if any.
37    ///
38    /// When a bit is set, the corresponding child node's hash is stored in the trie.
39    ///
40    /// This mask enables selective hashing of child nodes.
41    pub hash_mask: Option<TrieMask>,
42    /// Branch node tree mask, if any.
43    ///
44    /// When a bit is set, the corresponding child subtree is stored in the database.
45    pub tree_mask: Option<TrieMask>,
46}
47
48impl TrieMasks {
49    /// Helper function, returns both fields `hash_mask` and `tree_mask` as [`None`]
50    pub const fn none() -> Self {
51        Self { hash_mask: None, tree_mask: None }
52    }
53}
54
55/// A sparse trie that is either in a "blind" state (no nodes are revealed, root node hash is
56/// unknown) or in a "revealed" state (root node has been revealed and the trie can be updated).
57///
58/// In blind mode the trie does not contain any decoded node data, which saves memory but
59/// prevents direct access to node contents. The revealed mode stores decoded nodes along
60/// with additional information such as values, allowing direct manipulation.
61///
62/// The sparse trie design is optimised for:
63/// 1. Memory efficiency - only revealed nodes are loaded into memory
64/// 2. Update tracking - changes to the trie structure can be tracked and selectively persisted
65/// 3. Incremental operations - nodes can be revealed as needed without loading the entire trie.
66///    This is what gives rise to the notion of a "sparse" trie.
67#[derive(PartialEq, Eq, Default)]
68pub enum SparseTrie<P = DefaultBlindedProvider> {
69    /// The trie is blind -- no nodes have been revealed
70    ///
71    /// This is the default state. In this state,
72    /// the trie cannot be directly queried or modified until nodes are revealed.
73    #[default]
74    Blind,
75    /// Some nodes in the Trie have been revealed.
76    ///
77    /// In this state, the trie can be queried and modified for the parts
78    /// that have been revealed. Other parts remain blind and require revealing
79    /// before they can be accessed.
80    Revealed(Box<RevealedSparseTrie<P>>),
81}
82
83impl<P> fmt::Debug for SparseTrie<P> {
84    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
85        match self {
86            Self::Blind => write!(f, "Blind"),
87            Self::Revealed(revealed) => write!(f, "Revealed({revealed:?})"),
88        }
89    }
90}
91
92impl SparseTrie {
93    /// Creates a new blind sparse trie.
94    ///
95    /// # Examples
96    ///
97    /// ```
98    /// use reth_trie_sparse::{blinded::DefaultBlindedProvider, SparseTrie};
99    ///
100    /// let trie: SparseTrie<DefaultBlindedProvider> = SparseTrie::blind();
101    /// assert!(trie.is_blind());
102    /// let trie: SparseTrie<DefaultBlindedProvider> = SparseTrie::default();
103    /// assert!(trie.is_blind());
104    /// ```
105    pub const fn blind() -> Self {
106        Self::Blind
107    }
108
109    /// Creates a new revealed but empty sparse trie with `SparseNode::Empty` as root node.
110    ///
111    /// # Examples
112    ///
113    /// ```
114    /// use reth_trie_sparse::{blinded::DefaultBlindedProvider, SparseTrie};
115    ///
116    /// let trie: SparseTrie<DefaultBlindedProvider> = SparseTrie::revealed_empty();
117    /// assert!(!trie.is_blind());
118    /// ```
119    pub fn revealed_empty() -> Self {
120        Self::Revealed(Box::default())
121    }
122
123    /// Reveals the root node, converting a blind trie into a revealed one.
124    ///
125    /// If the trie is blinded, its root node is replaced with `root`.
126    ///
127    /// The `masks` are used to determine how the node's children are stored.
128    /// The `retain_updates` flag controls whether changes to the trie structure
129    /// should be tracked.
130    ///
131    /// # Returns
132    ///
133    /// A mutable reference to the underlying [`RevealedSparseTrie`].
134    pub fn reveal_root(
135        &mut self,
136        root: TrieNode,
137        masks: TrieMasks,
138        retain_updates: bool,
139    ) -> SparseTrieResult<&mut RevealedSparseTrie> {
140        self.reveal_root_with_provider(Default::default(), root, masks, retain_updates)
141    }
142}
143
144impl<P> SparseTrie<P> {
145    /// Returns `true` if the sparse trie has no revealed nodes.
146    pub const fn is_blind(&self) -> bool {
147        matches!(self, Self::Blind)
148    }
149
150    /// Returns an immutable reference to the underlying revealed sparse trie.
151    ///
152    /// Returns `None` if the trie is blinded.
153    pub const fn as_revealed_ref(&self) -> Option<&RevealedSparseTrie<P>> {
154        if let Self::Revealed(revealed) = self {
155            Some(revealed)
156        } else {
157            None
158        }
159    }
160
161    /// Returns a mutable reference to the underlying revealed sparse trie.
162    ///
163    /// Returns `None` if the trie is blinded.
164    pub fn as_revealed_mut(&mut self) -> Option<&mut RevealedSparseTrie<P>> {
165        if let Self::Revealed(revealed) = self {
166            Some(revealed)
167        } else {
168            None
169        }
170    }
171
172    /// Reveals the root node using a specified provider.
173    ///
174    /// This function is similar to [`Self::reveal_root`] but allows the caller to provide
175    /// a custom provider for fetching blinded nodes.
176    ///
177    /// # Returns
178    ///
179    /// Mutable reference to [`RevealedSparseTrie`].
180    pub fn reveal_root_with_provider(
181        &mut self,
182        provider: P,
183        root: TrieNode,
184        masks: TrieMasks,
185        retain_updates: bool,
186    ) -> SparseTrieResult<&mut RevealedSparseTrie<P>> {
187        if self.is_blind() {
188            *self = Self::Revealed(Box::new(RevealedSparseTrie::from_provider_and_root(
189                provider,
190                root,
191                masks,
192                retain_updates,
193            )?))
194        }
195        Ok(self.as_revealed_mut().unwrap())
196    }
197
198    /// Wipes the trie by removing all nodes and values,
199    /// and resetting the trie to only contain an empty root node.
200    ///
201    /// Note: This method will error if the trie is blinded.
202    pub fn wipe(&mut self) -> SparseTrieResult<()> {
203        let revealed = self.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?;
204        revealed.wipe();
205        Ok(())
206    }
207
208    /// Calculates the root hash of the trie.
209    ///
210    /// This will update any remaining dirty nodes before computing the root hash.
211    /// "dirty" nodes are nodes that need their hashes to be recomputed because one or more of their
212    /// children's hashes have changed.
213    ///
214    /// # Returns
215    ///
216    /// - `Some(B256)` with the calculated root hash if the trie is revealed.
217    /// - `None` if the trie is still blind.
218    pub fn root(&mut self) -> Option<B256> {
219        Some(self.as_revealed_mut()?.root())
220    }
221
222    /// Returns the root hash along with any accumulated update information.
223    ///
224    /// This is useful for when you need both the root hash and information about
225    /// what nodes were modified, which can be used to efficiently update
226    /// an external database.
227    ///
228    /// # Returns
229    ///
230    /// An `Option` tuple consisting of:
231    ///  - The trie root hash (`B256`).
232    ///  - A [`SparseTrieUpdates`] structure containing information about updated nodes.
233    ///  - `None` if the trie is still blind.
234    pub fn root_with_updates(&mut self) -> Option<(B256, SparseTrieUpdates)> {
235        let revealed = self.as_revealed_mut()?;
236        Some((revealed.root(), revealed.take_updates()))
237    }
238}
239
240impl<P: BlindedProvider> SparseTrie<P> {
241    /// Update the leaf node.
242    pub fn update_leaf(
243        &mut self,
244        path: Nibbles,
245        value: Vec<u8>,
246        is_private: bool,
247    ) -> SparseTrieResult<()> {
248        let revealed = self.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?;
249        revealed.update_leaf(path, value, is_private)?;
250        Ok(())
251    }
252
253    /// Removes a leaf node at the specified key path.
254    ///
255    /// # Errors
256    ///
257    /// Returns an error if the trie is still blind, or if the leaf cannot be removed
258    pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> {
259        let revealed = self.as_revealed_mut().ok_or(SparseTrieErrorKind::Blind)?;
260        revealed.remove_leaf(path)?;
261        Ok(())
262    }
263}
264
265/// The representation of revealed sparse trie.
266///
267/// The revealed sparse trie contains the actual trie structure with nodes, values, and
268/// tracking for changes. It supports operations like inserting, updating, and removing
269/// nodes.
270///
271///
272/// ## Invariants
273///
274/// - The root node is always present in `nodes` collection.
275/// - Each leaf entry in `nodes` collection must have a corresponding entry in `values` collection.
276///   The opposite is also true.
277/// - All keys in `values` collection are full leaf paths.
278#[derive(Clone, PartialEq, Eq)]
279pub struct RevealedSparseTrie<P = DefaultBlindedProvider> {
280    /// Provider used for retrieving blinded nodes.
281    /// This allows lazily loading parts of the trie from an external source.
282    provider: P,
283    /// Map from a path (nibbles) to its corresponding sparse trie node.
284    /// This contains all of the revealed nodes in trie.
285    nodes: HashMap<Nibbles, SparseNode>,
286    /// When a branch is set, the corresponding child subtree is stored in the database.
287    branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
288    /// When a bit is set, the corresponding child is stored as a hash in the database.
289    branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
290    /// Map from leaf key paths to their values.
291    /// All values are stored here instead of directly in leaf nodes.
292    values: HashMap<Nibbles, Vec<u8>>,
293    /// Set of prefixes (key paths) that have been marked as updated.
294    /// This is used to track which parts of the trie need to be recalculated.
295    prefix_set: PrefixSetMut,
296    /// Optional tracking of trie updates for later use.
297    updates: Option<SparseTrieUpdates>,
298    /// Reusable buffer for RLP encoding of nodes.
299    rlp_buf: Vec<u8>,
300}
301
302impl<P> fmt::Debug for RevealedSparseTrie<P> {
303    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
304        f.debug_struct("RevealedSparseTrie")
305            .field("nodes", &self.nodes)
306            .field("branch_tree_masks", &self.branch_node_tree_masks)
307            .field("branch_hash_masks", &self.branch_node_hash_masks)
308            .field("values", &self.values)
309            .field("prefix_set", &self.prefix_set)
310            .field("updates", &self.updates)
311            .field("rlp_buf", &hex::encode(&self.rlp_buf))
312            .finish_non_exhaustive()
313    }
314}
315
316/// Turns a [`Nibbles`] into a [`String`] by concatenating each nibbles' hex character.
317fn encode_nibbles(nibbles: &Nibbles) -> String {
318    let encoded = hex::encode(nibbles.pack());
319    encoded[..nibbles.len()].to_string()
320}
321
322impl<P: BlindedProvider> fmt::Display for RevealedSparseTrie<P> {
323    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
324        // This prints the trie in preorder traversal, using a stack
325        let mut stack = Vec::new();
326        let mut visited = HashSet::new();
327
328        // 4 spaces as indent per level
329        const INDENT: &str = "    ";
330
331        // Track both path and depth
332        stack.push((Nibbles::default(), self.nodes_ref().get(&Nibbles::default()).unwrap(), 0));
333
334        while let Some((path, node, depth)) = stack.pop() {
335            if !visited.insert(path.clone()) {
336                continue;
337            }
338
339            // Add indentation if alternate flag (#) is set
340            if f.alternate() {
341                write!(f, "{}", INDENT.repeat(depth))?;
342            }
343
344            let packed_path = if depth == 0 { String::from("Root") } else { encode_nibbles(&path) };
345
346            match node {
347                SparseNode::Empty | SparseNode::Hash(_) => {
348                    writeln!(f, "{packed_path} -> {node:?}")?;
349                }
350                SparseNode::Leaf { key, .. } => {
351                    // we want to append the key to the path
352                    let mut full_path = path.clone();
353                    full_path.extend_from_slice_unchecked(key);
354                    let packed_path = encode_nibbles(&full_path);
355
356                    writeln!(f, "{packed_path} -> {node:?}")?;
357                }
358                SparseNode::Extension { key, .. } => {
359                    writeln!(f, "{packed_path} -> {node:?}")?;
360
361                    // push the child node onto the stack with increased depth
362                    let mut child_path = path.clone();
363                    child_path.extend_from_slice_unchecked(key);
364                    if let Some(child_node) = self.nodes_ref().get(&child_path) {
365                        stack.push((child_path, child_node, depth + 1));
366                    }
367                }
368                SparseNode::Branch { state_mask, .. } => {
369                    writeln!(f, "{packed_path} -> {node:?}")?;
370
371                    for i in CHILD_INDEX_RANGE.rev() {
372                        if state_mask.is_bit_set(i) {
373                            let mut child_path = path.clone();
374                            child_path.push_unchecked(i);
375                            if let Some(child_node) = self.nodes_ref().get(&child_path) {
376                                stack.push((child_path, child_node, depth + 1));
377                            }
378                        }
379                    }
380                }
381            }
382        }
383
384        Ok(())
385    }
386}
387
388impl Default for RevealedSparseTrie {
389    fn default() -> Self {
390        Self {
391            provider: Default::default(),
392            nodes: HashMap::from_iter([(Nibbles::default(), SparseNode::Empty)]),
393            branch_node_tree_masks: HashMap::default(),
394            branch_node_hash_masks: HashMap::default(),
395            values: HashMap::default(),
396            prefix_set: PrefixSetMut::default(),
397            updates: None,
398            rlp_buf: Vec::new(),
399        }
400    }
401}
402
403impl RevealedSparseTrie {
404    /// Creates a new revealed sparse trie from the given root node.
405    ///
406    /// This function initializes the internal structures and then reveals the root.
407    /// It is a convenient method to create a [`RevealedSparseTrie`] when you already have
408    /// the root node available.
409    ///
410    /// # Returns
411    ///
412    /// A [`RevealedSparseTrie`] if successful, or an error if revealing fails.
413    pub fn from_root(
414        root: TrieNode,
415        masks: TrieMasks,
416        retain_updates: bool,
417    ) -> SparseTrieResult<Self> {
418        let mut this = Self {
419            provider: Default::default(),
420            nodes: HashMap::default(),
421            branch_node_tree_masks: HashMap::default(),
422            branch_node_hash_masks: HashMap::default(),
423            values: HashMap::default(),
424            prefix_set: PrefixSetMut::default(),
425            rlp_buf: Vec::new(),
426            updates: None,
427        }
428        .with_updates(retain_updates);
429        this.reveal_node(Nibbles::default(), root, masks)?;
430        Ok(this)
431    }
432}
433
434impl<P> RevealedSparseTrie<P> {
435    /// Creates a new revealed sparse trie from the given provider and root node.
436    ///
437    /// Similar to `from_root`, but allows specifying a custom provider for
438    /// retrieving blinded nodes.
439    ///
440    /// # Returns
441    ///
442    /// A [`RevealedSparseTrie`] if successful, or an error if revealing fails.
443    pub fn from_provider_and_root(
444        provider: P,
445        node: TrieNode,
446        masks: TrieMasks,
447        retain_updates: bool,
448    ) -> SparseTrieResult<Self> {
449        let mut this = Self {
450            provider,
451            nodes: HashMap::default(),
452            branch_node_tree_masks: HashMap::default(),
453            branch_node_hash_masks: HashMap::default(),
454            values: HashMap::default(),
455            prefix_set: PrefixSetMut::default(),
456            updates: None,
457            rlp_buf: Vec::new(),
458        }
459        .with_updates(retain_updates);
460        this.reveal_node(Nibbles::default(), node, masks)?;
461        Ok(this)
462    }
463
464    /// Replaces the current provider with a new provider.
465    ///
466    /// This allows changing how blinded nodes are retrieved without
467    /// rebuilding the entire trie structure.
468    ///
469    /// # Returns
470    ///
471    /// A new [`RevealedSparseTrie`] with the updated provider.
472    pub fn with_provider<BP>(self, provider: BP) -> RevealedSparseTrie<BP> {
473        RevealedSparseTrie {
474            provider,
475            nodes: self.nodes,
476            branch_node_tree_masks: self.branch_node_tree_masks,
477            branch_node_hash_masks: self.branch_node_hash_masks,
478            values: self.values,
479            prefix_set: self.prefix_set,
480            updates: self.updates,
481            rlp_buf: self.rlp_buf,
482        }
483    }
484
485    /// Configures the trie to retain information about updates.
486    ///
487    /// If `retain_updates` is true, the trie will record branch node updates and deletions.
488    /// This information can then be used to efficiently update an external database.
489    pub fn with_updates(mut self, retain_updates: bool) -> Self {
490        if retain_updates {
491            self.updates = Some(SparseTrieUpdates::default());
492        }
493        self
494    }
495
496    /// Returns a reference to the current sparse trie updates.
497    ///
498    /// If no updates have been made/recorded, returns an empty update set.
499    pub fn updates_ref(&self) -> Cow<'_, SparseTrieUpdates> {
500        self.updates.as_ref().map_or(Cow::Owned(SparseTrieUpdates::default()), Cow::Borrowed)
501    }
502
503    /// Returns an immutable reference to all nodes in the sparse trie.
504    pub const fn nodes_ref(&self) -> &HashMap<Nibbles, SparseNode> {
505        &self.nodes
506    }
507
508    /// Retrieves a reference to the leaf value stored at the given key path, if it is revealed.
509    ///
510    /// This method efficiently retrieves values from the trie without traversing
511    /// the entire node structure, as values are stored in a separate map.
512    ///
513    /// Note: a value can exist in the full trie and this function still returns `None`
514    /// because the value has not been revealed.
515    /// Hence a `None` indicates two possibilities:
516    /// - The value does not exists in the trie, so it cannot be revealed
517    /// - The value has not yet been revealed. In order to determine which is true, one would need
518    ///   an exclusion proof.
519    pub fn get_leaf_value(&self, path: &Nibbles) -> Option<&Vec<u8>> {
520        self.values.get(path)
521    }
522
523    /// Consumes and returns the currently accumulated trie updates.
524    ///
525    /// This is useful when you want to apply the updates to an external database,
526    /// and then start tracking a new set of updates.
527    pub fn take_updates(&mut self) -> SparseTrieUpdates {
528        self.updates.take().unwrap_or_default()
529    }
530
531    /// Reserves capacity in the nodes map for at least `additional` more nodes.
532    pub fn reserve_nodes(&mut self, additional: usize) {
533        self.nodes.reserve(additional);
534    }
535
536    /// Reveals a trie node if it has not been revealed before.
537    ///
538    /// This internal function decodes a trie node and inserts it into the nodes map.
539    /// It handles different node types (leaf, extension, branch) by appropriately
540    /// adding them to the trie structure and recursively revealing their children.
541    ///
542    ///
543    /// # Returns
544    ///
545    /// `Ok(())` if successful, or an error if node was not revealed.
546    pub fn reveal_node(
547        &mut self,
548        path: Nibbles,
549        node: TrieNode,
550        masks: TrieMasks,
551    ) -> SparseTrieResult<()> {
552        // If the node is already revealed and it's not a hash node, do nothing.
553        if self.nodes.get(&path).is_some_and(|node| !node.is_hash()) {
554            return Ok(());
555        }
556
557        if let Some(tree_mask) = masks.tree_mask {
558            self.branch_node_tree_masks.insert(path.clone(), tree_mask);
559        }
560        if let Some(hash_mask) = masks.hash_mask {
561            self.branch_node_hash_masks.insert(path.clone(), hash_mask);
562        }
563
564        match node {
565            TrieNode::EmptyRoot => {
566                // For an empty root, ensure that we are at the root path.
567                debug_assert!(path.is_empty());
568                self.nodes.insert(path, SparseNode::Empty);
569            }
570            TrieNode::Branch(branch) => {
571                // For a branch node, iterate over all potential children
572                let mut stack_ptr = branch.as_ref().first_child_index();
573                for idx in CHILD_INDEX_RANGE {
574                    if branch.state_mask.is_bit_set(idx) {
575                        let mut child_path = path.clone();
576                        child_path.push_unchecked(idx);
577                        // Reveal each child node or hash it has
578                        self.reveal_node_or_hash(child_path, &branch.stack[stack_ptr])?;
579                        stack_ptr += 1;
580                    }
581                }
582                // Update the branch node entry in the nodes map, handling cases where a blinded
583                // node is now replaced with a revealed node.
584                match self.nodes.entry(path) {
585                    Entry::Occupied(mut entry) => match entry.get() {
586                        // Replace a hash node with a fully revealed branch node.
587                        SparseNode::Hash(hash) => {
588                            entry.insert(SparseNode::Branch {
589                                state_mask: branch.state_mask,
590                                // Memoize the hash of a previously blinded node in a new branch
591                                // node.
592                                hash: Some(*hash),
593                                store_in_db_trie: Some(
594                                    masks.hash_mask.is_some_and(|mask| !mask.is_empty()) ||
595                                        masks.tree_mask.is_some_and(|mask| !mask.is_empty()),
596                                ),
597                            });
598                        }
599                        // Branch node already exists, or an extension node was placed where a
600                        // branch node was before.
601                        SparseNode::Branch { .. } | SparseNode::Extension { .. } => {}
602                        // All other node types can't be handled.
603                        node @ (SparseNode::Empty | SparseNode::Leaf { .. }) => {
604                            return Err(SparseTrieErrorKind::Reveal {
605                                path: entry.key().clone(),
606                                node: Box::new(node.clone()),
607                            }
608                            .into())
609                        }
610                    },
611                    Entry::Vacant(entry) => {
612                        entry.insert(SparseNode::new_branch(branch.state_mask));
613                    }
614                }
615            }
616            TrieNode::Extension(ext) => match self.nodes.entry(path) {
617                Entry::Occupied(mut entry) => match entry.get() {
618                    // Replace a hash node with a revealed extension node.
619                    SparseNode::Hash(hash) => {
620                        let mut child_path = entry.key().clone();
621                        child_path.extend_from_slice_unchecked(&ext.key);
622                        entry.insert(SparseNode::Extension {
623                            key: ext.key,
624                            // Memoize the hash of a previously blinded node in a new extension
625                            // node.
626                            hash: Some(*hash),
627                            store_in_db_trie: None,
628                        });
629                        self.reveal_node_or_hash(child_path, &ext.child)?;
630                    }
631                    // Extension node already exists, or an extension node was placed where a branch
632                    // node was before.
633                    SparseNode::Extension { .. } | SparseNode::Branch { .. } => {}
634                    // All other node types can't be handled.
635                    node @ (SparseNode::Empty | SparseNode::Leaf { .. }) => {
636                        return Err(SparseTrieErrorKind::Reveal {
637                            path: entry.key().clone(),
638                            node: Box::new(node.clone()),
639                        }
640                        .into())
641                    }
642                },
643                Entry::Vacant(entry) => {
644                    let mut child_path = entry.key().clone();
645                    child_path.extend_from_slice_unchecked(&ext.key);
646                    entry.insert(SparseNode::new_ext(ext.key));
647                    self.reveal_node_or_hash(child_path, &ext.child)?;
648                }
649            },
650            TrieNode::Leaf(leaf) => match self.nodes.entry(path) {
651                Entry::Occupied(mut entry) => match entry.get() {
652                    // Replace a hash node with a revealed leaf node and store leaf node value.
653                    SparseNode::Hash(hash) => {
654                        let mut full = entry.key().clone();
655                        full.extend_from_slice_unchecked(&leaf.key);
656                        self.values.insert(full, leaf.value);
657                        entry.insert(SparseNode::Leaf {
658                            key: leaf.key,
659                            // Memoize the hash of a previously blinded node in a new leaf
660                            // node.
661                            hash: Some(*hash),
662                            is_private: leaf.is_private,
663                        });
664                    }
665                    // Left node already exists.
666                    SparseNode::Leaf { .. } => {}
667                    // All other node types can't be handled.
668                    node @ (SparseNode::Empty |
669                    SparseNode::Extension { .. } |
670                    SparseNode::Branch { .. }) => {
671                        return Err(SparseTrieErrorKind::Reveal {
672                            path: entry.key().clone(),
673                            node: Box::new(node.clone()),
674                        }
675                        .into())
676                    }
677                },
678                Entry::Vacant(entry) => {
679                    let mut full = entry.key().clone();
680                    full.extend_from_slice_unchecked(&leaf.key);
681                    entry.insert(SparseNode::new_leaf(leaf.key, leaf.is_private));
682                    self.values.insert(full, leaf.value);
683                }
684            },
685        }
686
687        Ok(())
688    }
689
690    /// Reveals either a node or its hash placeholder based on the provided child data.
691    ///
692    /// When traversing the trie, we often encounter references to child nodes that
693    /// are either directly embedded or represented by their hash. This method
694    /// handles both cases:
695    ///
696    /// 1. If the child data represents a hash (32+1=33 bytes), store it as a hash node
697    /// 2. Otherwise, decode the data as a [`TrieNode`] and recursively reveal it using
698    ///    `reveal_node`
699    ///
700    /// # Returns
701    ///
702    /// Returns `Ok(())` if successful, or an error if the node cannot be revealed.
703    ///
704    /// # Error Handling
705    ///
706    /// Will error if there's a conflict between a new hash node and an existing one
707    /// at the same path
708    fn reveal_node_or_hash(&mut self, path: Nibbles, child: &[u8]) -> SparseTrieResult<()> {
709        if child.len() == B256::len_bytes() + 1 {
710            let hash = B256::from_slice(&child[1..]);
711            match self.nodes.entry(path) {
712                Entry::Occupied(entry) => match entry.get() {
713                    // Hash node with a different hash can't be handled.
714                    SparseNode::Hash(previous_hash) if previous_hash != &hash => {
715                        return Err(SparseTrieErrorKind::Reveal {
716                            path: entry.key().clone(),
717                            node: Box::new(SparseNode::Hash(hash)),
718                        }
719                        .into())
720                    }
721                    _ => {}
722                },
723                Entry::Vacant(entry) => {
724                    entry.insert(SparseNode::Hash(hash));
725                }
726            }
727            return Ok(());
728        }
729
730        self.reveal_node(path, TrieNode::decode(&mut &child[..])?, TrieMasks::none())
731    }
732
733    /// Traverse the trie from the root down to the leaf at the given path,
734    /// removing and collecting all nodes along that path.
735    ///
736    /// This helper function is used during leaf removal to extract the nodes of the trie
737    /// that will be affected by the deletion. These nodes are then re-inserted and modified
738    /// as needed (collapsing extension nodes etc) given that the leaf has now been removed.
739    ///
740    /// # Returns
741    ///
742    /// Returns a vector of [`RemovedSparseNode`] representing the nodes removed during the
743    /// traversal.
744    ///
745    /// # Errors
746    ///
747    /// Returns an error if a blinded node or an empty node is encountered unexpectedly,
748    /// as these prevent proper removal of the leaf.
749    fn take_nodes_for_path(&mut self, path: &Nibbles) -> SparseTrieResult<Vec<RemovedSparseNode>> {
750        let mut current = Nibbles::default(); // Start traversal from the root
751        let mut nodes = Vec::new(); // Collect traversed nodes
752
753        while let Some(node) = self.nodes.remove(&current) {
754            match &node {
755                SparseNode::Empty => return Err(SparseTrieErrorKind::Blind.into()),
756                &SparseNode::Hash(hash) => {
757                    return Err(SparseTrieErrorKind::BlindedNode { path: current, hash }.into())
758                }
759                SparseNode::Leaf { key: _key, .. } => {
760                    // Leaf node is always the one that we're deleting, and no other leaf nodes can
761                    // be found during traversal.
762
763                    #[cfg(debug_assertions)]
764                    {
765                        let mut current = current.clone();
766                        current.extend_from_slice_unchecked(_key);
767                        assert_eq!(&current, path);
768                    }
769
770                    nodes.push(RemovedSparseNode {
771                        path: current.clone(),
772                        node,
773                        unset_branch_nibble: None,
774                    });
775                    break;
776                }
777                SparseNode::Extension { key, .. } => {
778                    #[cfg(debug_assertions)]
779                    {
780                        let mut current = current.clone();
781                        current.extend_from_slice_unchecked(key);
782                        assert!(
783                            path.starts_with(&current),
784                            "path: {path:?}, current: {current:?}, key: {key:?}",
785                        );
786                    }
787
788                    let path = current.clone();
789                    current.extend_from_slice_unchecked(key);
790                    nodes.push(RemovedSparseNode { path, node, unset_branch_nibble: None });
791                }
792                SparseNode::Branch { state_mask, .. } => {
793                    let nibble = path[current.len()];
794                    debug_assert!(
795                        state_mask.is_bit_set(nibble),
796                        "current: {current:?}, path: {path:?}, nibble: {nibble:?}, state_mask: {state_mask:?}",
797                    );
798
799                    // If the branch node has a child that is a leaf node that we're removing,
800                    // we need to unset this nibble.
801                    // Any other branch nodes will not require unsetting the nibble, because
802                    // deleting one leaf node can not remove the whole path
803                    // where the branch node is located.
804                    let mut child_path =
805                        Nibbles::from_nibbles([current.as_slice(), &[nibble]].concat());
806                    let unset_branch_nibble = self
807                        .nodes
808                        .get(&child_path)
809                        .is_some_and(move |node| match node {
810                            SparseNode::Leaf { key, .. } => {
811                                // Get full path of the leaf node
812                                child_path.extend_from_slice_unchecked(key);
813                                &child_path == path
814                            }
815                            _ => false,
816                        })
817                        .then_some(nibble);
818
819                    nodes.push(RemovedSparseNode {
820                        path: current.clone(),
821                        node,
822                        unset_branch_nibble,
823                    });
824
825                    current.push_unchecked(nibble);
826                }
827            }
828        }
829
830        Ok(nodes)
831    }
832
833    /// Removes all nodes and values from the trie, resetting it to a blank state
834    /// with only an empty root node.
835    ///
836    /// Note: All previously tracked changes to the trie are also removed.
837    pub fn wipe(&mut self) {
838        self.nodes = HashMap::from_iter([(Nibbles::default(), SparseNode::Empty)]);
839        self.values = HashMap::default();
840        self.prefix_set = PrefixSetMut::all();
841        self.updates = self.updates.is_some().then(SparseTrieUpdates::wiped);
842    }
843
844    /// Calculates and returns the root hash of the trie.
845    ///
846    /// Before computing the hash, this function processes any remaining (dirty) nodes by
847    /// updating their RLP encodings. The root hash is either:
848    /// 1. The cached hash (if no dirty nodes were found)
849    /// 2. The keccak256 hash of the root node's RLP representation
850    pub fn root(&mut self) -> B256 {
851        // Take the current prefix set
852        let mut prefix_set = core::mem::take(&mut self.prefix_set).freeze();
853        let rlp_node = self.rlp_node_allocate(&mut prefix_set);
854        if let Some(root_hash) = rlp_node.as_hash() {
855            root_hash
856        } else {
857            keccak256(rlp_node)
858        }
859    }
860
861    /// Recalculates and updates the RLP hashes of nodes deeper than or equal to the specified
862    /// `depth`.
863    ///
864    /// The root node is considered to be at level 0. This method is useful for optimizing
865    /// hash recalculations after localized changes to the trie structure:
866    ///
867    /// This function identifies all nodes that have changed (based on the prefix set) at the given
868    /// depth and recalculates their RLP representation.
869    pub fn update_rlp_node_level(&mut self, depth: usize) {
870        // Take the current prefix set
871        let mut prefix_set = core::mem::take(&mut self.prefix_set).freeze();
872        let mut buffers = RlpNodeBuffers::default();
873
874        // Get the nodes that have changed at the given depth.
875        let (targets, new_prefix_set) = self.get_changed_nodes_at_depth(&mut prefix_set, depth);
876        // Update the prefix set to the prefix set of the nodes that still need to be updated.
877        self.prefix_set = new_prefix_set;
878
879        trace!(target: "trie::sparse", ?depth, ?targets, "Updating nodes at depth");
880
881        let mut temp_rlp_buf = core::mem::take(&mut self.rlp_buf);
882        for (level, path) in targets {
883            buffers.path_stack.push(RlpNodePathStackItem {
884                level,
885                path,
886                is_in_prefix_set: Some(true),
887            });
888            self.rlp_node(&mut prefix_set, &mut buffers, &mut temp_rlp_buf);
889        }
890        self.rlp_buf = temp_rlp_buf;
891    }
892
893    /// Returns a list of (level, path) tuples identifying the nodes that have changed at the
894    /// specified depth, along with a new prefix set for the paths above the provided depth that
895    /// remain unchanged.
896    ///
897    /// Leaf nodes with a depth less than `depth` are returned too.
898    ///
899    /// This method helps optimize hash recalculations by identifying which specific
900    /// nodes need to be updated at each level of the trie.
901    ///
902    /// # Parameters
903    ///
904    /// - `prefix_set`: The current prefix set tracking which paths need updates.
905    /// - `depth`: The minimum depth (relative to the root) to include nodes in the targets.
906    ///
907    /// # Returns
908    ///
909    /// A tuple containing:
910    /// - A vector of `(level, Nibbles)` pairs for nodes that require updates at or below the
911    ///   specified depth.
912    /// - A `PrefixSetMut` containing paths shallower than the specified depth that still need to be
913    ///   tracked for future updates.
914    fn get_changed_nodes_at_depth(
915        &self,
916        prefix_set: &mut PrefixSet,
917        depth: usize,
918    ) -> (Vec<(usize, Nibbles)>, PrefixSetMut) {
919        let mut unchanged_prefix_set = PrefixSetMut::default();
920        let mut paths = Vec::from([(Nibbles::default(), 0)]);
921        let mut targets = Vec::new();
922
923        while let Some((mut path, level)) = paths.pop() {
924            match self.nodes.get(&path).unwrap() {
925                SparseNode::Empty | SparseNode::Hash(_) => {}
926                SparseNode::Leaf { key: _, hash, is_private: _ } => {
927                    if hash.is_some() && !prefix_set.contains(&path) {
928                        continue;
929                    }
930
931                    targets.push((level, path));
932                }
933                SparseNode::Extension { key, hash, store_in_db_trie: _ } => {
934                    if hash.is_some() && !prefix_set.contains(&path) {
935                        continue;
936                    }
937
938                    if level >= depth {
939                        targets.push((level, path));
940                    } else {
941                        unchanged_prefix_set.insert(path.clone());
942
943                        path.extend_from_slice_unchecked(key);
944                        paths.push((path, level + 1));
945                    }
946                }
947                SparseNode::Branch { state_mask, hash, store_in_db_trie: _ } => {
948                    if hash.is_some() && !prefix_set.contains(&path) {
949                        continue;
950                    }
951
952                    if level >= depth {
953                        targets.push((level, path));
954                    } else {
955                        unchanged_prefix_set.insert(path.clone());
956
957                        for bit in CHILD_INDEX_RANGE.rev() {
958                            if state_mask.is_bit_set(bit) {
959                                let mut child_path = path.clone();
960                                child_path.push_unchecked(bit);
961                                paths.push((child_path, level + 1));
962                            }
963                        }
964                    }
965                }
966            }
967        }
968
969        (targets, unchanged_prefix_set)
970    }
971
972    /// Look up or calculate the RLP of the node at the root path.
973    ///
974    /// # Panics
975    ///
976    /// If the node at provided path does not exist.
977    pub fn rlp_node_allocate(&mut self, prefix_set: &mut PrefixSet) -> RlpNode {
978        let mut buffers = RlpNodeBuffers::new_with_root_path();
979        let mut temp_rlp_buf = core::mem::take(&mut self.rlp_buf);
980        let result = self.rlp_node(prefix_set, &mut buffers, &mut temp_rlp_buf);
981        self.rlp_buf = temp_rlp_buf;
982
983        result
984    }
985
986    /// Looks up or computes the RLP encoding of the node specified by the current
987    /// path in the provided buffers.
988    ///
989    /// The function uses a stack (`RlpNodeBuffers::path_stack`) to track the traversal and
990    /// accumulate RLP encodings.
991    ///
992    /// # Parameters
993    ///
994    /// - `prefix_set`: The set of trie paths that need their nodes updated.
995    /// - `buffers`: The reusable buffers for stack management and temporary RLP values.
996    ///
997    /// # Panics
998    ///
999    /// If the node at provided path does not exist.
1000    pub fn rlp_node(
1001        &mut self,
1002        prefix_set: &mut PrefixSet,
1003        buffers: &mut RlpNodeBuffers,
1004        rlp_buf: &mut Vec<u8>,
1005    ) -> RlpNode {
1006        let _starting_path = buffers.path_stack.last().map(|item| item.path.clone());
1007
1008        'main: while let Some(RlpNodePathStackItem { level, path, mut is_in_prefix_set }) =
1009            buffers.path_stack.pop()
1010        {
1011            let node = self.nodes.get_mut(&path).unwrap();
1012            trace!(
1013                target: "trie::sparse",
1014                ?_starting_path,
1015                ?level,
1016                ?path,
1017                ?is_in_prefix_set,
1018                ?node,
1019                "Popped node from path stack"
1020            );
1021
1022            // Check if the path is in the prefix set.
1023            // First, check the cached value. If it's `None`, then check the prefix set, and update
1024            // the cached value.
1025            let mut prefix_set_contains =
1026                |path: &Nibbles| *is_in_prefix_set.get_or_insert_with(|| prefix_set.contains(path));
1027
1028            let (rlp_node, node_type) = match node {
1029                SparseNode::Empty => (RlpNode::word_rlp(&EMPTY_ROOT_HASH), SparseNodeType::Empty),
1030                SparseNode::Hash(hash) => (RlpNode::word_rlp(hash), SparseNodeType::Hash),
1031                SparseNode::Leaf { key, hash, is_private } => {
1032                    let mut path = path.clone();
1033                    path.extend_from_slice_unchecked(key);
1034                    if let Some(hash) = hash.filter(|_| !prefix_set_contains(&path)) {
1035                        (RlpNode::word_rlp(&hash), SparseNodeType::Leaf)
1036                    } else {
1037                        let value = self.values.get(&path).unwrap();
1038                        rlp_buf.clear();
1039                        let rlp_node = LeafNodeRef { key, value, is_private }.rlp(rlp_buf);
1040                        *hash = rlp_node.as_hash();
1041                        (rlp_node, SparseNodeType::Leaf)
1042                    }
1043                }
1044                SparseNode::Extension { key, hash, store_in_db_trie } => {
1045                    let mut child_path = path.clone();
1046                    child_path.extend_from_slice_unchecked(key);
1047                    if let Some((hash, store_in_db_trie)) =
1048                        hash.zip(*store_in_db_trie).filter(|_| !prefix_set_contains(&path))
1049                    {
1050                        (
1051                            RlpNode::word_rlp(&hash),
1052                            SparseNodeType::Extension { store_in_db_trie: Some(store_in_db_trie) },
1053                        )
1054                    } else if buffers.rlp_node_stack.last().is_some_and(|e| e.path == child_path) {
1055                        let RlpNodeStackItem {
1056                            path: _,
1057                            rlp_node: child,
1058                            node_type: child_node_type,
1059                        } = buffers.rlp_node_stack.pop().unwrap();
1060                        rlp_buf.clear();
1061                        let rlp_node = ExtensionNodeRef::new(key, &child).rlp(rlp_buf);
1062                        *hash = rlp_node.as_hash();
1063
1064                        let store_in_db_trie_value = child_node_type.store_in_db_trie();
1065
1066                        trace!(
1067                            target: "trie::sparse",
1068                            ?path,
1069                            ?child_path,
1070                            ?child_node_type,
1071                            "Extension node"
1072                        );
1073
1074                        *store_in_db_trie = store_in_db_trie_value;
1075
1076                        (
1077                            rlp_node,
1078                            SparseNodeType::Extension {
1079                                // Inherit the `store_in_db_trie` flag from the child node, which is
1080                                // always the branch node
1081                                store_in_db_trie: store_in_db_trie_value,
1082                            },
1083                        )
1084                    } else {
1085                        // need to get rlp node for child first
1086                        buffers.path_stack.extend([
1087                            RlpNodePathStackItem { level, path, is_in_prefix_set },
1088                            RlpNodePathStackItem {
1089                                level: level + 1,
1090                                path: child_path,
1091                                is_in_prefix_set: None,
1092                            },
1093                        ]);
1094                        continue;
1095                    }
1096                }
1097                SparseNode::Branch { state_mask, hash, store_in_db_trie } => {
1098                    if let Some((hash, store_in_db_trie)) =
1099                        hash.zip(*store_in_db_trie).filter(|_| !prefix_set_contains(&path))
1100                    {
1101                        buffers.rlp_node_stack.push(RlpNodeStackItem {
1102                            path,
1103                            rlp_node: RlpNode::word_rlp(&hash),
1104                            node_type: SparseNodeType::Branch {
1105                                store_in_db_trie: Some(store_in_db_trie),
1106                            },
1107                        });
1108                        continue;
1109                    }
1110                    let retain_updates = self.updates.is_some() && prefix_set_contains(&path);
1111
1112                    buffers.branch_child_buf.clear();
1113                    // Walk children in a reverse order from `f` to `0`, so we pop the `0` first
1114                    // from the stack and keep walking in the sorted order.
1115                    for bit in CHILD_INDEX_RANGE.rev() {
1116                        if state_mask.is_bit_set(bit) {
1117                            let mut child = path.clone();
1118                            child.push_unchecked(bit);
1119                            buffers.branch_child_buf.push(child);
1120                        }
1121                    }
1122
1123                    buffers
1124                        .branch_value_stack_buf
1125                        .resize(buffers.branch_child_buf.len(), Default::default());
1126                    let mut added_children = false;
1127
1128                    let mut tree_mask = TrieMask::default();
1129                    let mut hash_mask = TrieMask::default();
1130                    let mut hashes = Vec::new();
1131                    for (i, child_path) in buffers.branch_child_buf.iter().enumerate() {
1132                        if buffers.rlp_node_stack.last().is_some_and(|e| &e.path == child_path) {
1133                            let RlpNodeStackItem {
1134                                path: _,
1135                                rlp_node: child,
1136                                node_type: child_node_type,
1137                            } = buffers.rlp_node_stack.pop().unwrap();
1138
1139                            // Update the masks only if we need to retain trie updates
1140                            if retain_updates {
1141                                // SAFETY: it's a child, so it's never empty
1142                                let last_child_nibble = child_path.last().unwrap();
1143
1144                                // Determine whether we need to set trie mask bit.
1145                                let should_set_tree_mask_bit = if let Some(store_in_db_trie) =
1146                                    child_node_type.store_in_db_trie()
1147                                {
1148                                    // A branch or an extension node explicitly set the
1149                                    // `store_in_db_trie` flag
1150                                    store_in_db_trie
1151                                } else {
1152                                    // A blinded node has the tree mask bit set
1153                                    child_node_type.is_hash() &&
1154                                        self.branch_node_tree_masks.get(&path).is_some_and(
1155                                            |mask| mask.is_bit_set(last_child_nibble),
1156                                        )
1157                                };
1158                                if should_set_tree_mask_bit {
1159                                    tree_mask.set_bit(last_child_nibble);
1160                                }
1161
1162                                // Set the hash mask. If a child node is a revealed branch node OR
1163                                // is a blinded node that has its hash mask bit set according to the
1164                                // database, set the hash mask bit and save the hash.
1165                                let hash = child.as_hash().filter(|_| {
1166                                    child_node_type.is_branch() ||
1167                                        (child_node_type.is_hash() &&
1168                                            self.branch_node_hash_masks
1169                                                .get(&path)
1170                                                .is_some_and(|mask| {
1171                                                    mask.is_bit_set(last_child_nibble)
1172                                                }))
1173                                });
1174                                if let Some(hash) = hash {
1175                                    hash_mask.set_bit(last_child_nibble);
1176                                    hashes.push(hash);
1177                                }
1178                            }
1179
1180                            // Insert children in the resulting buffer in a normal order,
1181                            // because initially we iterated in reverse.
1182                            // SAFETY: i < len and len is never 0
1183                            let original_idx = buffers.branch_child_buf.len() - i - 1;
1184                            buffers.branch_value_stack_buf[original_idx] = child;
1185                            added_children = true;
1186                        } else {
1187                            debug_assert!(!added_children);
1188                            buffers.path_stack.push(RlpNodePathStackItem {
1189                                level,
1190                                path,
1191                                is_in_prefix_set,
1192                            });
1193                            buffers.path_stack.extend(buffers.branch_child_buf.drain(..).map(
1194                                |path| RlpNodePathStackItem {
1195                                    level: level + 1,
1196                                    path,
1197                                    is_in_prefix_set: None,
1198                                },
1199                            ));
1200                            continue 'main;
1201                        }
1202                    }
1203
1204                    trace!(
1205                        target: "trie::sparse",
1206                        ?path,
1207                        ?tree_mask,
1208                        ?hash_mask,
1209                        "Branch node masks"
1210                    );
1211
1212                    rlp_buf.clear();
1213                    let branch_node_ref =
1214                        BranchNodeRef::new(&buffers.branch_value_stack_buf, *state_mask);
1215                    let rlp_node = branch_node_ref.rlp(rlp_buf);
1216                    *hash = rlp_node.as_hash();
1217
1218                    // Save a branch node update only if it's not a root node, and we need to
1219                    // persist updates.
1220                    let store_in_db_trie_value = if let Some(updates) =
1221                        self.updates.as_mut().filter(|_| retain_updates && !path.is_empty())
1222                    {
1223                        let store_in_db_trie = !tree_mask.is_empty() || !hash_mask.is_empty();
1224                        if store_in_db_trie {
1225                            // Store in DB trie if there are either any children that are stored in
1226                            // the DB trie, or any children represent hashed values
1227                            hashes.reverse();
1228                            let branch_node = BranchNodeCompact::new(
1229                                *state_mask,
1230                                tree_mask,
1231                                hash_mask,
1232                                hashes,
1233                                hash.filter(|_| path.is_empty()),
1234                            );
1235                            updates.updated_nodes.insert(path.clone(), branch_node);
1236                        } else if self
1237                            .branch_node_tree_masks
1238                            .get(&path)
1239                            .is_some_and(|mask| !mask.is_empty()) ||
1240                            self.branch_node_hash_masks
1241                                .get(&path)
1242                                .is_some_and(|mask| !mask.is_empty())
1243                        {
1244                            // If new tree and hash masks are empty, but previously they weren't, we
1245                            // need to remove the node update and add the node itself to the list of
1246                            // removed nodes.
1247                            updates.updated_nodes.remove(&path);
1248                            updates.removed_nodes.insert(path.clone());
1249                        } else if self
1250                            .branch_node_hash_masks
1251                            .get(&path)
1252                            .is_none_or(|mask| mask.is_empty()) &&
1253                            self.branch_node_hash_masks
1254                                .get(&path)
1255                                .is_none_or(|mask| mask.is_empty())
1256                        {
1257                            // If new tree and hash masks are empty, and they were previously empty
1258                            // as well, we need to remove the node update.
1259                            updates.updated_nodes.remove(&path);
1260                        }
1261
1262                        store_in_db_trie
1263                    } else {
1264                        false
1265                    };
1266                    *store_in_db_trie = Some(store_in_db_trie_value);
1267
1268                    (
1269                        rlp_node,
1270                        SparseNodeType::Branch { store_in_db_trie: Some(store_in_db_trie_value) },
1271                    )
1272                }
1273            };
1274
1275            trace!(
1276                target: "trie::sparse",
1277                ?_starting_path,
1278                ?level,
1279                ?path,
1280                ?node,
1281                ?node_type,
1282                ?is_in_prefix_set,
1283                "Added node to rlp node stack"
1284            );
1285
1286            buffers.rlp_node_stack.push(RlpNodeStackItem { path, rlp_node, node_type });
1287        }
1288
1289        debug_assert_eq!(buffers.rlp_node_stack.len(), 1);
1290        buffers.rlp_node_stack.pop().unwrap().rlp_node
1291    }
1292}
1293
1294/// Error type for a leaf lookup operation
1295#[derive(Debug, Clone, PartialEq, Eq)]
1296pub enum LeafLookupError {
1297    /// The path leads to a blinded node, cannot determine if leaf exists.
1298    /// This means the witness is not complete.
1299    BlindedNode {
1300        /// Path to the blinded node.
1301        path: Nibbles,
1302        /// Hash of the blinded node.
1303        hash: B256,
1304    },
1305    /// The path leads to a leaf with a different value than expected.
1306    /// This means the witness is malformed.
1307    ValueMismatch {
1308        /// Path to the leaf.
1309        path: Nibbles,
1310        /// Expected value.
1311        expected: Option<Vec<u8>>,
1312        /// Actual value found.
1313        actual: Vec<u8>,
1314    },
1315}
1316
1317/// Success value for a leaf lookup operation
1318#[derive(Debug, Clone, PartialEq, Eq)]
1319pub enum LeafLookup {
1320    /// Leaf exists with expected value.
1321    Exists,
1322    /// Leaf does not exist (exclusion proof found).
1323    NonExistent {
1324        /// Path where the search diverged from the target path.
1325        diverged_at: Nibbles,
1326    },
1327}
1328
1329impl<P: BlindedProvider> RevealedSparseTrie<P> {
1330    /// This clears all data structures in the sparse trie, keeping the backing data structures
1331    /// allocated.
1332    ///
1333    /// This is useful for reusing the trie without needing to reallocate memory.
1334    pub fn clear(&mut self) {
1335        self.nodes.clear();
1336        self.branch_node_tree_masks.clear();
1337        self.branch_node_hash_masks.clear();
1338        self.values.clear();
1339        self.prefix_set.clear();
1340        if let Some(updates) = self.updates.as_mut() {
1341            updates.clear()
1342        }
1343        self.rlp_buf.clear();
1344    }
1345
1346    /// Attempts to find a leaf node at the specified path.
1347    ///
1348    /// This method traverses the trie from the root down to the given path, checking
1349    /// if a leaf exists at that path. It can be used to verify the existence of a leaf
1350    /// or to generate an exclusion proof (proof that a leaf does not exist).
1351    ///
1352    /// # Parameters
1353    ///
1354    /// - `path`: The path to search for.
1355    /// - `expected_value`: Optional expected value. If provided, will verify the leaf value
1356    ///   matches.
1357    ///
1358    /// # Returns
1359    ///
1360    /// - `Ok(LeafLookup::Exists)` if the leaf exists with the expected value.
1361    /// - `Ok(LeafLookup::NonExistent)` if the leaf definitely does not exist (exclusion proof).
1362    /// - `Err(LeafLookupError)` if the search encountered a blinded node or found a different
1363    ///   value.
1364    pub fn find_leaf(
1365        &self,
1366        path: &Nibbles,
1367        expected_value: Option<&Vec<u8>>,
1368    ) -> Result<LeafLookup, LeafLookupError> {
1369        // Helper function to check if a value matches the expected value
1370        fn check_value_match(
1371            actual_value: &Vec<u8>,
1372            expected_value: Option<&Vec<u8>>,
1373            path: &Nibbles,
1374        ) -> Result<(), LeafLookupError> {
1375            if let Some(expected) = expected_value {
1376                if actual_value != expected {
1377                    return Err(LeafLookupError::ValueMismatch {
1378                        path: path.clone(),
1379                        expected: Some(expected.clone()),
1380                        actual: actual_value.clone(),
1381                    });
1382                }
1383            }
1384            Ok(())
1385        }
1386
1387        let mut current = Nibbles::default(); // Start at the root
1388
1389        // Inclusion proof
1390        //
1391        // First, do a quick check if the value exists in our values map.
1392        // We assume that if there exists a leaf node, then its value will
1393        // be in the `values` map.
1394        if let Some(actual_value) = self.values.get(path) {
1395            // We found the leaf, check if the value matches (if expected value was provided)
1396            check_value_match(actual_value, expected_value, path)?;
1397            return Ok(LeafLookup::Exists);
1398        }
1399
1400        // If the value does not exist in the `values` map, then this means that the leaf either:
1401        // - Does not exist in the trie
1402        // - Is missing from the witness
1403        // We traverse the trie to find the location where this leaf would have been, showing
1404        // that it is not in the trie. Or we find a blinded node, showing that the witness is
1405        // not complete.
1406        while current.len() < path.len() {
1407            match self.nodes.get(&current) {
1408                Some(SparseNode::Empty) | None => {
1409                    // None implies no node is at the current path (even in the full trie)
1410                    // Empty node means there is a node at this path and it is "Empty"
1411                    return Ok(LeafLookup::NonExistent { diverged_at: current });
1412                }
1413                Some(&SparseNode::Hash(hash)) => {
1414                    // We hit a blinded node - cannot determine if leaf exists
1415                    return Err(LeafLookupError::BlindedNode { path: current.clone(), hash });
1416                }
1417                Some(SparseNode::Leaf { key, .. }) => {
1418                    // We found a leaf node before reaching our target depth
1419
1420                    // Temporarily append the leaf key to `current`
1421                    let saved_len = current.len();
1422                    current.extend_from_slice_unchecked(key);
1423
1424                    if &current == path {
1425                        // This should have been handled by our initial values map check
1426                        if let Some(value) = self.values.get(path) {
1427                            check_value_match(value, expected_value, path)?;
1428                            return Ok(LeafLookup::Exists);
1429                        }
1430                    }
1431
1432                    let diverged_at = current.slice(..saved_len);
1433
1434                    // The leaf node's path doesn't match our target path,
1435                    // providing an exclusion proof
1436                    return Ok(LeafLookup::NonExistent { diverged_at });
1437                }
1438                Some(SparseNode::Extension { key, .. }) => {
1439                    // Temporarily append the extension key to `current`
1440                    let saved_len = current.len();
1441                    current.extend_from_slice_unchecked(key);
1442
1443                    if path.len() < current.len() || !path.starts_with(&current) {
1444                        let diverged_at = current.slice(..saved_len);
1445                        current.truncate(saved_len); // restore
1446                        return Ok(LeafLookup::NonExistent { diverged_at });
1447                    }
1448                    // Prefix matched, so we keep walking with the longer `current`.
1449                }
1450                Some(SparseNode::Branch { state_mask, .. }) => {
1451                    // Check if branch has a child at the next nibble in our path
1452                    let nibble = path[current.len()];
1453                    if !state_mask.is_bit_set(nibble) {
1454                        // No child at this nibble - exclusion proof
1455                        return Ok(LeafLookup::NonExistent { diverged_at: current });
1456                    }
1457
1458                    // Continue down the branch
1459                    current.push_unchecked(nibble);
1460                }
1461            }
1462        }
1463
1464        // We've traversed to the end of the path and didn't find a leaf
1465        // Check if there's a node exactly at our target path
1466        match self.nodes.get(path) {
1467            Some(SparseNode::Leaf { key, .. }) if key.is_empty() => {
1468                // We found a leaf with an empty key (exact match)
1469                // This should be handled by the values map check above
1470                if let Some(value) = self.values.get(path) {
1471                    check_value_match(value, expected_value, path)?;
1472                    return Ok(LeafLookup::Exists);
1473                }
1474            }
1475            Some(&SparseNode::Hash(hash)) => {
1476                return Err(LeafLookupError::BlindedNode { path: path.clone(), hash });
1477            }
1478            _ => {
1479                // No leaf at exactly the target path
1480                let parent_path = if path.is_empty() {
1481                    Nibbles::default()
1482                } else {
1483                    path.slice(0..path.len() - 1)
1484                };
1485                return Ok(LeafLookup::NonExistent { diverged_at: parent_path });
1486            }
1487        }
1488
1489        // If we get here, there's no leaf at the target path
1490        Ok(LeafLookup::NonExistent { diverged_at: current })
1491    }
1492
1493    /// Updates or inserts a leaf node at the specified key path with the provided RLP-encoded
1494    /// value.
1495    ///
1496    /// This method updates the internal prefix set and, if the leaf did not previously exist,
1497    /// adjusts the trie structure by inserting new leaf nodes, splitting branch nodes, or
1498    /// collapsing extension nodes as needed.
1499    ///
1500    /// # Returns
1501    ///
1502    /// Returns `Ok(())` if the update is successful.
1503    ///
1504    /// Note: If an update requires revealing a blinded node, an error is returned if the blinded
1505    /// provider returns an error.
1506    pub fn update_leaf(
1507        &mut self,
1508        path: Nibbles,
1509        value: Vec<u8>,
1510        is_private: bool,
1511    ) -> SparseTrieResult<()> {
1512        self.prefix_set.insert(path.clone());
1513        let existing = self.values.insert(path.clone(), value);
1514        if existing.is_some() {
1515            // trie structure unchanged, return immediately
1516            return Ok(());
1517        }
1518
1519        let mut current = Nibbles::default();
1520        while let Some(node) = self.nodes.get_mut(&current) {
1521            let current_is_private = node.is_private();
1522            match node {
1523                SparseNode::Empty => {
1524                    *node = SparseNode::new_leaf(path, is_private);
1525                    break;
1526                }
1527                &mut SparseNode::Hash(hash) => {
1528                    return Err(SparseTrieErrorKind::BlindedNode { path: current, hash }.into())
1529                }
1530                SparseNode::Leaf { key: current_key, .. } => {
1531                    current.extend_from_slice_unchecked(current_key);
1532
1533                    // this leaf is being updated
1534                    if current == path {
1535                        unreachable!("we already checked leaf presence in the beginning");
1536                    }
1537
1538                    // find the common prefix
1539                    let common = current.common_prefix_length(&path);
1540
1541                    // update existing node
1542                    let new_ext_key = current.slice(current.len() - current_key.len()..common);
1543                    *node = SparseNode::new_ext(new_ext_key);
1544
1545                    // create a branch node and corresponding leaves
1546                    self.nodes.reserve(3);
1547                    self.nodes.insert(
1548                        current.slice(..common),
1549                        SparseNode::new_split_branch(current[common], path[common]),
1550                    );
1551                    self.nodes.insert(
1552                        path.slice(..=common),
1553                        SparseNode::new_leaf(path.slice(common + 1..), is_private),
1554                    );
1555                    self.nodes.insert(
1556                        current.slice(..=common),
1557                        SparseNode::new_leaf(
1558                            current.slice(common + 1..),
1559                            current_is_private.clone(),
1560                        ),
1561                    );
1562
1563                    break;
1564                }
1565                SparseNode::Extension { key, .. } => {
1566                    current.extend_from_slice(key);
1567
1568                    if !path.starts_with(&current) {
1569                        // find the common prefix
1570                        let common = current.common_prefix_length(&path);
1571                        *key = current.slice(current.len() - key.len()..common);
1572
1573                        // If branch node updates retention is enabled, we need to query the
1574                        // extension node child to later set the hash mask for a parent branch node
1575                        // correctly.
1576                        if self.updates.is_some() {
1577                            // Check if the extension node child is a hash that needs to be revealed
1578                            if self.nodes.get(&current).unwrap().is_hash() {
1579                                if let Some(RevealedNode { node, tree_mask, hash_mask }) =
1580                                    self.provider.blinded_node(&current)?
1581                                {
1582                                    let decoded = TrieNode::decode(&mut &node[..])?;
1583                                    trace!(
1584                                        target: "trie::sparse",
1585                                        ?current,
1586                                        ?decoded,
1587                                        ?tree_mask,
1588                                        ?hash_mask,
1589                                        "Revealing extension node child",
1590                                    );
1591                                    self.reveal_node(
1592                                        current.clone(),
1593                                        decoded,
1594                                        TrieMasks { hash_mask, tree_mask },
1595                                    )?;
1596                                }
1597                            }
1598                        }
1599
1600                        // create state mask for new branch node
1601                        // NOTE: this might overwrite the current extension node
1602                        self.nodes.reserve(3);
1603                        let branch = SparseNode::new_split_branch(current[common], path[common]);
1604                        self.nodes.insert(current.slice(..common), branch);
1605
1606                        // create new leaf
1607                        let new_leaf = SparseNode::new_leaf(path.slice(common + 1..), is_private);
1608                        self.nodes.insert(path.slice(..=common), new_leaf);
1609
1610                        // recreate extension to previous child if needed
1611                        let key = current.slice(common + 1..);
1612                        if !key.is_empty() {
1613                            self.nodes.insert(current.slice(..=common), SparseNode::new_ext(key));
1614                        }
1615
1616                        break;
1617                    }
1618                }
1619                SparseNode::Branch { state_mask, .. } => {
1620                    let nibble = path[current.len()];
1621                    current.push_unchecked(nibble);
1622                    if !state_mask.is_bit_set(nibble) {
1623                        state_mask.set_bit(nibble);
1624                        let new_leaf =
1625                            SparseNode::new_leaf(path.slice(current.len()..), is_private);
1626                        self.nodes.insert(current, new_leaf);
1627                        break;
1628                    }
1629                }
1630            };
1631        }
1632
1633        Ok(())
1634    }
1635
1636    /// Removes a leaf node from the trie at the specified key path.
1637    ///
1638    /// This function removes the leaf value from the internal values map and then traverses
1639    /// the trie to remove or adjust intermediate nodes, merging or collapsing them as necessary.
1640    ///
1641    /// # Returns
1642    ///
1643    /// Returns `Ok(())` if the leaf is successfully removed, otherwise returns an error
1644    /// if the leaf is not present or if a blinded node prevents removal.
1645    pub fn remove_leaf(&mut self, path: &Nibbles) -> SparseTrieResult<()> {
1646        if self.values.remove(path).is_none() {
1647            if let Some(&SparseNode::Hash(hash)) = self.nodes.get(path) {
1648                // Leaf is present in the trie, but it's blinded.
1649                return Err(SparseTrieErrorKind::BlindedNode { path: path.clone(), hash }.into());
1650            }
1651
1652            trace!(target: "trie::sparse", ?path, "Leaf node is not present in the trie");
1653            // Leaf is not present in the trie.
1654            return Ok(());
1655        }
1656        self.prefix_set.insert(path.clone());
1657
1658        // If the path wasn't present in `values`, we still need to walk the trie and ensure that
1659        // there is no node at the path. When a leaf node is a blinded `Hash`, it will have an entry
1660        // in `nodes`, but not in the `values`.
1661
1662        let mut removed_nodes = self.take_nodes_for_path(path)?;
1663        trace!(target: "trie::sparse", ?path, ?removed_nodes, "Removed nodes for path");
1664        // Pop the first node from the stack which is the leaf node we want to remove.
1665        let mut child = removed_nodes.pop().expect("leaf exists");
1666        #[cfg(debug_assertions)]
1667        {
1668            let mut child_path = child.path.clone();
1669            let SparseNode::Leaf { key, .. } = &child.node else { panic!("expected leaf node") };
1670            child_path.extend_from_slice_unchecked(key);
1671            assert_eq!(&child_path, path);
1672        }
1673
1674        // If we don't have any other removed nodes, insert an empty node at the root.
1675        if removed_nodes.is_empty() {
1676            debug_assert!(self.nodes.is_empty());
1677            self.nodes.insert(Nibbles::default(), SparseNode::Empty);
1678
1679            return Ok(());
1680        }
1681
1682        // Walk the stack of removed nodes from the back and re-insert them back into the trie,
1683        // adjusting the node type as needed.
1684        while let Some(removed_node) = removed_nodes.pop() {
1685            let removed_path = removed_node.path;
1686
1687            let new_node = match &removed_node.node {
1688                SparseNode::Empty => return Err(SparseTrieErrorKind::Blind.into()),
1689                &SparseNode::Hash(hash) => {
1690                    return Err(SparseTrieErrorKind::BlindedNode { path: removed_path, hash }.into())
1691                }
1692                SparseNode::Leaf { .. } => {
1693                    unreachable!("we already popped the leaf node")
1694                }
1695                SparseNode::Extension { key, .. } => {
1696                    // If the node is an extension node, we need to look at its child to see if we
1697                    // need to merge them.
1698                    match &child.node {
1699                        SparseNode::Empty => return Err(SparseTrieErrorKind::Blind.into()),
1700                        &SparseNode::Hash(hash) => {
1701                            return Err(
1702                                SparseTrieErrorKind::BlindedNode { path: child.path, hash }.into()
1703                            )
1704                        }
1705                        // For a leaf node, we collapse the extension node into a leaf node,
1706                        // extending the key. While it's impossible to encounter an extension node
1707                        // followed by a leaf node in a complete trie, it's possible here because we
1708                        // could have downgraded the extension node's child into a leaf node from
1709                        // another node type.
1710                        SparseNode::Leaf { key: leaf_key, hash: _, is_private } => {
1711                            self.nodes.remove(&child.path);
1712
1713                            let mut new_key = key.clone();
1714                            new_key.extend_from_slice_unchecked(leaf_key);
1715                            SparseNode::new_leaf(new_key, *is_private)
1716                        }
1717                        // For an extension node, we collapse them into one extension node,
1718                        // extending the key
1719                        SparseNode::Extension { key: extension_key, .. } => {
1720                            self.nodes.remove(&child.path);
1721
1722                            let mut new_key = key.clone();
1723                            new_key.extend_from_slice_unchecked(extension_key);
1724                            SparseNode::new_ext(new_key)
1725                        }
1726                        // For a branch node, we just leave the extension node as-is.
1727                        SparseNode::Branch { .. } => removed_node.node,
1728                    }
1729                }
1730                &SparseNode::Branch { mut state_mask, hash: _, store_in_db_trie: _ } => {
1731                    // If the node is a branch node, we need to check the number of children left
1732                    // after deleting the child at the given nibble.
1733
1734                    if let Some(removed_nibble) = removed_node.unset_branch_nibble {
1735                        state_mask.unset_bit(removed_nibble);
1736                    }
1737
1738                    // If only one child is left set in the branch node, we need to collapse it.
1739                    if state_mask.count_bits() == 1 {
1740                        let child_nibble =
1741                            state_mask.first_set_bit_index().expect("state mask is not empty");
1742
1743                        // Get full path of the only child node left.
1744                        let mut child_path = removed_path.clone();
1745                        child_path.push_unchecked(child_nibble);
1746
1747                        trace!(target: "trie::sparse", ?removed_path, ?child_path, "Branch node has only one child");
1748
1749                        if self.nodes.get(&child_path).unwrap().is_hash() {
1750                            trace!(target: "trie::sparse", ?child_path, "Retrieving remaining blinded branch child");
1751                            if let Some(RevealedNode { node, tree_mask, hash_mask }) =
1752                                self.provider.blinded_node(&child_path)?
1753                            {
1754                                let decoded = TrieNode::decode(&mut &node[..])?;
1755                                trace!(
1756                                    target: "trie::sparse",
1757                                    ?child_path,
1758                                    ?decoded,
1759                                    ?tree_mask,
1760                                    ?hash_mask,
1761                                    "Revealing remaining blinded branch child"
1762                                );
1763                                self.reveal_node(
1764                                    child_path.clone(),
1765                                    decoded,
1766                                    TrieMasks { hash_mask, tree_mask },
1767                                )?;
1768                            }
1769                        }
1770
1771                        // Get the only child node.
1772                        let child = self.nodes.get(&child_path).unwrap();
1773
1774                        let mut delete_child = false;
1775                        let new_node = match child {
1776                            SparseNode::Empty => return Err(SparseTrieErrorKind::Blind.into()),
1777                            &SparseNode::Hash(hash) => {
1778                                return Err(SparseTrieErrorKind::BlindedNode {
1779                                    path: child_path,
1780                                    hash,
1781                                }
1782                                .into())
1783                            }
1784                            // If the only child is a leaf node, we downgrade the branch node into a
1785                            // leaf node, prepending the nibble to the key, and delete the old
1786                            // child.
1787                            SparseNode::Leaf { key, hash: _, is_private } => {
1788                                delete_child = true;
1789
1790                                let mut new_key = Nibbles::from_nibbles_unchecked([child_nibble]);
1791                                new_key.extend_from_slice_unchecked(key);
1792                                SparseNode::new_leaf(new_key, *is_private)
1793                            }
1794                            // If the only child node is an extension node, we downgrade the branch
1795                            // node into an even longer extension node, prepending the nibble to the
1796                            // key, and delete the old child.
1797                            SparseNode::Extension { key, .. } => {
1798                                delete_child = true;
1799
1800                                let mut new_key = Nibbles::from_nibbles_unchecked([child_nibble]);
1801                                new_key.extend_from_slice_unchecked(key);
1802                                SparseNode::new_ext(new_key)
1803                            }
1804                            // If the only child is a branch node, we downgrade the current branch
1805                            // node into a one-nibble extension node.
1806                            SparseNode::Branch { .. } => {
1807                                SparseNode::new_ext(Nibbles::from_nibbles_unchecked([child_nibble]))
1808                            }
1809                        };
1810
1811                        if delete_child {
1812                            self.nodes.remove(&child_path);
1813                        }
1814
1815                        if let Some(updates) = self.updates.as_mut() {
1816                            updates.updated_nodes.remove(&removed_path);
1817                            updates.removed_nodes.insert(removed_path.clone());
1818                        }
1819
1820                        new_node
1821                    }
1822                    // If more than one child is left set in the branch, we just re-insert it as-is.
1823                    else {
1824                        SparseNode::new_branch(state_mask)
1825                    }
1826                }
1827            };
1828
1829            child = RemovedSparseNode {
1830                path: removed_path.clone(),
1831                node: new_node.clone(),
1832                unset_branch_nibble: None,
1833            };
1834            trace!(target: "trie::sparse", ?removed_path, ?new_node, "Re-inserting the node");
1835            self.nodes.insert(removed_path, new_node);
1836        }
1837
1838        Ok(())
1839    }
1840}
1841
1842/// Enum representing sparse trie node type.
1843#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1844enum SparseNodeType {
1845    /// Empty trie node.
1846    Empty,
1847    /// A placeholder that stores only the hash for a node that has not been fully revealed.
1848    Hash,
1849    /// Sparse leaf node.
1850    Leaf,
1851    /// Sparse extension node.
1852    Extension {
1853        /// A flag indicating whether the extension node should be stored in the database.
1854        store_in_db_trie: Option<bool>,
1855    },
1856    /// Sparse branch node.
1857    Branch {
1858        /// A flag indicating whether the branch node should be stored in the database.
1859        store_in_db_trie: Option<bool>,
1860    },
1861}
1862
1863impl SparseNodeType {
1864    const fn is_hash(&self) -> bool {
1865        matches!(self, Self::Hash)
1866    }
1867
1868    const fn is_branch(&self) -> bool {
1869        matches!(self, Self::Branch { .. })
1870    }
1871
1872    const fn store_in_db_trie(&self) -> Option<bool> {
1873        match *self {
1874            Self::Extension { store_in_db_trie } | Self::Branch { store_in_db_trie } => {
1875                store_in_db_trie
1876            }
1877            _ => None,
1878        }
1879    }
1880}
1881
1882/// Enum representing trie nodes in sparse trie.
1883#[derive(Debug, Clone, PartialEq, Eq)]
1884pub enum SparseNode {
1885    /// Empty trie node.
1886    Empty,
1887    /// The hash of the node that was not revealed.
1888    Hash(B256),
1889    /// Sparse leaf node with remaining key suffix.
1890    Leaf {
1891        /// Remaining key suffix for the leaf node.
1892        key: Nibbles,
1893        /// Pre-computed hash of the sparse node.
1894        /// Can be reused unless this trie path has been updated.
1895        hash: Option<B256>,
1896        /// whether the leaf node is private
1897        is_private: bool,
1898    },
1899    /// Sparse extension node with key.
1900    Extension {
1901        /// The key slice stored by this extension node.
1902        key: Nibbles,
1903        /// Pre-computed hash of the sparse node.
1904        /// Can be reused unless this trie path has been updated.
1905        ///
1906        /// If [`None`], then the value is not known and should be calculated from scratch.
1907        hash: Option<B256>,
1908        /// Pre-computed flag indicating whether the trie node should be stored in the database.
1909        /// Can be reused unless this trie path has been updated.
1910        ///
1911        /// If [`None`], then the value is not known and should be calculated from scratch.
1912        store_in_db_trie: Option<bool>,
1913    },
1914    /// Sparse branch node with state mask.
1915    Branch {
1916        /// The bitmask representing children present in the branch node.
1917        state_mask: TrieMask,
1918        /// Pre-computed hash of the sparse node.
1919        /// Can be reused unless this trie path has been updated.
1920        ///
1921        /// If [`None`], then the value is not known and should be calculated from scratch.
1922        hash: Option<B256>,
1923        /// Pre-computed flag indicating whether the trie node should be stored in the database.
1924        /// Can be reused unless this trie path has been updated.
1925        ///
1926        /// If [`None`], then the value is not known and should be calculated from scratch.
1927        store_in_db_trie: Option<bool>,
1928    },
1929}
1930
1931impl SparseNode {
1932    /// Create new sparse node from [`TrieNode`].
1933    pub fn from_node(node: TrieNode) -> Self {
1934        match node {
1935            TrieNode::EmptyRoot => Self::Empty,
1936            TrieNode::Leaf(leaf) => Self::new_leaf(leaf.key, leaf.is_private),
1937            TrieNode::Extension(ext) => Self::new_ext(ext.key),
1938            TrieNode::Branch(branch) => Self::new_branch(branch.state_mask),
1939        }
1940    }
1941
1942    /// Create new [`SparseNode::Branch`] from state mask.
1943    pub const fn new_branch(state_mask: TrieMask) -> Self {
1944        Self::Branch { state_mask, hash: None, store_in_db_trie: None }
1945    }
1946
1947    /// Create new [`SparseNode::Branch`] with two bits set.
1948    pub const fn new_split_branch(bit_a: u8, bit_b: u8) -> Self {
1949        let state_mask = TrieMask::new(
1950            // set bits for both children
1951            (1u16 << bit_a) | (1u16 << bit_b),
1952        );
1953        Self::Branch { state_mask, hash: None, store_in_db_trie: None }
1954    }
1955
1956    /// Create new [`SparseNode::Extension`] from the key slice.
1957    pub const fn new_ext(key: Nibbles) -> Self {
1958        Self::Extension { key, hash: None, store_in_db_trie: None }
1959    }
1960
1961    /// Create new [`SparseNode::Leaf`] from leaf key and value.
1962    pub const fn new_leaf(key: Nibbles, is_private: bool) -> Self {
1963        Self::Leaf { key, hash: None, is_private }
1964    }
1965
1966    /// Returns `true` if the node is a hash node.
1967    pub const fn is_hash(&self) -> bool {
1968        matches!(self, Self::Hash(_))
1969    }
1970
1971    /// returns if the node holds private state
1972    /// node is always public unless it is a leaf
1973    pub fn is_private(&self) -> bool {
1974        match self {
1975            Self::Leaf { is_private, .. } => *is_private,
1976            _ => false,
1977        }
1978    }
1979}
1980
1981/// A helper struct used to store information about a node that has been removed
1982/// during a deletion operation.
1983#[derive(Debug)]
1984struct RemovedSparseNode {
1985    /// The path at which the node was located.
1986    path: Nibbles,
1987    /// The removed node
1988    node: SparseNode,
1989    /// For branch nodes, an optional nibble that should be unset due to the node being removed.
1990    ///
1991    /// During leaf deletion, this identifies the specific branch nibble path that
1992    /// connects to the leaf being deleted. Then when restructuring the trie after deletion,
1993    /// this nibble position will be cleared from the branch node's to
1994    /// indicate that the child no longer exists.
1995    ///
1996    /// This is only set for branch nodes that have a direct path to the leaf being deleted.
1997    unset_branch_nibble: Option<u8>,
1998}
1999
2000/// Collection of reusable buffers for [`RevealedSparseTrie::rlp_node`] calculations.
2001///
2002/// These buffers reduce allocations when computing RLP representations during trie updates.
2003#[derive(Debug, Default)]
2004pub struct RlpNodeBuffers {
2005    /// Stack of RLP node paths
2006    path_stack: Vec<RlpNodePathStackItem>,
2007    /// Stack of RLP nodes
2008    rlp_node_stack: Vec<RlpNodeStackItem>,
2009    /// Reusable branch child path
2010    branch_child_buf: SmallVec<[Nibbles; 16]>,
2011    /// Reusable branch value stack
2012    branch_value_stack_buf: SmallVec<[RlpNode; 16]>,
2013}
2014
2015impl RlpNodeBuffers {
2016    /// Creates a new instance of buffers with the root path on the stack.
2017    fn new_with_root_path() -> Self {
2018        Self {
2019            path_stack: vec![RlpNodePathStackItem {
2020                level: 0,
2021                path: Nibbles::default(),
2022                is_in_prefix_set: None,
2023            }],
2024            rlp_node_stack: Vec::new(),
2025            branch_child_buf: SmallVec::<[Nibbles; 16]>::new_const(),
2026            branch_value_stack_buf: SmallVec::<[RlpNode; 16]>::new_const(),
2027        }
2028    }
2029}
2030
2031/// RLP node path stack item.
2032#[derive(Debug)]
2033struct RlpNodePathStackItem {
2034    /// Level at which the node is located. Higher numbers correspond to lower levels in the trie.
2035    level: usize,
2036    /// Path to the node.
2037    path: Nibbles,
2038    /// Whether the path is in the prefix set. If [`None`], then unknown yet.
2039    is_in_prefix_set: Option<bool>,
2040}
2041
2042/// RLP node stack item.
2043#[derive(Debug)]
2044struct RlpNodeStackItem {
2045    /// Path to the node.
2046    path: Nibbles,
2047    /// RLP node.
2048    rlp_node: RlpNode,
2049    /// Type of the node.
2050    node_type: SparseNodeType,
2051}
2052
2053/// Tracks modifications to the sparse trie structure.
2054///
2055/// Maintains references to both modified and pruned/removed branches, enabling
2056/// one to make batch updates to a persistent database.
2057#[derive(Debug, Clone, Default, PartialEq, Eq)]
2058pub struct SparseTrieUpdates {
2059    pub(crate) updated_nodes: HashMap<Nibbles, BranchNodeCompact>,
2060    pub(crate) removed_nodes: HashSet<Nibbles>,
2061    pub(crate) wiped: bool,
2062}
2063
2064impl SparseTrieUpdates {
2065    /// Create new wiped sparse trie updates.
2066    pub fn wiped() -> Self {
2067        Self { wiped: true, ..Default::default() }
2068    }
2069
2070    /// Clears the updates, but keeps the backing data structures allocated.
2071    ///
2072    /// Sets `wiped` to `false`.
2073    pub fn clear(&mut self) {
2074        self.updated_nodes.clear();
2075        self.removed_nodes.clear();
2076        self.wiped = false;
2077    }
2078}
2079
2080#[cfg(test)]
2081mod find_leaf_tests {
2082    use super::*;
2083    use crate::blinded::DefaultBlindedProvider;
2084    use alloy_primitives::map::foldhash::fast::RandomState;
2085    // Assuming this exists
2086    use alloy_rlp::Encodable;
2087    use assert_matches::assert_matches;
2088    use reth_primitives_traits::Account;
2089    use reth_trie_common::LeafNode;
2090
2091    // Helper to create some test values
2092    fn encode_value(nonce: u64) -> Vec<u8> {
2093        let account = Account { nonce, ..Default::default() };
2094        let trie_account = account.into_trie_account(EMPTY_ROOT_HASH);
2095        let mut buf = Vec::new();
2096        trie_account.encode(&mut buf);
2097        buf
2098    }
2099
2100    const VALUE_A: fn() -> Vec<u8> = || encode_value(1);
2101    const VALUE_B: fn() -> Vec<u8> = || encode_value(2);
2102
2103    #[test]
2104    fn find_leaf_existing_leaf() {
2105        // Create a simple trie with one leaf
2106        let mut sparse = RevealedSparseTrie::default();
2107        let path = Nibbles::from_nibbles([0x1, 0x2, 0x3]);
2108        let value = b"test_value".to_vec();
2109
2110        let is_private = false; // hardcode to false for legacy test
2111        sparse.update_leaf(path.clone(), value.clone(), is_private).unwrap();
2112
2113        // Check that the leaf exists
2114        let result = sparse.find_leaf(&path, None);
2115        assert_matches!(result, Ok(LeafLookup::Exists));
2116
2117        // Check with expected value matching
2118        let result = sparse.find_leaf(&path, Some(&value));
2119        assert_matches!(result, Ok(LeafLookup::Exists));
2120    }
2121
2122    #[test]
2123    fn find_leaf_value_mismatch() {
2124        // Create a simple trie with one leaf
2125        let mut sparse = RevealedSparseTrie::default();
2126        let path = Nibbles::from_nibbles([0x1, 0x2, 0x3]);
2127        let value = b"test_value".to_vec();
2128        let wrong_value = b"wrong_value".to_vec();
2129
2130        let is_private = false; // hardcode to false for legacy test
2131        sparse.update_leaf(path.clone(), value, is_private).unwrap();
2132
2133        // Check with wrong expected value
2134        let result = sparse.find_leaf(&path, Some(&wrong_value));
2135        assert_matches!(
2136            result,
2137            Err(LeafLookupError::ValueMismatch { path: p, expected: Some(e), actual: _a }) if p == path && e == wrong_value
2138        );
2139    }
2140
2141    #[test]
2142    fn find_leaf_not_found_empty_trie() {
2143        // Empty trie
2144        let sparse = RevealedSparseTrie::default();
2145        let path = Nibbles::from_nibbles([0x1, 0x2, 0x3]);
2146
2147        // Leaf should not exist
2148        let result = sparse.find_leaf(&path, None);
2149        assert_matches!(
2150            result,
2151            Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == Nibbles::default()
2152        );
2153    }
2154
2155    #[test]
2156    fn find_leaf_empty_trie() {
2157        let sparse = RevealedSparseTrie::<DefaultBlindedProvider>::default();
2158        let path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
2159
2160        let result = sparse.find_leaf(&path, None);
2161
2162        // In an empty trie, the search diverges immediately at the root.
2163        assert_matches!(result, Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == Nibbles::default());
2164    }
2165
2166    #[test]
2167    fn find_leaf_exists_no_value_check() {
2168        let mut sparse = RevealedSparseTrie::<DefaultBlindedProvider>::default();
2169        let path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
2170        let is_private = false; // hardcode to false for legacy test
2171        sparse.update_leaf(path.clone(), VALUE_A(), is_private).unwrap();
2172
2173        let result = sparse.find_leaf(&path, None);
2174        assert_matches!(result, Ok(LeafLookup::Exists));
2175    }
2176
2177    #[test]
2178    fn find_leaf_exists_with_value_check_ok() {
2179        let mut sparse = RevealedSparseTrie::<DefaultBlindedProvider>::default();
2180        let path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
2181        let value = VALUE_A();
2182        let is_private = false; // hardcode to false for legacy test
2183        sparse.update_leaf(path.clone(), value.clone(), is_private).unwrap();
2184
2185        let result = sparse.find_leaf(&path, Some(&value));
2186        assert_matches!(result, Ok(LeafLookup::Exists));
2187    }
2188
2189    #[test]
2190    fn find_leaf_exclusion_branch_divergence() {
2191        let mut sparse = RevealedSparseTrie::<DefaultBlindedProvider>::default();
2192        let path1 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); // Creates branch at 0x12
2193        let path2 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x5, 0x6]); // Belongs to same branch
2194        let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x7, 0x8]); // Diverges at nibble 7
2195
2196        let is_private = false; // hardcode to false for legacy test
2197        sparse.update_leaf(path1, VALUE_A(), is_private).unwrap();
2198        sparse.update_leaf(path2, VALUE_B(), is_private).unwrap();
2199
2200        let result = sparse.find_leaf(&search_path, None);
2201
2202        // Diverged at the branch node because nibble '7' is not present.
2203        let expected_divergence = Nibbles::from_nibbles_unchecked([0x1, 0x2]);
2204        assert_matches!(result, Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == expected_divergence);
2205    }
2206
2207    #[test]
2208    fn find_leaf_exclusion_extension_divergence() {
2209        let mut sparse = RevealedSparseTrie::<DefaultBlindedProvider>::default();
2210        // This will create an extension node at root with key 0x12
2211        let path1 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4, 0x5, 0x6]);
2212        // This path diverges from the extension key
2213        let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x7, 0x8]);
2214
2215        let is_private = false; // hardcode to false for legacy test
2216        sparse.update_leaf(path1, VALUE_A(), is_private).unwrap();
2217
2218        let result = sparse.find_leaf(&search_path, None);
2219
2220        // Diverged where the extension node started because the path doesn't match its key prefix.
2221        let expected_divergence = Nibbles::default();
2222        assert_matches!(result, Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == expected_divergence);
2223    }
2224
2225    #[test]
2226    fn find_leaf_exclusion_leaf_divergence() {
2227        let mut sparse = RevealedSparseTrie::<DefaultBlindedProvider>::default();
2228        let existing_leaf_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
2229        let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4, 0x5, 0x6]);
2230
2231        let is_private = false; // hardcode to false for legacy test
2232        sparse.update_leaf(existing_leaf_path, VALUE_A(), is_private).unwrap();
2233
2234        let result = sparse.find_leaf(&search_path, None);
2235
2236        // Diverged when it hit the leaf node at the root, because the search path is longer
2237        // than the leaf's key stored there. The code returns the path of the node (root)
2238        // where the divergence occurred.
2239        let expected_divergence = Nibbles::default();
2240        assert_matches!(result, Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == expected_divergence);
2241    }
2242
2243    #[test]
2244    fn find_leaf_exclusion_path_ends_at_branch() {
2245        let mut sparse = RevealedSparseTrie::<DefaultBlindedProvider>::default();
2246        let path1 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); // Creates branch at 0x12
2247        let path2 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x5, 0x6]);
2248        let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2]); // Path of the branch itself
2249
2250        let is_private = false; // hardcode to false for legacy test
2251        sparse.update_leaf(path1, VALUE_A(), is_private).unwrap();
2252        sparse.update_leaf(path2, VALUE_B(), is_private).unwrap();
2253
2254        let result = sparse.find_leaf(&search_path, None);
2255
2256        // The path ends, but the node at the path is a branch, not a leaf.
2257        // Diverged at the parent of the node found at the search path.
2258        let expected_divergence = Nibbles::from_nibbles_unchecked([0x1]);
2259        assert_matches!(result, Ok(LeafLookup::NonExistent { diverged_at }) if diverged_at == expected_divergence);
2260    }
2261
2262    #[test]
2263    fn find_leaf_error_blinded_node_at_leaf_path() {
2264        // Scenario: The node *at* the leaf path is blinded.
2265        let blinded_hash = B256::repeat_byte(0xBB);
2266        let leaf_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
2267
2268        let mut nodes = alloy_primitives::map::HashMap::with_hasher(RandomState::default());
2269        // Create path to the blinded node
2270        nodes.insert(
2271            Nibbles::default(),
2272            SparseNode::new_ext(Nibbles::from_nibbles_unchecked([0x1, 0x2])),
2273        ); // Ext 0x12
2274        nodes.insert(
2275            Nibbles::from_nibbles_unchecked([0x1, 0x2]),
2276            SparseNode::new_ext(Nibbles::from_nibbles_unchecked([0x3])),
2277        ); // Ext 0x123
2278        nodes.insert(
2279            Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3]),
2280            SparseNode::new_branch(TrieMask::new(0b10000)),
2281        ); // Branch at 0x123, child 4
2282        nodes.insert(leaf_path.clone(), SparseNode::Hash(blinded_hash)); // Blinded node at 0x1234
2283
2284        let sparse = RevealedSparseTrie {
2285            provider: DefaultBlindedProvider,
2286            nodes,
2287            branch_node_tree_masks: Default::default(),
2288            branch_node_hash_masks: Default::default(),
2289            /* The value is not in the values map, or else it would early return */
2290            values: Default::default(),
2291            prefix_set: Default::default(),
2292            updates: None,
2293            rlp_buf: Vec::new(),
2294        };
2295
2296        let result = sparse.find_leaf(&leaf_path, None);
2297
2298        // Should error because it hit the blinded node exactly at the leaf path
2299        assert_matches!(result, Err(LeafLookupError::BlindedNode { path, hash })
2300            if path == leaf_path && hash == blinded_hash
2301        );
2302    }
2303
2304    #[test]
2305    fn find_leaf_error_blinded_node() {
2306        let is_private = false; // legacy test does not use private storage
2307        let blinded_hash = B256::repeat_byte(0xAA);
2308        let path_to_blind = Nibbles::from_nibbles_unchecked([0x1]);
2309        let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
2310
2311        let mut nodes = HashMap::with_hasher(RandomState::default());
2312
2313        // Root is a branch with child 0x1 (blinded) and 0x5 (revealed leaf)
2314        // So we set Bit 1 and Bit 5 in the state_mask
2315        let state_mask = TrieMask::new(0b100010);
2316        nodes.insert(Nibbles::default(), SparseNode::new_branch(state_mask));
2317
2318        nodes.insert(path_to_blind.clone(), SparseNode::Hash(blinded_hash));
2319        let path_revealed = Nibbles::from_nibbles_unchecked([0x5]);
2320        let path_revealed_leaf = Nibbles::from_nibbles_unchecked([0x5, 0x6, 0x7, 0x8]);
2321        nodes.insert(
2322            path_revealed,
2323            SparseNode::new_leaf(Nibbles::from_nibbles_unchecked([0x6, 0x7, 0x8]), is_private),
2324        );
2325
2326        let mut values = HashMap::with_hasher(RandomState::default());
2327        values.insert(path_revealed_leaf, VALUE_A());
2328
2329        let sparse = RevealedSparseTrie {
2330            provider: DefaultBlindedProvider,
2331            nodes,
2332            branch_node_tree_masks: Default::default(),
2333            branch_node_hash_masks: Default::default(),
2334            values,
2335            prefix_set: Default::default(),
2336            updates: None,
2337            rlp_buf: Vec::new(),
2338        };
2339
2340        let result = sparse.find_leaf(&search_path, None);
2341
2342        // Should error because it hit the blinded node at path 0x1
2343        assert_matches!(result, Err(LeafLookupError::BlindedNode { path, hash })
2344            if path == path_to_blind && hash == blinded_hash
2345        );
2346    }
2347
2348    #[test]
2349    fn find_leaf_error_blinded_node_via_reveal() {
2350        let blinded_hash = B256::repeat_byte(0xAA);
2351        let path_to_blind = Nibbles::from_nibbles_unchecked([0x1]); // Path of the blinded node itself
2352        let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); // Path we will search for
2353
2354        let revealed_leaf_prefix = Nibbles::from_nibbles_unchecked([0x5]);
2355        let revealed_leaf_suffix = Nibbles::from_nibbles_unchecked([0x6, 0x7, 0x8]);
2356        let revealed_leaf_full_path = Nibbles::from_nibbles_unchecked([0x5, 0x6, 0x7, 0x8]);
2357        let revealed_value = VALUE_A();
2358
2359        // 1. Construct the RLP representation of the children for the root branch
2360        let rlp_node_child1 = RlpNode::word_rlp(&blinded_hash); // Blinded node
2361
2362        let is_private = false; // legacy test does not use private storage
2363        let leaf_node_child5 =
2364            LeafNode::new(revealed_leaf_suffix.clone(), revealed_value.clone(), is_private);
2365        let leaf_node_child5_rlp_buf = alloy_rlp::encode(&leaf_node_child5);
2366        let hash_of_child5 = keccak256(&leaf_node_child5_rlp_buf);
2367        let rlp_node_child5 = RlpNode::word_rlp(&hash_of_child5);
2368
2369        // 2. Construct the root BranchNode using the RLP of its children
2370        // The stack order depends on the bit indices (1 and 5)
2371        let root_branch_node = reth_trie_common::BranchNode::new(
2372            vec![rlp_node_child1, rlp_node_child5], // Child 1 first, then Child 5
2373            TrieMask::new(0b100010),                // Mask with bits 1 and 5 set
2374        );
2375        let root_trie_node = TrieNode::Branch(root_branch_node);
2376
2377        // 3. Initialize the sparse trie using from_root
2378        // This will internally create Hash nodes for paths "1" and "5" initially.
2379        let mut sparse = RevealedSparseTrie::from_root(root_trie_node, TrieMasks::none(), false)
2380            .expect("Failed to create trie from root");
2381
2382        // Assertions before we reveal child5
2383        assert_matches!(sparse.nodes.get(&Nibbles::default()), Some(SparseNode::Branch { state_mask, .. }) if *state_mask == TrieMask::new(0b100010)); // Here we check that 1 and 5 are set in the state_mask
2384        assert_matches!(sparse.nodes.get(&path_to_blind), Some(SparseNode::Hash(h)) if *h == blinded_hash );
2385        assert!(sparse.nodes.get(&revealed_leaf_prefix).unwrap().is_hash()); // Child 5 is initially a hash of its RLP
2386        assert!(sparse.values.is_empty());
2387
2388        // 4. Explicitly reveal the leaf node for child 5
2389        sparse
2390            .reveal_node(
2391                revealed_leaf_prefix.clone(),
2392                TrieNode::Leaf(leaf_node_child5),
2393                TrieMasks::none(),
2394            )
2395            .expect("Failed to reveal leaf node");
2396
2397        // Assertions after we reveal child 5
2398        assert_matches!(sparse.nodes.get(&Nibbles::default()), Some(SparseNode::Branch { state_mask, .. }) if *state_mask == TrieMask::new(0b100010));
2399        assert_matches!(sparse.nodes.get(&path_to_blind), Some(SparseNode::Hash(h)) if *h == blinded_hash );
2400        assert_matches!(sparse.nodes.get(&revealed_leaf_prefix), Some(SparseNode::Leaf { key, .. }) if *key == revealed_leaf_suffix);
2401        assert_eq!(sparse.values.get(&revealed_leaf_full_path), Some(&revealed_value));
2402
2403        let result = sparse.find_leaf(&search_path, None);
2404
2405        // 5. Assert the expected error
2406        // Should error because it hit the blinded node at path "1" only node at "5" was revealed
2407        assert_matches!(result, Err(LeafLookupError::BlindedNode { path, hash })
2408            if path == path_to_blind && hash == blinded_hash
2409        );
2410    }
2411}
2412
2413#[cfg(test)]
2414mod tests {
2415    use super::*;
2416    use alloy_primitives::{map::B256Set, U256};
2417    use alloy_rlp::Encodable;
2418    use assert_matches::assert_matches;
2419    use itertools::Itertools;
2420    use prop::sample::SizeRange;
2421    use proptest::prelude::*;
2422    use proptest_arbitrary_interop::arb;
2423    use reth_primitives_traits::Account;
2424    use reth_provider::{test_utils::create_test_provider_factory, TrieWriter};
2425    use reth_trie::{
2426        hashed_cursor::{noop::NoopHashedAccountCursor, HashedPostStateAccountCursor},
2427        node_iter::{TrieElement, TrieNodeIter},
2428        trie_cursor::{noop::NoopAccountTrieCursor, TrieCursor, TrieCursorFactory},
2429        walker::TrieWalker,
2430        BranchNode, ExtensionNode, HashedPostState, LeafNode,
2431    };
2432    use reth_trie_common::{
2433        proof::{ProofNodes, ProofRetainer},
2434        updates::TrieUpdates,
2435        HashBuilder,
2436    };
2437    use reth_trie_db::DatabaseTrieCursorFactory;
2438    use std::collections::{BTreeMap, BTreeSet};
2439
2440    /// Pad nibbles to the length of a B256 hash with zeros on the left.
2441    fn pad_nibbles_left(nibbles: Nibbles) -> Nibbles {
2442        let mut base =
2443            Nibbles::from_nibbles_unchecked(vec![0; B256::len_bytes() * 2 - nibbles.len()]);
2444        base.extend_from_slice_unchecked(&nibbles);
2445        base
2446    }
2447
2448    /// Pad nibbles to the length of a B256 hash with zeros on the right.
2449    fn pad_nibbles_right(mut nibbles: Nibbles) -> Nibbles {
2450        nibbles.extend_from_slice_unchecked(&vec![0; B256::len_bytes() * 2 - nibbles.len()]);
2451        nibbles
2452    }
2453
2454    /// Calculate the state root by feeding the provided state to the hash builder and retaining the
2455    /// proofs for the provided targets.
2456    ///
2457    /// Returns the state root and the retained proof nodes.
2458    fn run_hash_builder(
2459        state: impl IntoIterator<Item = (Nibbles, Account)> + Clone,
2460        trie_cursor: impl TrieCursor,
2461        destroyed_accounts: B256Set,
2462        proof_targets: impl IntoIterator<Item = Nibbles>,
2463    ) -> (B256, TrieUpdates, ProofNodes, HashMap<Nibbles, TrieMask>, HashMap<Nibbles, TrieMask>)
2464    {
2465        let is_private = false; // hardcode to false for legacy test, TODO: make a private equivalent
2466        let mut account_rlp = Vec::new();
2467
2468        let mut hash_builder = HashBuilder::default()
2469            .with_updates(true)
2470            .with_proof_retainer(ProofRetainer::from_iter(proof_targets));
2471
2472        let mut prefix_set = PrefixSetMut::default();
2473        prefix_set.extend_keys(state.clone().into_iter().map(|(nibbles, _)| nibbles));
2474        prefix_set.extend_keys(destroyed_accounts.iter().map(Nibbles::unpack));
2475        let walker =
2476            TrieWalker::state_trie(trie_cursor, prefix_set.freeze()).with_deletions_retained(true);
2477        let hashed_post_state = HashedPostState::default()
2478            .with_accounts(state.into_iter().map(|(nibbles, account)| {
2479                (nibbles.pack().into_inner().unwrap().into(), Some(account))
2480            }))
2481            .into_sorted();
2482        let mut node_iter = TrieNodeIter::state_trie(
2483            walker,
2484            HashedPostStateAccountCursor::new(
2485                NoopHashedAccountCursor::default(),
2486                hashed_post_state.accounts(),
2487            ),
2488        );
2489
2490        while let Some(node) = node_iter.try_next().unwrap() {
2491            match node {
2492                TrieElement::Branch(branch) => {
2493                    hash_builder.add_branch(branch.key, branch.value, branch.children_are_in_trie);
2494                }
2495                TrieElement::Leaf(key, account) => {
2496                    let account = account.into_trie_account(EMPTY_ROOT_HASH);
2497                    account.encode(&mut account_rlp);
2498
2499                    hash_builder.add_leaf(Nibbles::unpack(key), &account_rlp, is_private);
2500                    account_rlp.clear();
2501                }
2502            }
2503        }
2504        let root = hash_builder.root();
2505        let proof_nodes = hash_builder.take_proof_nodes();
2506        let branch_node_hash_masks = hash_builder
2507            .updated_branch_nodes
2508            .clone()
2509            .unwrap_or_default()
2510            .iter()
2511            .map(|(path, node)| (path.clone(), node.hash_mask))
2512            .collect();
2513        let branch_node_tree_masks = hash_builder
2514            .updated_branch_nodes
2515            .clone()
2516            .unwrap_or_default()
2517            .iter()
2518            .map(|(path, node)| (path.clone(), node.tree_mask))
2519            .collect();
2520
2521        let mut trie_updates = TrieUpdates::default();
2522        let removed_keys = node_iter.walker.take_removed_keys();
2523        trie_updates.finalize(hash_builder, removed_keys, destroyed_accounts);
2524
2525        (root, trie_updates, proof_nodes, branch_node_hash_masks, branch_node_tree_masks)
2526    }
2527
2528    /// Assert that the sparse trie nodes and the proof nodes from the hash builder are equal.
2529    fn assert_eq_sparse_trie_proof_nodes(
2530        sparse_trie: &RevealedSparseTrie,
2531        proof_nodes: ProofNodes,
2532    ) {
2533        let proof_nodes = proof_nodes
2534            .into_nodes_sorted()
2535            .into_iter()
2536            .map(|(path, node)| (path, TrieNode::decode(&mut node.as_ref()).unwrap()));
2537
2538        let sparse_nodes = sparse_trie.nodes.iter().sorted_by_key(|(path, _)| *path);
2539
2540        for ((proof_node_path, proof_node), (sparse_node_path, sparse_node)) in
2541            proof_nodes.zip(sparse_nodes)
2542        {
2543            assert_eq!(&proof_node_path, sparse_node_path);
2544
2545            let equals = match (&proof_node, &sparse_node) {
2546                // Both nodes are empty
2547                (TrieNode::EmptyRoot, SparseNode::Empty) => true,
2548                // Both nodes are branches and have the same state mask
2549                (
2550                    TrieNode::Branch(BranchNode { state_mask: proof_state_mask, .. }),
2551                    SparseNode::Branch { state_mask: sparse_state_mask, .. },
2552                ) => proof_state_mask == sparse_state_mask,
2553                // Both nodes are extensions and have the same key
2554                (
2555                    TrieNode::Extension(ExtensionNode { key: proof_key, .. }),
2556                    SparseNode::Extension { key: sparse_key, .. },
2557                ) => proof_key == sparse_key,
2558                // Both nodes are leaves and have the same key
2559                (
2560                    TrieNode::Leaf(LeafNode {
2561                        key: proof_key, is_private: proof_is_private, ..
2562                    }),
2563                    SparseNode::Leaf { key: sparse_key, is_private: sparse_is_private, .. },
2564                ) => proof_key == sparse_key && proof_is_private == sparse_is_private,
2565                // Empty and hash nodes are specific to the sparse trie, skip them
2566                (_, SparseNode::Empty | SparseNode::Hash(_)) => continue,
2567                _ => false,
2568            };
2569            assert!(
2570                equals,
2571                "path: {proof_node_path:?}\nproof node: {proof_node:?}\nsparse node: {sparse_node:?}"
2572            );
2573        }
2574    }
2575
2576    #[test]
2577    fn sparse_trie_is_blind() {
2578        assert!(SparseTrie::blind().is_blind());
2579        assert!(!SparseTrie::revealed_empty().is_blind());
2580    }
2581
2582    #[test]
2583    fn sparse_trie_empty_update_one() {
2584        let is_private = false; // hardcode to false for legacy test, TODO: make a private equivalent
2585        let key = Nibbles::unpack(B256::with_last_byte(42));
2586        let value = || Account::default();
2587        let value_encoded = || {
2588            let mut account_rlp = Vec::new();
2589            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
2590            account_rlp
2591        };
2592
2593        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
2594            run_hash_builder(
2595                [(key.clone(), value())],
2596                NoopAccountTrieCursor::default(),
2597                Default::default(),
2598                [key.clone()],
2599            );
2600
2601        let mut sparse = RevealedSparseTrie::default().with_updates(true);
2602        sparse.update_leaf(key, value_encoded(), is_private).unwrap();
2603        let sparse_root = sparse.root();
2604        let sparse_updates = sparse.take_updates();
2605
2606        assert_eq!(sparse_root, hash_builder_root);
2607        assert_eq!(sparse_updates.updated_nodes, hash_builder_updates.account_nodes);
2608        assert_eq_sparse_trie_proof_nodes(&sparse, hash_builder_proof_nodes);
2609    }
2610
2611    #[test]
2612    fn sparse_trie_empty_update_multiple_lower_nibbles() {
2613        let is_private = false; // hardcode to false for legacy test, TODO: make a private equivalent
2614        reth_tracing::init_test_tracing();
2615
2616        let paths = (0..=16).map(|b| Nibbles::unpack(B256::with_last_byte(b))).collect::<Vec<_>>();
2617        let value = || Account::default();
2618        let value_encoded = || {
2619            let mut account_rlp = Vec::new();
2620            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
2621            account_rlp
2622        };
2623
2624        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
2625            run_hash_builder(
2626                paths.iter().cloned().zip(std::iter::repeat_with(value)),
2627                NoopAccountTrieCursor::default(),
2628                Default::default(),
2629                paths.clone(),
2630            );
2631
2632        let mut sparse = RevealedSparseTrie::default().with_updates(true);
2633        for path in &paths {
2634            sparse.update_leaf(path.clone(), value_encoded(), is_private).unwrap();
2635        }
2636        let sparse_root = sparse.root();
2637        let sparse_updates = sparse.take_updates();
2638
2639        assert_eq!(sparse_root, hash_builder_root);
2640        assert_eq!(sparse_updates.updated_nodes, hash_builder_updates.account_nodes);
2641        assert_eq_sparse_trie_proof_nodes(&sparse, hash_builder_proof_nodes);
2642    }
2643
2644    #[test]
2645    fn sparse_trie_empty_update_multiple_upper_nibbles() {
2646        let is_private = false; // hardcode to false for legacy test, TODO: make a private equivalent
2647        let paths = (239..=255).map(|b| Nibbles::unpack(B256::repeat_byte(b))).collect::<Vec<_>>();
2648        let value = || Account::default();
2649        let value_encoded = || {
2650            let mut account_rlp = Vec::new();
2651            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
2652            account_rlp
2653        };
2654
2655        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
2656            run_hash_builder(
2657                paths.iter().cloned().zip(std::iter::repeat_with(value)),
2658                NoopAccountTrieCursor::default(),
2659                Default::default(),
2660                paths.clone(),
2661            );
2662
2663        let mut sparse = RevealedSparseTrie::default().with_updates(true);
2664        for path in &paths {
2665            sparse.update_leaf(path.clone(), value_encoded(), is_private).unwrap();
2666        }
2667        let sparse_root = sparse.root();
2668        let sparse_updates = sparse.take_updates();
2669
2670        assert_eq!(sparse_root, hash_builder_root);
2671        assert_eq!(sparse_updates.updated_nodes, hash_builder_updates.account_nodes);
2672        assert_eq_sparse_trie_proof_nodes(&sparse, hash_builder_proof_nodes);
2673    }
2674
2675    #[test]
2676    fn sparse_trie_empty_update_multiple() {
2677        let is_private = false; // hardcode to false for legacy test, TODO: make a private equivalent
2678        let paths = (0..=255)
2679            .map(|b| {
2680                Nibbles::unpack(if b % 2 == 0 {
2681                    B256::repeat_byte(b)
2682                } else {
2683                    B256::with_last_byte(b)
2684                })
2685            })
2686            .collect::<Vec<_>>();
2687        let value = || Account::default();
2688        let value_encoded = || {
2689            let mut account_rlp = Vec::new();
2690            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
2691            account_rlp
2692        };
2693
2694        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
2695            run_hash_builder(
2696                paths.iter().sorted_unstable().cloned().zip(std::iter::repeat_with(value)),
2697                NoopAccountTrieCursor::default(),
2698                Default::default(),
2699                paths.clone(),
2700            );
2701
2702        let mut sparse = RevealedSparseTrie::default().with_updates(true);
2703        for path in &paths {
2704            sparse.update_leaf(path.clone(), value_encoded(), is_private).unwrap();
2705        }
2706        let sparse_root = sparse.root();
2707        let sparse_updates = sparse.take_updates();
2708
2709        assert_eq!(sparse_root, hash_builder_root);
2710        pretty_assertions::assert_eq!(
2711            BTreeMap::from_iter(sparse_updates.updated_nodes),
2712            BTreeMap::from_iter(hash_builder_updates.account_nodes)
2713        );
2714        assert_eq_sparse_trie_proof_nodes(&sparse, hash_builder_proof_nodes);
2715    }
2716
2717    #[test]
2718    fn sparse_trie_empty_update_repeated() {
2719        let is_private = false; // legacy test does not use private storage
2720        let paths = (0..=255).map(|b| Nibbles::unpack(B256::repeat_byte(b))).collect::<Vec<_>>();
2721        let old_value = Account { nonce: 1, ..Default::default() };
2722        let old_value_encoded = {
2723            let mut account_rlp = Vec::new();
2724            old_value.into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
2725            account_rlp
2726        };
2727        let new_value = Account { nonce: 2, ..Default::default() };
2728        let new_value_encoded = {
2729            let mut account_rlp = Vec::new();
2730            new_value.into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
2731            account_rlp
2732        };
2733
2734        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
2735            run_hash_builder(
2736                paths.iter().cloned().zip(std::iter::repeat_with(|| old_value)),
2737                NoopAccountTrieCursor::default(),
2738                Default::default(),
2739                paths.clone(),
2740            );
2741
2742        let mut sparse = RevealedSparseTrie::default().with_updates(true);
2743        for path in &paths {
2744            sparse.update_leaf(path.clone(), old_value_encoded.clone(), is_private).unwrap();
2745        }
2746        let sparse_root = sparse.root();
2747        let sparse_updates = sparse.updates_ref();
2748
2749        assert_eq!(sparse_root, hash_builder_root);
2750        assert_eq!(sparse_updates.updated_nodes, hash_builder_updates.account_nodes);
2751        assert_eq_sparse_trie_proof_nodes(&sparse, hash_builder_proof_nodes);
2752
2753        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
2754            run_hash_builder(
2755                paths.iter().cloned().zip(std::iter::repeat_with(|| new_value)),
2756                NoopAccountTrieCursor::default(),
2757                Default::default(),
2758                paths.clone(),
2759            );
2760
2761        for path in &paths {
2762            sparse.update_leaf(path.clone(), new_value_encoded.clone(), is_private).unwrap();
2763        }
2764        let sparse_root = sparse.root();
2765        let sparse_updates = sparse.take_updates();
2766
2767        assert_eq!(sparse_root, hash_builder_root);
2768        assert_eq!(sparse_updates.updated_nodes, hash_builder_updates.account_nodes);
2769        assert_eq_sparse_trie_proof_nodes(&sparse, hash_builder_proof_nodes);
2770    }
2771
2772    #[test]
2773    fn sparse_trie_remove_leaf() {
2774        let is_private = false; // hardcoded to false all nodes except first
2775
2776        reth_tracing::init_test_tracing();
2777
2778        let mut sparse = RevealedSparseTrie::default();
2779
2780        let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec();
2781
2782        sparse
2783            .update_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]), value.clone(), true)
2784            .unwrap();
2785        sparse
2786            .update_leaf(
2787                Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
2788                value.clone(),
2789                is_private,
2790            )
2791            .unwrap();
2792        sparse
2793            .update_leaf(
2794                Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3]),
2795                value.clone(),
2796                is_private,
2797            )
2798            .unwrap();
2799        sparse
2800            .update_leaf(
2801                Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2]),
2802                value.clone(),
2803                is_private,
2804            )
2805            .unwrap();
2806        sparse
2807            .update_leaf(
2808                Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]),
2809                value.clone(),
2810                is_private,
2811            )
2812            .unwrap();
2813        sparse
2814            .update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]), value, is_private)
2815            .unwrap();
2816
2817        // Extension (Key = 5)
2818        // └── Branch (Mask = 1011)
2819        //     ├── 0 -> Extension (Key = 23)
2820        //     │        └── Branch (Mask = 0101)
2821        //     │              ├── 1 -> Leaf (Key = 1, Path = 50231)
2822        //     │              └── 3 -> Leaf (Key = 3, Path = 50233)
2823        //     ├── 2 -> Leaf (Key = 013, Path = 52013)
2824        //     └── 3 -> Branch (Mask = 0101)
2825        //                ├── 1 -> Leaf (Key = 3102, Path = 53102)
2826        //                └── 3 -> Branch (Mask = 1010)
2827        //                       ├── 0 -> Leaf (Key = 3302, Path = 53302)
2828        //                       └── 2 -> Leaf (Key = 3320, Path = 53320)
2829        pretty_assertions::assert_eq!(
2830            sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
2831            BTreeMap::from_iter([
2832                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
2833                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1101.into())),
2834                (
2835                    Nibbles::from_nibbles([0x5, 0x0]),
2836                    SparseNode::new_ext(Nibbles::from_nibbles([0x2, 0x3]))
2837                ),
2838                (
2839                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3]),
2840                    SparseNode::new_branch(0b1010.into())
2841                ),
2842                (
2843                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]),
2844                    SparseNode::new_leaf(Nibbles::default(), true)
2845                ),
2846                (
2847                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
2848                    SparseNode::new_leaf(Nibbles::default(), is_private)
2849                ),
2850                (
2851                    Nibbles::from_nibbles([0x5, 0x2]),
2852                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0, 0x1, 0x3]), is_private),
2853                ),
2854                (Nibbles::from_nibbles([0x5, 0x3]), SparseNode::new_branch(0b1010.into())),
2855                (
2856                    Nibbles::from_nibbles([0x5, 0x3, 0x1]),
2857                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0, 0x2]), is_private)
2858                ),
2859                (Nibbles::from_nibbles([0x5, 0x3, 0x3]), SparseNode::new_branch(0b0101.into())),
2860                (
2861                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0]),
2862                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2]), is_private)
2863                ),
2864                (
2865                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2]),
2866                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0]), is_private)
2867                )
2868            ])
2869        );
2870
2871        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3])).unwrap();
2872
2873        // Extension (Key = 5)
2874        // └── Branch (Mask = 1001)
2875        //     ├── 0 -> Extension (Key = 23)
2876        //     │        └── Branch (Mask = 0101)
2877        //     │              ├── 1 -> Leaf (Key = 0231, Path = 50231)
2878        //     │              └── 3 -> Leaf (Key = 0233, Path = 50233)
2879        //     └── 3 -> Branch (Mask = 0101)
2880        //                ├── 1 -> Leaf (Key = 3102, Path = 53102)
2881        //                └── 3 -> Branch (Mask = 1010)
2882        //                       ├── 0 -> Leaf (Key = 3302, Path = 53302)
2883        //                       └── 2 -> Leaf (Key = 3320, Path = 53320)
2884        pretty_assertions::assert_eq!(
2885            sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
2886            BTreeMap::from_iter([
2887                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
2888                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
2889                (
2890                    Nibbles::from_nibbles([0x5, 0x0]),
2891                    SparseNode::new_ext(Nibbles::from_nibbles([0x2, 0x3]))
2892                ),
2893                (
2894                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3]),
2895                    SparseNode::new_branch(0b1010.into())
2896                ),
2897                (
2898                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]),
2899                    SparseNode::new_leaf(Nibbles::default(), true)
2900                ),
2901                (
2902                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
2903                    SparseNode::new_leaf(Nibbles::default(), is_private)
2904                ),
2905                (Nibbles::from_nibbles([0x5, 0x3]), SparseNode::new_branch(0b1010.into())),
2906                (
2907                    Nibbles::from_nibbles([0x5, 0x3, 0x1]),
2908                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0, 0x2]), is_private)
2909                ),
2910                (Nibbles::from_nibbles([0x5, 0x3, 0x3]), SparseNode::new_branch(0b0101.into())),
2911                (
2912                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0]),
2913                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2]), is_private)
2914                ),
2915                (
2916                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2]),
2917                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0]), is_private)
2918                )
2919            ])
2920        );
2921
2922        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1])).unwrap();
2923
2924        // Extension (Key = 5)
2925        // └── Branch (Mask = 1001)
2926        //     ├── 0 -> Leaf (Key = 0233, Path = 50233)
2927        //     └── 3 -> Branch (Mask = 0101)
2928        //                ├── 1 -> Leaf (Key = 3102, Path = 53102)
2929        //                └── 3 -> Branch (Mask = 1010)
2930        //                       ├── 0 -> Leaf (Key = 3302, Path = 53302)
2931        //                       └── 2 -> Leaf (Key = 3320, Path = 53320)
2932        pretty_assertions::assert_eq!(
2933            sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
2934            BTreeMap::from_iter([
2935                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
2936                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
2937                (
2938                    Nibbles::from_nibbles([0x5, 0x0]),
2939                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2, 0x3, 0x3]), is_private)
2940                ),
2941                (Nibbles::from_nibbles([0x5, 0x3]), SparseNode::new_branch(0b1010.into())),
2942                (
2943                    Nibbles::from_nibbles([0x5, 0x3, 0x1]),
2944                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0, 0x2]), is_private)
2945                ),
2946                (Nibbles::from_nibbles([0x5, 0x3, 0x3]), SparseNode::new_branch(0b0101.into())),
2947                (
2948                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0]),
2949                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2]), is_private)
2950                ),
2951                (
2952                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2]),
2953                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0]), is_private)
2954                )
2955            ])
2956        );
2957
2958        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2])).unwrap();
2959
2960        // Extension (Key = 5)
2961        // └── Branch (Mask = 1001)
2962        //     ├── 0 -> Leaf (Key = 0233, Path = 50233)
2963        //     └── 3 -> Branch (Mask = 1010)
2964        //                ├── 0 -> Leaf (Key = 3302, Path = 53302)
2965        //                └── 2 -> Leaf (Key = 3320, Path = 53320)
2966        pretty_assertions::assert_eq!(
2967            sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
2968            BTreeMap::from_iter([
2969                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
2970                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
2971                (
2972                    Nibbles::from_nibbles([0x5, 0x0]),
2973                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2, 0x3, 0x3]), is_private)
2974                ),
2975                (
2976                    Nibbles::from_nibbles([0x5, 0x3]),
2977                    SparseNode::new_ext(Nibbles::from_nibbles([0x3]))
2978                ),
2979                (Nibbles::from_nibbles([0x5, 0x3, 0x3]), SparseNode::new_branch(0b0101.into())),
2980                (
2981                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0]),
2982                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2]), is_private)
2983                ),
2984                (
2985                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2]),
2986                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0]), is_private)
2987                )
2988            ])
2989        );
2990
2991        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0])).unwrap();
2992
2993        // Extension (Key = 5)
2994        // └── Branch (Mask = 1001)
2995        //     ├── 0 -> Leaf (Key = 0233, Path = 50233)
2996        //     └── 3 -> Leaf (Key = 3302, Path = 53302)
2997        pretty_assertions::assert_eq!(
2998            sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
2999            BTreeMap::from_iter([
3000                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
3001                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
3002                (
3003                    Nibbles::from_nibbles([0x5, 0x0]),
3004                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2, 0x3, 0x3]), is_private)
3005                ),
3006                (
3007                    Nibbles::from_nibbles([0x5, 0x3]),
3008                    SparseNode::new_leaf(Nibbles::from_nibbles([0x3, 0x0, 0x2]), is_private)
3009                ),
3010            ])
3011        );
3012
3013        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3])).unwrap();
3014
3015        // Leaf (Key = 53302)
3016        pretty_assertions::assert_eq!(
3017            sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
3018            BTreeMap::from_iter([(
3019                Nibbles::default(),
3020                SparseNode::new_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]), is_private)
3021            ),])
3022        );
3023
3024        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2])).unwrap();
3025
3026        // Empty
3027        pretty_assertions::assert_eq!(
3028            sparse.nodes.clone().into_iter().collect::<BTreeMap<_, _>>(),
3029            BTreeMap::from_iter([(Nibbles::default(), SparseNode::Empty)])
3030        );
3031    }
3032
3033    #[test]
3034    fn sparse_trie_remove_leaf_blinded() {
3035        // legacy test does not use private storage
3036        // not relevant to removing a blinded leaf
3037        let is_private = false;
3038
3039        let leaf = LeafNode::new(
3040            Nibbles::default(),
3041            alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec(),
3042            is_private,
3043        );
3044        let branch = TrieNode::Branch(BranchNode::new(
3045            vec![
3046                RlpNode::word_rlp(&B256::repeat_byte(1)),
3047                RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(),
3048            ],
3049            TrieMask::new(0b11),
3050        ));
3051
3052        let mut sparse = RevealedSparseTrie::from_root(
3053            branch.clone(),
3054            TrieMasks { hash_mask: Some(TrieMask::new(0b01)), tree_mask: None },
3055            false,
3056        )
3057        .unwrap();
3058
3059        // Reveal a branch node and one of its children
3060        //
3061        // Branch (Mask = 11)
3062        // ├── 0 -> Hash (Path = 0)
3063        // └── 1 -> Leaf (Path = 1)
3064        sparse
3065            .reveal_node(
3066                Nibbles::default(),
3067                branch,
3068                TrieMasks { hash_mask: None, tree_mask: Some(TrieMask::new(0b01)) },
3069            )
3070            .unwrap();
3071        sparse
3072            .reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), TrieMasks::none())
3073            .unwrap();
3074
3075        // Removing a blinded leaf should result in an error
3076        assert_matches!(
3077            sparse.remove_leaf(&Nibbles::from_nibbles([0x0])).map_err(|e| e.into_kind()),
3078            Err(SparseTrieErrorKind::BlindedNode { path, hash }) if path == Nibbles::from_nibbles([0x0]) && hash == B256::repeat_byte(1)
3079        );
3080    }
3081
3082    #[test]
3083    fn sparse_trie_remove_leaf_non_existent() {
3084        // legacy test does not use private storage
3085        // not relevant to removing a non-existent leaf
3086        let is_private = false;
3087
3088        let leaf = LeafNode::new(
3089            Nibbles::default(),
3090            alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec(),
3091            is_private,
3092        );
3093        let branch = TrieNode::Branch(BranchNode::new(
3094            vec![
3095                RlpNode::word_rlp(&B256::repeat_byte(1)),
3096                RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(),
3097            ],
3098            TrieMask::new(0b11),
3099        ));
3100
3101        let mut sparse = RevealedSparseTrie::from_root(
3102            branch.clone(),
3103            TrieMasks { hash_mask: Some(TrieMask::new(0b01)), tree_mask: None },
3104            false,
3105        )
3106        .unwrap();
3107
3108        // Reveal a branch node and one of its children
3109        //
3110        // Branch (Mask = 11)
3111        // ├── 0 -> Hash (Path = 0)
3112        // └── 1 -> Leaf (Path = 1)
3113        sparse
3114            .reveal_node(
3115                Nibbles::default(),
3116                branch,
3117                TrieMasks { hash_mask: None, tree_mask: Some(TrieMask::new(0b01)) },
3118            )
3119            .unwrap();
3120        sparse
3121            .reveal_node(Nibbles::from_nibbles([0x1]), TrieNode::Leaf(leaf), TrieMasks::none())
3122            .unwrap();
3123
3124        // Removing a non-existent leaf should be a noop
3125        let sparse_old = sparse.clone();
3126        assert_matches!(sparse.remove_leaf(&Nibbles::from_nibbles([0x2])), Ok(()));
3127        assert_eq!(sparse, sparse_old);
3128    }
3129
3130    #[test]
3131    fn sparse_trie_fuzz() {
3132        // Having only the first 3 nibbles set, we narrow down the range of keys
3133        // to 4096 different hashes. It allows us to generate collisions more likely
3134        // to test the sparse trie updates.
3135        const KEY_NIBBLES_LEN: usize = 3;
3136
3137        fn test(updates: Vec<(BTreeMap<Nibbles, Account>, BTreeSet<Nibbles>)>) {
3138            {
3139                let mut state = BTreeMap::default();
3140                let provider_factory = create_test_provider_factory();
3141                let mut sparse = RevealedSparseTrie::default().with_updates(true);
3142
3143                for (update, keys_to_delete) in updates {
3144                    // Insert state updates into the sparse trie and calculate the root
3145                    for (key, account) in update.clone() {
3146                        let account = account.into_trie_account(EMPTY_ROOT_HASH);
3147                        let mut account_rlp = Vec::new();
3148                        account.encode(&mut account_rlp);
3149                        let is_private = false; // hardcode to false for legacy test, TODO: make a private equivalent
3150                        sparse.update_leaf(key, account_rlp, is_private).unwrap();
3151                    }
3152                    // We need to clone the sparse trie, so that all updated branch nodes are
3153                    // preserved, and not only those that were changed after the last call to
3154                    // `root()`.
3155                    let mut updated_sparse = sparse.clone();
3156                    let sparse_root = updated_sparse.root();
3157                    let sparse_updates = updated_sparse.take_updates();
3158
3159                    // Insert state updates into the hash builder and calculate the root
3160                    state.extend(update);
3161                    let provider = provider_factory.provider().unwrap();
3162                    let trie_cursor = DatabaseTrieCursorFactory::new(provider.tx_ref());
3163                    let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
3164                        run_hash_builder(
3165                            state.clone(),
3166                            trie_cursor.account_trie_cursor().unwrap(),
3167                            Default::default(),
3168                            state.keys().cloned().collect::<Vec<_>>(),
3169                        );
3170
3171                    // Write trie updates to the database
3172                    let provider_rw = provider_factory.provider_rw().unwrap();
3173                    provider_rw.write_trie_updates(&hash_builder_updates).unwrap();
3174                    provider_rw.commit().unwrap();
3175
3176                    // Assert that the sparse trie root matches the hash builder root
3177                    assert_eq!(sparse_root, hash_builder_root);
3178                    // Assert that the sparse trie updates match the hash builder updates
3179                    pretty_assertions::assert_eq!(
3180                        BTreeMap::from_iter(sparse_updates.updated_nodes),
3181                        BTreeMap::from_iter(hash_builder_updates.account_nodes)
3182                    );
3183                    // Assert that the sparse trie nodes match the hash builder proof nodes
3184                    assert_eq_sparse_trie_proof_nodes(&updated_sparse, hash_builder_proof_nodes);
3185
3186                    // Delete some keys from both the hash builder and the sparse trie and check
3187                    // that the sparse trie root still matches the hash builder root
3188                    for key in &keys_to_delete {
3189                        state.remove(key).unwrap();
3190                        sparse.remove_leaf(key).unwrap();
3191                    }
3192
3193                    // We need to clone the sparse trie, so that all updated branch nodes are
3194                    // preserved, and not only those that were changed after the last call to
3195                    // `root()`.
3196                    let mut updated_sparse = sparse.clone();
3197                    let sparse_root = updated_sparse.root();
3198                    let sparse_updates = updated_sparse.take_updates();
3199
3200                    let provider = provider_factory.provider().unwrap();
3201                    let trie_cursor = DatabaseTrieCursorFactory::new(provider.tx_ref());
3202                    let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
3203                        run_hash_builder(
3204                            state.clone(),
3205                            trie_cursor.account_trie_cursor().unwrap(),
3206                            keys_to_delete
3207                                .iter()
3208                                .map(|nibbles| B256::from_slice(&nibbles.pack()))
3209                                .collect(),
3210                            state.keys().cloned().collect::<Vec<_>>(),
3211                        );
3212
3213                    // Write trie updates to the database
3214                    let provider_rw = provider_factory.provider_rw().unwrap();
3215                    provider_rw.write_trie_updates(&hash_builder_updates).unwrap();
3216                    provider_rw.commit().unwrap();
3217
3218                    // Assert that the sparse trie root matches the hash builder root
3219                    assert_eq!(sparse_root, hash_builder_root);
3220                    // Assert that the sparse trie updates match the hash builder updates
3221                    pretty_assertions::assert_eq!(
3222                        BTreeMap::from_iter(sparse_updates.updated_nodes),
3223                        BTreeMap::from_iter(hash_builder_updates.account_nodes)
3224                    );
3225                    // Assert that the sparse trie nodes match the hash builder proof nodes
3226                    assert_eq_sparse_trie_proof_nodes(&updated_sparse, hash_builder_proof_nodes);
3227                }
3228            }
3229        }
3230
3231        fn transform_updates(
3232            updates: Vec<BTreeMap<Nibbles, Account>>,
3233            mut rng: impl rand_08::Rng,
3234        ) -> Vec<(BTreeMap<Nibbles, Account>, BTreeSet<Nibbles>)> {
3235            let mut keys = BTreeSet::new();
3236            updates
3237                .into_iter()
3238                .map(|update| {
3239                    keys.extend(update.keys().cloned());
3240
3241                    let keys_to_delete_len = update.len() / 2;
3242                    let keys_to_delete = (0..keys_to_delete_len)
3243                        .map(|_| {
3244                            let key = rand_08::seq::IteratorRandom::choose(keys.iter(), &mut rng)
3245                                .unwrap()
3246                                .clone();
3247                            keys.take(&key).unwrap()
3248                        })
3249                        .collect();
3250
3251                    (update, keys_to_delete)
3252                })
3253                .collect::<Vec<_>>()
3254        }
3255
3256        proptest!(ProptestConfig::with_cases(10), |(
3257            updates in proptest::collection::vec(
3258                proptest::collection::btree_map(
3259                    any_with::<Nibbles>(SizeRange::new(KEY_NIBBLES_LEN..=KEY_NIBBLES_LEN)).prop_map(pad_nibbles_right),
3260                    arb::<Account>(),
3261                    1..50,
3262                ),
3263                1..50,
3264            ).prop_perturb(transform_updates)
3265        )| {
3266            test(updates)
3267        });
3268    }
3269
3270    /// We have three leaves that share the same prefix: 0x00, 0x01 and 0x02. Hash builder trie has
3271    /// only nodes 0x00 and 0x01, and we have proofs for them. Node B is new and inserted in the
3272    /// sparse trie first.
3273    ///
3274    /// 1. Reveal the hash builder proof to leaf 0x00 in the sparse trie.
3275    /// 2. Insert leaf 0x01 into the sparse trie.
3276    /// 3. Reveal the hash builder proof to leaf 0x02 in the sparse trie.
3277    ///
3278    /// The hash builder proof to the leaf 0x02 didn't have the leaf 0x01 at the corresponding
3279    /// nibble of the branch node, so we need to adjust the branch node instead of fully
3280    /// replacing it.
3281    #[test]
3282    fn sparse_trie_reveal_node_1() {
3283        let is_private = false; // hardcode to false for legacy test, TODO: make a private equivalent
3284
3285        let key1 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x00]));
3286        let key2 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x01]));
3287        let key3 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x02]));
3288        let value = || Account::default();
3289        let value_encoded = || {
3290            let mut account_rlp = Vec::new();
3291            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
3292            account_rlp
3293        };
3294
3295        // Generate the proof for the root node and initialize the sparse trie with it
3296        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
3297            run_hash_builder(
3298                [(key1(), value()), (key3(), value())],
3299                NoopAccountTrieCursor::default(),
3300                Default::default(),
3301                [Nibbles::default()],
3302            );
3303        let mut sparse = RevealedSparseTrie::from_root(
3304            TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(),
3305            TrieMasks {
3306                hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(),
3307                tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(),
3308            },
3309            false,
3310        )
3311        .unwrap();
3312
3313        // Generate the proof for the first key and reveal it in the sparse trie
3314        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
3315            run_hash_builder(
3316                [(key1(), value()), (key3(), value())],
3317                NoopAccountTrieCursor::default(),
3318                Default::default(),
3319                [key1()],
3320            );
3321        for (path, node) in hash_builder_proof_nodes.nodes_sorted() {
3322            let hash_mask = branch_node_hash_masks.get(&path).copied();
3323            let tree_mask = branch_node_tree_masks.get(&path).copied();
3324            sparse
3325                .reveal_node(
3326                    path,
3327                    TrieNode::decode(&mut &node[..]).unwrap(),
3328                    TrieMasks { hash_mask, tree_mask },
3329                )
3330                .unwrap();
3331        }
3332
3333        // Check that the branch node exists with only two nibbles set
3334        assert_eq!(
3335            sparse.nodes.get(&Nibbles::default()),
3336            Some(&SparseNode::new_branch(0b101.into()))
3337        );
3338
3339        // Insert the leaf for the second key
3340        sparse.update_leaf(key2(), value_encoded(), is_private).unwrap();
3341
3342        // Check that the branch node was updated and another nibble was set
3343        assert_eq!(
3344            sparse.nodes.get(&Nibbles::default()),
3345            Some(&SparseNode::new_branch(0b111.into()))
3346        );
3347
3348        // Generate the proof for the third key and reveal it in the sparse trie
3349        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
3350            run_hash_builder(
3351                [(key1(), value()), (key3(), value())],
3352                NoopAccountTrieCursor::default(),
3353                Default::default(),
3354                [key3()],
3355            );
3356        for (path, node) in hash_builder_proof_nodes.nodes_sorted() {
3357            let hash_mask = branch_node_hash_masks.get(&path).copied();
3358            let tree_mask = branch_node_tree_masks.get(&path).copied();
3359            sparse
3360                .reveal_node(
3361                    path,
3362                    TrieNode::decode(&mut &node[..]).unwrap(),
3363                    TrieMasks { hash_mask, tree_mask },
3364                )
3365                .unwrap();
3366        }
3367
3368        // Check that nothing changed in the branch node
3369        assert_eq!(
3370            sparse.nodes.get(&Nibbles::default()),
3371            Some(&SparseNode::new_branch(0b111.into()))
3372        );
3373
3374        // Generate the nodes for the full trie with all three key using the hash builder, and
3375        // compare them to the sparse trie
3376        let (_, _, hash_builder_proof_nodes, _, _) = run_hash_builder(
3377            [(key1(), value()), (key2(), value()), (key3(), value())],
3378            NoopAccountTrieCursor::default(),
3379            Default::default(),
3380            [key1(), key2(), key3()],
3381        );
3382
3383        assert_eq_sparse_trie_proof_nodes(&sparse, hash_builder_proof_nodes);
3384    }
3385
3386    /// We have three leaves: 0x0000, 0x0101, and 0x0102. Hash builder trie has all nodes, and we
3387    /// have proofs for them.
3388    ///
3389    /// 1. Reveal the hash builder proof to leaf 0x00 in the sparse trie.
3390    /// 2. Remove leaf 0x00 from the sparse trie (that will remove the branch node and create an
3391    ///    extension node with the key 0x0000).
3392    /// 3. Reveal the hash builder proof to leaf 0x0101 in the sparse trie.
3393    ///
3394    /// The hash builder proof to the leaf 0x0101 had a branch node in the path, but we turned it
3395    /// into an extension node, so it should ignore this node.
3396    #[test]
3397    fn sparse_trie_reveal_node_2() {
3398        let key1 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x00, 0x00]));
3399        let key2 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x01, 0x01]));
3400        let key3 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x01, 0x02]));
3401        let value = || Account::default();
3402
3403        // Generate the proof for the root node and initialize the sparse trie with it
3404        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
3405            run_hash_builder(
3406                [(key1(), value()), (key2(), value()), (key3(), value())],
3407                NoopAccountTrieCursor::default(),
3408                Default::default(),
3409                [Nibbles::default()],
3410            );
3411        let mut sparse = RevealedSparseTrie::from_root(
3412            TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(),
3413            TrieMasks {
3414                hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(),
3415                tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(),
3416            },
3417            false,
3418        )
3419        .unwrap();
3420
3421        // Generate the proof for the children of the root branch node and reveal it in the sparse
3422        // trie
3423        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
3424            run_hash_builder(
3425                [(key1(), value()), (key2(), value()), (key3(), value())],
3426                NoopAccountTrieCursor::default(),
3427                Default::default(),
3428                [key1(), Nibbles::from_nibbles_unchecked([0x01])],
3429            );
3430        for (path, node) in hash_builder_proof_nodes.nodes_sorted() {
3431            let hash_mask = branch_node_hash_masks.get(&path).copied();
3432            let tree_mask = branch_node_tree_masks.get(&path).copied();
3433            sparse
3434                .reveal_node(
3435                    path,
3436                    TrieNode::decode(&mut &node[..]).unwrap(),
3437                    TrieMasks { hash_mask, tree_mask },
3438                )
3439                .unwrap();
3440        }
3441
3442        // Check that the branch node exists
3443        assert_eq!(
3444            sparse.nodes.get(&Nibbles::default()),
3445            Some(&SparseNode::new_branch(0b11.into()))
3446        );
3447
3448        // Remove the leaf for the first key
3449        sparse.remove_leaf(&key1()).unwrap();
3450
3451        // Check that the branch node was turned into an extension node
3452        assert_eq!(
3453            sparse.nodes.get(&Nibbles::default()),
3454            Some(&SparseNode::new_ext(Nibbles::from_nibbles_unchecked([0x01])))
3455        );
3456
3457        // Generate the proof for the third key and reveal it in the sparse trie
3458        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
3459            run_hash_builder(
3460                [(key1(), value()), (key2(), value()), (key3(), value())],
3461                NoopAccountTrieCursor::default(),
3462                Default::default(),
3463                [key2()],
3464            );
3465        for (path, node) in hash_builder_proof_nodes.nodes_sorted() {
3466            let hash_mask = branch_node_hash_masks.get(&path).copied();
3467            let tree_mask = branch_node_tree_masks.get(&path).copied();
3468            sparse
3469                .reveal_node(
3470                    path,
3471                    TrieNode::decode(&mut &node[..]).unwrap(),
3472                    TrieMasks { hash_mask, tree_mask },
3473                )
3474                .unwrap();
3475        }
3476
3477        // Check that nothing changed in the extension node
3478        assert_eq!(
3479            sparse.nodes.get(&Nibbles::default()),
3480            Some(&SparseNode::new_ext(Nibbles::from_nibbles_unchecked([0x01])))
3481        );
3482    }
3483
3484    /// We have two leaves that share the same prefix: 0x0001 and 0x0002, and a leaf with a
3485    /// different prefix: 0x0100. Hash builder trie has only the first two leaves, and we have
3486    /// proofs for them.
3487    ///
3488    /// 1. Insert the leaf 0x0100 into the sparse trie, and check that the root extension node was
3489    ///    turned into a branch node.
3490    /// 2. Reveal the leaf 0x0001 in the sparse trie, and check that the root branch node wasn't
3491    ///    overwritten with the extension node from the proof.
3492    #[test]
3493    fn sparse_trie_reveal_node_3() {
3494        let is_private = false; // hardcode to false for legacy test, TODO: make a private equivalent
3495        let key1 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x00, 0x01]));
3496        let key2 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x00, 0x02]));
3497        let key3 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x01, 0x00]));
3498        let value = || Account::default();
3499        let value_encoded = || {
3500            let mut account_rlp = Vec::new();
3501            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
3502            account_rlp
3503        };
3504
3505        // Generate the proof for the root node and initialize the sparse trie with it
3506        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
3507            run_hash_builder(
3508                [(key1(), value()), (key2(), value())],
3509                NoopAccountTrieCursor::default(),
3510                Default::default(),
3511                [Nibbles::default()],
3512            );
3513        let mut sparse = RevealedSparseTrie::from_root(
3514            TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(),
3515            TrieMasks {
3516                hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(),
3517                tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(),
3518            },
3519            false,
3520        )
3521        .unwrap();
3522
3523        // Check that the root extension node exists
3524        assert_matches!(
3525            sparse.nodes.get(&Nibbles::default()),
3526            Some(SparseNode::Extension { key, hash: None, store_in_db_trie: None }) if *key == Nibbles::from_nibbles([0x00])
3527        );
3528
3529        // Insert the leaf with a different prefix
3530        sparse.update_leaf(key3(), value_encoded(), is_private).unwrap();
3531
3532        // Check that the extension node was turned into a branch node
3533        assert_matches!(
3534            sparse.nodes.get(&Nibbles::default()),
3535            Some(SparseNode::Branch { state_mask, hash: None, store_in_db_trie: None }) if *state_mask == TrieMask::new(0b11)
3536        );
3537
3538        // Generate the proof for the first key and reveal it in the sparse trie
3539        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
3540            run_hash_builder(
3541                [(key1(), value()), (key2(), value())],
3542                NoopAccountTrieCursor::default(),
3543                Default::default(),
3544                [key1()],
3545            );
3546        for (path, node) in hash_builder_proof_nodes.nodes_sorted() {
3547            let hash_mask = branch_node_hash_masks.get(&path).copied();
3548            let tree_mask = branch_node_tree_masks.get(&path).copied();
3549            sparse
3550                .reveal_node(
3551                    path,
3552                    TrieNode::decode(&mut &node[..]).unwrap(),
3553                    TrieMasks { hash_mask, tree_mask },
3554                )
3555                .unwrap();
3556        }
3557
3558        // Check that the branch node wasn't overwritten by the extension node in the proof
3559        assert_matches!(
3560            sparse.nodes.get(&Nibbles::default()),
3561            Some(SparseNode::Branch { state_mask, hash: None, store_in_db_trie: None }) if *state_mask == TrieMask::new(0b11)
3562        );
3563    }
3564
3565    #[test]
3566    fn sparse_trie_get_changed_nodes_at_depth() {
3567        let is_private = false; // hardcoded to false all nodes except first
3568        let mut sparse = RevealedSparseTrie::default();
3569
3570        let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec();
3571
3572        // Extension (Key = 5) – Level 0
3573        // └── Branch (Mask = 1011) – Level 1
3574        //     ├── 0 -> Extension (Key = 23) – Level 2
3575        //     │        └── Branch (Mask = 0101) – Level 3
3576        //     │              ├── 1 -> Leaf (Key = 1, Path = 50231) – Level 4
3577        //     │              └── 3 -> Leaf (Key = 3, Path = 50233) – Level 4
3578        //     ├── 2 -> Leaf (Key = 013, Path = 52013) – Level 2
3579        //     └── 3 -> Branch (Mask = 0101) – Level 2
3580        //                ├── 1 -> Leaf (Key = 3102, Path = 53102) – Level 3
3581        //                └── 3 -> Branch (Mask = 1010) – Level 3
3582        //                       ├── 0 -> Leaf (Key = 3302, Path = 53302) – Level 4
3583        //                       └── 2 -> Leaf (Key = 3320, Path = 53320) – Level 4
3584        sparse
3585            .update_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]), value.clone(), true)
3586            .unwrap();
3587        sparse
3588            .update_leaf(
3589                Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
3590                value.clone(),
3591                is_private,
3592            )
3593            .unwrap();
3594        sparse
3595            .update_leaf(
3596                Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3]),
3597                value.clone(),
3598                is_private,
3599            )
3600            .unwrap();
3601        sparse
3602            .update_leaf(
3603                Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2]),
3604                value.clone(),
3605                is_private,
3606            )
3607            .unwrap();
3608        sparse
3609            .update_leaf(
3610                Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]),
3611                value.clone(),
3612                is_private,
3613            )
3614            .unwrap();
3615        sparse
3616            .update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]), value, is_private)
3617            .unwrap();
3618
3619        assert_eq!(
3620            sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 0),
3621            (vec![(0, Nibbles::default())], PrefixSetMut::default())
3622        );
3623        assert_eq!(
3624            sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 1),
3625            (vec![(1, Nibbles::from_nibbles_unchecked([0x5]))], [Nibbles::default()].into())
3626        );
3627        assert_eq!(
3628            sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 2),
3629            (
3630                vec![
3631                    (2, Nibbles::from_nibbles_unchecked([0x5, 0x0])),
3632                    (2, Nibbles::from_nibbles_unchecked([0x5, 0x2])),
3633                    (2, Nibbles::from_nibbles_unchecked([0x5, 0x3]))
3634                ],
3635                [Nibbles::default(), Nibbles::from_nibbles_unchecked([0x5])].into()
3636            )
3637        );
3638        assert_eq!(
3639            sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 3),
3640            (
3641                vec![
3642                    (3, Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3])),
3643                    (2, Nibbles::from_nibbles_unchecked([0x5, 0x2])),
3644                    (3, Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x1])),
3645                    (3, Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3]))
3646                ],
3647                [
3648                    Nibbles::default(),
3649                    Nibbles::from_nibbles_unchecked([0x5]),
3650                    Nibbles::from_nibbles_unchecked([0x5, 0x0]),
3651                    Nibbles::from_nibbles_unchecked([0x5, 0x3])
3652                ]
3653                .into()
3654            )
3655        );
3656        assert_eq!(
3657            sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 4),
3658            (
3659                vec![
3660                    (4, Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x1])),
3661                    (4, Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x3])),
3662                    (2, Nibbles::from_nibbles_unchecked([0x5, 0x2])),
3663                    (3, Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x1])),
3664                    (4, Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3, 0x0])),
3665                    (4, Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3, 0x2]))
3666                ],
3667                [
3668                    Nibbles::default(),
3669                    Nibbles::from_nibbles_unchecked([0x5]),
3670                    Nibbles::from_nibbles_unchecked([0x5, 0x0]),
3671                    Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3]),
3672                    Nibbles::from_nibbles_unchecked([0x5, 0x3]),
3673                    Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3])
3674                ]
3675                .into()
3676            )
3677        );
3678    }
3679
3680    #[test]
3681    fn hash_builder_branch_hash_mask() {
3682        let is_private = false; // legacy test does not use private storage
3683        let key1 = || pad_nibbles_left(Nibbles::from_nibbles_unchecked([0x00]));
3684        let key2 = || pad_nibbles_left(Nibbles::from_nibbles_unchecked([0x01]));
3685        let value = || Account { bytecode_hash: Some(B256::repeat_byte(1)), ..Default::default() };
3686        let value_encoded = || {
3687            let mut account_rlp = Vec::new();
3688            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
3689            account_rlp
3690        };
3691
3692        let (hash_builder_root, hash_builder_updates, _, _, _) = run_hash_builder(
3693            [(key1(), value()), (key2(), value())],
3694            NoopAccountTrieCursor::default(),
3695            Default::default(),
3696            [Nibbles::default()],
3697        );
3698        let mut sparse = RevealedSparseTrie::default();
3699        sparse.update_leaf(key1(), value_encoded(), is_private).unwrap();
3700        sparse.update_leaf(key2(), value_encoded(), is_private).unwrap();
3701        let sparse_root = sparse.root();
3702        let sparse_updates = sparse.take_updates();
3703
3704        assert_eq!(sparse_root, hash_builder_root);
3705        assert_eq!(sparse_updates.updated_nodes, hash_builder_updates.account_nodes);
3706    }
3707
3708    #[test]
3709    fn sparse_trie_wipe() {
3710        let is_private = false; // hardcoded to false all nodes except first
3711        let mut sparse = RevealedSparseTrie::default().with_updates(true);
3712
3713        let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec();
3714
3715        // Extension (Key = 5) – Level 0
3716        // └── Branch (Mask = 1011) – Level 1
3717        //     ├── 0 -> Extension (Key = 23) – Level 2
3718        //     │        └── Branch (Mask = 0101) – Level 3
3719        //     │              ├── 1 -> Leaf (Key = 1, Path = 50231) – Level 4
3720        //     │              └── 3 -> Leaf (Key = 3, Path = 50233) – Level 4
3721        //     ├── 2 -> Leaf (Key = 013, Path = 52013) – Level 2
3722        //     └── 3 -> Branch (Mask = 0101) – Level 2
3723        //                ├── 1 -> Leaf (Key = 3102, Path = 53102) – Level 3
3724        //                └── 3 -> Branch (Mask = 1010) – Level 3
3725        //                       ├── 0 -> Leaf (Key = 3302, Path = 53302) – Level 4
3726        //                       └── 2 -> Leaf (Key = 3320, Path = 53320) – Level 4
3727        sparse
3728            .update_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]), value.clone(), true)
3729            .unwrap();
3730        sparse
3731            .update_leaf(
3732                Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
3733                value.clone(),
3734                is_private,
3735            )
3736            .unwrap();
3737        sparse
3738            .update_leaf(
3739                Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3]),
3740                value.clone(),
3741                is_private,
3742            )
3743            .unwrap();
3744        sparse
3745            .update_leaf(
3746                Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2]),
3747                value.clone(),
3748                is_private,
3749            )
3750            .unwrap();
3751        sparse
3752            .update_leaf(
3753                Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]),
3754                value.clone(),
3755                is_private,
3756            )
3757            .unwrap();
3758        sparse
3759            .update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]), value, is_private)
3760            .unwrap();
3761
3762        sparse.wipe();
3763
3764        assert_eq!(sparse.root(), EMPTY_ROOT_HASH);
3765    }
3766
3767    #[test]
3768    fn sparse_trie_clear() {
3769        let is_private = false; // hardcoded to false for legacy test
3770                                // tests that if we fill a sparse trie with some nodes and then clear it, it has the same
3771                                // contents as an empty sparse trie
3772        let mut sparse = RevealedSparseTrie::default();
3773        let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec();
3774        sparse
3775            .update_leaf(
3776                Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]),
3777                value.clone(),
3778                is_private,
3779            )
3780            .unwrap();
3781        sparse
3782            .update_leaf(
3783                Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
3784                value.clone(),
3785                is_private,
3786            )
3787            .unwrap();
3788        sparse
3789            .update_leaf(
3790                Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3]),
3791                value.clone(),
3792                is_private,
3793            )
3794            .unwrap();
3795        sparse
3796            .update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2]), value, is_private)
3797            .unwrap();
3798
3799        sparse.clear();
3800
3801        // we have to update the root hash to be an empty one, because the `Default` impl of
3802        // `RevealedSparseTrie` sets the root hash to `EMPTY_ROOT_HASH` in the constructor.
3803        //
3804        // The default impl is only used in tests.
3805        sparse.nodes.insert(Nibbles::default(), SparseNode::Empty);
3806
3807        let empty_trie = RevealedSparseTrie::default();
3808        assert_eq!(empty_trie, sparse);
3809    }
3810
3811    #[test]
3812    fn sparse_trie_display() {
3813        let is_private = false; // hardcoded to false all nodes except first
3814        let mut sparse = RevealedSparseTrie::default();
3815
3816        let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec();
3817
3818        // Extension (Key = 5) – Level 0
3819        // └── Branch (Mask = 1011) – Level 1
3820        //     ├── 0 -> Extension (Key = 23) – Level 2
3821        //     │        └── Branch (Mask = 0101) – Level 3
3822        //     │              ├── 1 -> Leaf (Key = 1, Path = 50231) – Level 4
3823        //     │              └── 3 -> Leaf (Key = 3, Path = 50233) – Level 4
3824        //     ├── 2 -> Leaf (Key = 013, Path = 52013) – Level 2
3825        //     └── 3 -> Branch (Mask = 0101) – Level 2
3826        //                ├── 1 -> Leaf (Key = 3102, Path = 53102) – Level 3
3827        //                └── 3 -> Branch (Mask = 1010) – Level 3
3828        //                       ├── 0 -> Leaf (Key = 3302, Path = 53302) – Level 4
3829        //                       └── 2 -> Leaf (Key = 3320, Path = 53320) – Level 4
3830        sparse
3831            .update_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]), value.clone(), true)
3832            .unwrap();
3833        sparse
3834            .update_leaf(
3835                Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
3836                value.clone(),
3837                is_private,
3838            )
3839            .unwrap();
3840        sparse
3841            .update_leaf(
3842                Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3]),
3843                value.clone(),
3844                is_private,
3845            )
3846            .unwrap();
3847        sparse
3848            .update_leaf(
3849                Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2]),
3850                value.clone(),
3851                is_private,
3852            )
3853            .unwrap();
3854        sparse
3855            .update_leaf(
3856                Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]),
3857                value.clone(),
3858                is_private,
3859            )
3860            .unwrap();
3861        sparse
3862            .update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]), value, is_private)
3863            .unwrap();
3864
3865        let normal_printed = format!("{sparse}");
3866        let expected = "\
3867Root -> Extension { key: Nibbles(0x05), hash: None, store_in_db_trie: None }
38685 -> Branch { state_mask: TrieMask(0000000000001101), hash: None, store_in_db_trie: None }
386950 -> Extension { key: Nibbles(0x0203), hash: None, store_in_db_trie: None }
38705023 -> Branch { state_mask: TrieMask(0000000000001010), hash: None, store_in_db_trie: None }
387150231 -> Leaf { key: Nibbles(0x), hash: None, is_private: true }
387250233 -> Leaf { key: Nibbles(0x), hash: None, is_private: false }
387352013 -> Leaf { key: Nibbles(0x000103), hash: None, is_private: false }
387453 -> Branch { state_mask: TrieMask(0000000000001010), hash: None, store_in_db_trie: None }
387553102 -> Leaf { key: Nibbles(0x0002), hash: None, is_private: false }
3876533 -> Branch { state_mask: TrieMask(0000000000000101), hash: None, store_in_db_trie: None }
387753302 -> Leaf { key: Nibbles(0x02), hash: None, is_private: false }
387853320 -> Leaf { key: Nibbles(0x00), hash: None, is_private: false }
3879";
3880        assert_eq!(normal_printed, expected);
3881
3882        let alternate_printed = format!("{sparse:#}");
3883        let expected = "\
3884Root -> Extension { key: Nibbles(0x05), hash: None, store_in_db_trie: None }
3885    5 -> Branch { state_mask: TrieMask(0000000000001101), hash: None, store_in_db_trie: None }
3886        50 -> Extension { key: Nibbles(0x0203), hash: None, store_in_db_trie: None }
3887            5023 -> Branch { state_mask: TrieMask(0000000000001010), hash: None, store_in_db_trie: None }
3888                50231 -> Leaf { key: Nibbles(0x), hash: None, is_private: true }
3889                50233 -> Leaf { key: Nibbles(0x), hash: None, is_private: false }
3890        52013 -> Leaf { key: Nibbles(0x000103), hash: None, is_private: false }
3891        53 -> Branch { state_mask: TrieMask(0000000000001010), hash: None, store_in_db_trie: None }
3892            53102 -> Leaf { key: Nibbles(0x0002), hash: None, is_private: false }
3893            533 -> Branch { state_mask: TrieMask(0000000000000101), hash: None, store_in_db_trie: None }
3894                53302 -> Leaf { key: Nibbles(0x02), hash: None, is_private: false }
3895                53320 -> Leaf { key: Nibbles(0x00), hash: None, is_private: false }
3896";
3897
3898        assert_eq!(alternate_printed, expected);
3899    }
3900}