reth_trie_sparse_parallel/
trie.rs

1use crate::LowerSparseSubtrie;
2use alloc::borrow::Cow;
3use alloy_primitives::{
4    map::{Entry, HashMap},
5    B256,
6};
7use alloy_rlp::Decodable;
8use alloy_trie::{BranchNodeCompact, TrieMask, EMPTY_ROOT_HASH};
9use reth_execution_errors::{SparseTrieErrorKind, SparseTrieResult};
10use reth_trie_common::{
11    prefix_set::{PrefixSet, PrefixSetMut},
12    BranchNodeRef, ExtensionNodeRef, LeafNodeRef, Nibbles, RlpNode, TrieNode, CHILD_INDEX_RANGE,
13};
14use reth_trie_sparse::{
15    provider::{RevealedNode, TrieNodeProvider},
16    LeafLookup, LeafLookupError, RevealedSparseNode, RlpNodeStackItem, SparseNode, SparseNodeType,
17    SparseTrieInterface, SparseTrieUpdates, TrieMasks,
18};
19use smallvec::SmallVec;
20use std::{
21    cmp::{Ord, Ordering, PartialOrd},
22    sync::mpsc,
23};
24use tracing::{instrument, trace, warn};
25
26/// The maximum length of a path, in nibbles, which belongs to the upper subtrie of a
27/// [`ParallelSparseTrie`]. All longer paths belong to a lower subtrie.
28pub const UPPER_TRIE_MAX_DEPTH: usize = 2;
29
30/// Number of lower subtries which are managed by the [`ParallelSparseTrie`].
31pub const NUM_LOWER_SUBTRIES: usize = 16usize.pow(UPPER_TRIE_MAX_DEPTH as u32);
32
33/// A revealed sparse trie with subtries that can be updated in parallel.
34///
35/// ## Structure
36///
37/// The trie is divided into two tiers for efficient parallel processing:
38/// - **Upper subtrie**: Contains nodes with paths shorter than [`UPPER_TRIE_MAX_DEPTH`]
39/// - **Lower subtries**: An array of [`NUM_LOWER_SUBTRIES`] subtries, each handling nodes with
40///   paths of at least [`UPPER_TRIE_MAX_DEPTH`] nibbles
41///
42/// Node placement is determined by path depth:
43/// - Paths with < [`UPPER_TRIE_MAX_DEPTH`] nibbles go to the upper subtrie
44/// - Paths with >= [`UPPER_TRIE_MAX_DEPTH`] nibbles go to lower subtries, indexed by their first
45///   [`UPPER_TRIE_MAX_DEPTH`] nibbles.
46///
47/// Each lower subtrie tracks its root via the `path` field, which represents the shortest path
48/// in that subtrie. This path will have at least [`UPPER_TRIE_MAX_DEPTH`] nibbles, but may be
49/// longer when an extension node in the upper trie "reaches into" the lower subtrie. For example,
50/// if the upper trie has an extension from `0x1` to `0x12345`, then the lower subtrie for prefix
51/// `0x12` will have its root at path `0x12345` rather than at `0x12`.
52///
53/// ## Node Revealing
54///
55/// The trie uses lazy loading to efficiently handle large state tries. Nodes can be:
56/// - **Blind nodes**: Stored as hashes ([`SparseNode::Hash`]), representing unloaded trie parts
57/// - **Revealed nodes**: Fully loaded nodes (Branch, Extension, Leaf) with complete structure
58///
59/// Note: An empty trie contains an `EmptyRoot` node at the root path, rather than no nodes at all.
60/// A trie with no nodes is blinded, its root may be `EmptyRoot` or some other node type.
61///
62/// Revealing is generally done using pre-loaded node data provided to via `reveal_nodes`. In
63/// certain cases, such as edge-cases when updating/removing leaves, nodes are revealed on-demand.
64///
65/// ## Leaf Operations
66///
67/// **Update**: When updating a leaf, the new value is stored in the appropriate subtrie's values
68/// map. If the leaf is new, the trie structure is updated by walking to the leaf from the root,
69/// creating necessary intermediate branch nodes.
70///
71/// **Removal**: Leaf removal may require parent node modifications. The algorithm walks up the
72/// trie, removing nodes that become empty and converting single-child branches to extensions.
73///
74/// During leaf operations the overall structure of the trie may change, causing nodes to be moved
75/// from the upper to lower trie or vice-versa.
76///
77/// The `prefix_set` is modified during both leaf updates and removals to track changed leaf paths.
78///
79/// ## Root Hash Calculation
80///
81/// Root hash computation follows a bottom-up approach:
82/// 1. Update hashes for all modified lower subtries (can be done in parallel)
83/// 2. Update hashes for the upper subtrie (which may reference lower subtrie hashes)
84/// 3. Calculate the final root hash from the upper subtrie's root node
85///
86/// The `prefix_set` tracks which paths have been modified, enabling incremental updates instead of
87/// recalculating the entire trie.
88///
89/// ## Invariants
90///
91/// - Each leaf entry in the `subtries` and `upper_trie` collection must have a corresponding entry
92///   in `values` collection. If the root node is a leaf, it must also have an entry in `values`.
93/// - All keys in `values` collection are full leaf paths.
94#[derive(Clone, PartialEq, Eq, Debug)]
95pub struct ParallelSparseTrie {
96    /// This contains the trie nodes for the upper part of the trie.
97    upper_subtrie: Box<SparseSubtrie>,
98    /// An array containing the subtries at the second level of the trie.
99    lower_subtries: [LowerSparseSubtrie; NUM_LOWER_SUBTRIES],
100    /// Set of prefixes (key paths) that have been marked as updated.
101    /// This is used to track which parts of the trie need to be recalculated.
102    prefix_set: PrefixSetMut,
103    /// Optional tracking of trie updates for later use.
104    updates: Option<SparseTrieUpdates>,
105    /// When a bit is set, the corresponding child subtree is stored in the database.
106    branch_node_tree_masks: HashMap<Nibbles, TrieMask>,
107    /// When a bit is set, the corresponding child is stored as a hash in the database.
108    branch_node_hash_masks: HashMap<Nibbles, TrieMask>,
109    /// Reusable buffer pool used for collecting [`SparseTrieUpdatesAction`]s during hash
110    /// computations.
111    update_actions_buffers: Vec<Vec<SparseTrieUpdatesAction>>,
112    /// Metrics for the parallel sparse trie.
113    #[cfg(feature = "metrics")]
114    metrics: crate::metrics::ParallelSparseTrieMetrics,
115}
116
117impl Default for ParallelSparseTrie {
118    fn default() -> Self {
119        Self {
120            upper_subtrie: Box::new(SparseSubtrie {
121                nodes: HashMap::from_iter([(Nibbles::default(), SparseNode::Empty)]),
122                ..Default::default()
123            }),
124            lower_subtries: [const { LowerSparseSubtrie::Blind(None) }; NUM_LOWER_SUBTRIES],
125            prefix_set: PrefixSetMut::default(),
126            updates: None,
127            branch_node_tree_masks: HashMap::default(),
128            branch_node_hash_masks: HashMap::default(),
129            update_actions_buffers: Vec::default(),
130            #[cfg(feature = "metrics")]
131            metrics: Default::default(),
132        }
133    }
134}
135
136impl SparseTrieInterface for ParallelSparseTrie {
137    fn with_root(
138        mut self,
139        root: TrieNode,
140        masks: TrieMasks,
141        retain_updates: bool,
142    ) -> SparseTrieResult<Self> {
143        // A fresh/cleared `ParallelSparseTrie` has a `SparseNode::Empty` at its root in the upper
144        // subtrie. Delete that so we can reveal the new root node.
145        let path = Nibbles::default();
146        let _removed_root = self.upper_subtrie.nodes.remove(&path).expect("root node should exist");
147        debug_assert_eq!(_removed_root, SparseNode::Empty);
148
149        self = self.with_updates(retain_updates);
150
151        self.reveal_upper_node(Nibbles::default(), &root, masks)?;
152        Ok(self)
153    }
154
155    fn with_updates(mut self, retain_updates: bool) -> Self {
156        self.updates = retain_updates.then(Default::default);
157        self
158    }
159
160    fn reveal_nodes(&mut self, mut nodes: Vec<RevealedSparseNode>) -> SparseTrieResult<()> {
161        if nodes.is_empty() {
162            return Ok(())
163        }
164
165        // Sort nodes first by their subtrie, and secondarily by their path. This allows for
166        // grouping nodes by their subtrie using `chunk_by`.
167        nodes.sort_unstable_by(
168            |RevealedSparseNode { path: path_a, .. }, RevealedSparseNode { path: path_b, .. }| {
169                let subtrie_type_a = SparseSubtrieType::from_path(path_a);
170                let subtrie_type_b = SparseSubtrieType::from_path(path_b);
171                subtrie_type_a.cmp(&subtrie_type_b).then(path_a.cmp(path_b))
172            },
173        );
174
175        // Update the top-level branch node masks. This is simple and can't be done in parallel.
176        for RevealedSparseNode { path, masks, .. } in &nodes {
177            if let Some(tree_mask) = masks.tree_mask {
178                self.branch_node_tree_masks.insert(*path, tree_mask);
179            }
180            if let Some(hash_mask) = masks.hash_mask {
181                self.branch_node_hash_masks.insert(*path, hash_mask);
182            }
183        }
184
185        // Due to the sorting all upper subtrie nodes will be at the front of the slice. We split
186        // them off from the rest to be handled specially by
187        // `ParallelSparseTrie::reveal_upper_node`.
188        let num_upper_nodes = nodes
189            .iter()
190            .position(|n| !SparseSubtrieType::path_len_is_upper(n.path.len()))
191            .unwrap_or(nodes.len());
192
193        let upper_nodes = &nodes[..num_upper_nodes];
194        let lower_nodes = &nodes[num_upper_nodes..];
195
196        // Reserve the capacity of the upper subtrie's `nodes` HashMap before iterating, so we don't
197        // end up making many small capacity changes as we loop.
198        self.upper_subtrie.nodes.reserve(upper_nodes.len());
199        for node in upper_nodes {
200            self.reveal_upper_node(node.path, &node.node, node.masks)?;
201        }
202
203        #[cfg(not(feature = "std"))]
204        // Reveal lower subtrie nodes serially if nostd
205        {
206            for node in lower_nodes {
207                if let Some(subtrie) = self.lower_subtrie_for_path_mut(&node.path) {
208                    subtrie.reveal_node(node.path, &node.node, &node.masks)?;
209                } else {
210                    panic!("upper subtrie node {node:?} found amongst lower nodes");
211                }
212            }
213            Ok(())
214        }
215
216        #[cfg(feature = "std")]
217        // Reveal lower subtrie nodes in parallel
218        {
219            use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
220
221            // Group the nodes by lower subtrie. This must be collected into a Vec in order for
222            // rayon's `zip` to be happy.
223            let node_groups: Vec<_> = lower_nodes
224                .chunk_by(|node_a, node_b| {
225                    SparseSubtrieType::from_path(&node_a.path) ==
226                        SparseSubtrieType::from_path(&node_b.path)
227                })
228                .collect();
229
230            // Take the lower subtries in the same order that the nodes were grouped into, so that
231            // the two can be zipped together. This also must be collected into a Vec for rayon's
232            // `zip` to be happy.
233            let lower_subtries: Vec<_> = node_groups
234                .iter()
235                .map(|nodes| {
236                    // NOTE: chunk_by won't produce empty groups
237                    let node = &nodes[0];
238                    let idx =
239                        SparseSubtrieType::from_path(&node.path).lower_index().unwrap_or_else(
240                            || panic!("upper subtrie node {node:?} found amongst lower nodes"),
241                        );
242                    // due to the nodes being sorted secondarily on their path, and chunk_by keeping
243                    // the first element of each group, the `path` here will necessarily be the
244                    // shortest path being revealed for each subtrie. Therefore we can reveal the
245                    // subtrie itself using this path and retain correct behavior.
246                    self.lower_subtries[idx].reveal(&node.path);
247                    (idx, self.lower_subtries[idx].take_revealed().expect("just revealed"))
248                })
249                .collect();
250
251            let (tx, rx) = mpsc::channel();
252
253            // Zip the lower subtries and their corresponding node groups, and reveal lower subtrie
254            // nodes in parallel
255            lower_subtries
256                .into_par_iter()
257                .zip(node_groups.into_par_iter())
258                .map(|((subtrie_idx, mut subtrie), nodes)| {
259                    // reserve space in the HashMap ahead of time; doing it on a node-by-node basis
260                    // can cause multiple re-allocations as the hashmap grows.
261                    subtrie.nodes.reserve(nodes.len());
262
263                    for node in nodes {
264                        // Reveal each node in the subtrie, returning early on any errors
265                        let res = subtrie.reveal_node(node.path, &node.node, node.masks);
266                        if res.is_err() {
267                            return (subtrie_idx, subtrie, res)
268                        }
269                    }
270                    (subtrie_idx, subtrie, Ok(()))
271                })
272                .for_each_init(|| tx.clone(), |tx, result| tx.send(result).unwrap());
273
274            drop(tx);
275
276            // Take back all lower subtries which were sent to the rayon pool, collecting the last
277            // seen error in the process and returning that. If we don't fully drain the channel
278            // then we lose lower sparse tries, putting the whole ParallelSparseTrie in an
279            // inconsistent state.
280            let mut any_err = Ok(());
281            for (subtrie_idx, subtrie, res) in rx {
282                self.lower_subtries[subtrie_idx] = LowerSparseSubtrie::Revealed(subtrie);
283                if res.is_err() {
284                    any_err = res;
285                }
286            }
287
288            any_err
289        }
290    }
291
292    fn update_leaf<P: TrieNodeProvider>(
293        &mut self,
294        full_path: Nibbles,
295        value: Vec<u8>,
296        provider: P,
297    ) -> SparseTrieResult<()> {
298        self.prefix_set.insert(full_path);
299        let existing = self.upper_subtrie.inner.values.insert(full_path, value.clone());
300        if existing.is_some() {
301            // upper trie structure unchanged, return immediately
302            return Ok(())
303        }
304
305        let retain_updates = self.updates_enabled();
306
307        // Start at the root, traversing until we find either the node to update or a subtrie to
308        // update.
309        //
310        // We first traverse the upper subtrie for two levels, and moving any created nodes to a
311        // lower subtrie if necessary.
312        //
313        // We use `next` to keep track of the next node that we need to traverse to, and
314        // `new_nodes` to keep track of any nodes that were created during the traversal.
315        let mut new_nodes = Vec::new();
316        let mut next = Some(Nibbles::default());
317
318        // Traverse the upper subtrie to find the node to update or the subtrie to update.
319        //
320        // We stop when the next node to traverse would be in a lower subtrie, or if there are no
321        // more nodes to traverse.
322        while let Some(current) =
323            next.filter(|next| SparseSubtrieType::path_len_is_upper(next.len()))
324        {
325            // Traverse the next node, keeping track of any changed nodes and the next step in the
326            // trie
327            match self.upper_subtrie.update_next_node(current, &full_path, retain_updates)? {
328                LeafUpdateStep::Continue { next_node } => {
329                    next = Some(next_node);
330                }
331                LeafUpdateStep::Complete { inserted_nodes, reveal_path } => {
332                    new_nodes.extend(inserted_nodes);
333
334                    if let Some(reveal_path) = reveal_path {
335                        let subtrie = self.subtrie_for_path_mut(&reveal_path);
336                        if subtrie.nodes.get(&reveal_path).expect("node must exist").is_hash() {
337                            warn!(
338                                target: "trie::parallel_sparse",
339                                child_path = ?reveal_path,
340                                leaf_full_path = ?full_path,
341                                "Extension node child not revealed in update_leaf, falling back to db",
342                            );
343                            if let Some(RevealedNode { node, tree_mask, hash_mask }) =
344                                provider.trie_node(&reveal_path)?
345                            {
346                                let decoded = TrieNode::decode(&mut &node[..])?;
347                                trace!(
348                                    target: "trie::parallel_sparse",
349                                    ?reveal_path,
350                                    ?decoded,
351                                    ?tree_mask,
352                                    ?hash_mask,
353                                    "Revealing child (from upper)",
354                                );
355                                subtrie.reveal_node(
356                                    reveal_path,
357                                    &decoded,
358                                    TrieMasks { hash_mask, tree_mask },
359                                )?;
360                            } else {
361                                return Err(SparseTrieErrorKind::NodeNotFoundInProvider {
362                                    path: reveal_path,
363                                }
364                                .into())
365                            }
366                        }
367                    }
368
369                    next = None;
370                }
371                LeafUpdateStep::NodeNotFound => {
372                    next = None;
373                }
374            }
375        }
376
377        // Move nodes from upper subtrie to lower subtries
378        for node_path in &new_nodes {
379            // Skip nodes that belong in the upper subtrie
380            if SparseSubtrieType::path_len_is_upper(node_path.len()) {
381                continue
382            }
383
384            let node =
385                self.upper_subtrie.nodes.remove(node_path).expect("node belongs to upper subtrie");
386
387            // If it's a leaf node, extract its value before getting mutable reference to subtrie.
388            let leaf_value = if let SparseNode::Leaf { key, .. } = &node {
389                let mut leaf_full_path = *node_path;
390                leaf_full_path.extend(key);
391                Some((
392                    leaf_full_path,
393                    self.upper_subtrie
394                        .inner
395                        .values
396                        .remove(&leaf_full_path)
397                        .expect("leaf nodes have associated values entries"),
398                ))
399            } else {
400                None
401            };
402
403            // Get or create the subtrie with the exact node path (not truncated to 2 nibbles).
404            let subtrie = self.subtrie_for_path_mut(node_path);
405
406            // Insert the leaf value if we have one
407            if let Some((leaf_full_path, value)) = leaf_value {
408                subtrie.inner.values.insert(leaf_full_path, value);
409            }
410
411            // Insert the node into the lower subtrie
412            subtrie.nodes.insert(*node_path, node);
413        }
414
415        // If we reached the max depth of the upper trie, we may have had more nodes to insert.
416        if let Some(next_path) = next.filter(|n| !SparseSubtrieType::path_len_is_upper(n.len())) {
417            // The value was inserted into the upper subtrie's `values` at the top of this method.
418            // At this point we know the value is not in the upper subtrie, and the call to
419            // `update_leaf` below will insert it into the lower subtrie. So remove it from the
420            // upper subtrie.
421            self.upper_subtrie.inner.values.remove(&full_path);
422
423            // Use subtrie_for_path to ensure the subtrie has the correct path.
424            //
425            // The next_path here represents where we need to continue traversal, which may
426            // be longer than 2 nibbles if we're following an extension node.
427            let subtrie = self.subtrie_for_path_mut(&next_path);
428
429            // Create an empty root at the subtrie path if the subtrie is empty
430            if subtrie.nodes.is_empty() {
431                subtrie.nodes.insert(subtrie.path, SparseNode::Empty);
432            }
433
434            // If we didn't update the target leaf, we need to call update_leaf on the subtrie
435            // to ensure that the leaf is updated correctly.
436            subtrie.update_leaf(full_path, value, provider, retain_updates)?;
437        }
438
439        Ok(())
440    }
441
442    fn remove_leaf<P: TrieNodeProvider>(
443        &mut self,
444        full_path: &Nibbles,
445        provider: P,
446    ) -> SparseTrieResult<()> {
447        // When removing a leaf node it's possibly necessary to modify its parent node, and possibly
448        // the parent's parent node. It is not ever necessary to descend further than that; once an
449        // extension node is hit it must terminate in a branch or the root, which won't need further
450        // updates. So the situation with maximum updates is:
451        //
452        // - Leaf
453        // - Branch with 2 children, one being this leaf
454        // - Extension
455        //
456        // ...which will result in just a leaf or extension, depending on what the branch's other
457        // child is.
458        //
459        // Therefore, first traverse the trie in order to find the leaf node and at most its parent
460        // and grandparent.
461
462        let leaf_path;
463        let leaf_subtrie;
464
465        let mut branch_parent_path: Option<Nibbles> = None;
466        let mut branch_parent_node: Option<SparseNode> = None;
467
468        let mut ext_grandparent_path: Option<Nibbles> = None;
469        let mut ext_grandparent_node: Option<SparseNode> = None;
470
471        let mut curr_path = Nibbles::new(); // start traversal from root
472        let mut curr_subtrie = self.upper_subtrie.as_mut();
473        let mut curr_subtrie_is_upper = true;
474
475        // List of node paths which need to have their hashes reset
476        let mut paths_to_reset_hashes = Vec::new();
477
478        loop {
479            let curr_node = curr_subtrie.nodes.get_mut(&curr_path).unwrap();
480
481            match Self::find_next_to_leaf(&curr_path, curr_node, full_path) {
482                FindNextToLeafOutcome::NotFound => return Ok(()), // leaf isn't in the trie
483                FindNextToLeafOutcome::BlindedNode(hash) => {
484                    return Err(SparseTrieErrorKind::BlindedNode { path: curr_path, hash }.into())
485                }
486                FindNextToLeafOutcome::Found => {
487                    // this node is the target leaf
488                    leaf_path = curr_path;
489                    leaf_subtrie = curr_subtrie;
490                    break;
491                }
492                FindNextToLeafOutcome::ContinueFrom(next_path) => {
493                    // Any branches/extensions along the path to the leaf will have their `hash`
494                    // field unset, as it will no longer be valid once the leaf is removed.
495                    match curr_node {
496                        SparseNode::Branch { hash, .. } => {
497                            if hash.is_some() {
498                                paths_to_reset_hashes
499                                    .push((SparseSubtrieType::from_path(&curr_path), curr_path));
500                            }
501
502                            // If there is already an extension leading into a branch, then that
503                            // extension is no longer relevant.
504                            match (&branch_parent_path, &ext_grandparent_path) {
505                                (Some(branch), Some(ext)) if branch.len() > ext.len() => {
506                                    ext_grandparent_path = None;
507                                    ext_grandparent_node = None;
508                                }
509                                _ => (),
510                            };
511                            branch_parent_path = Some(curr_path);
512                            branch_parent_node = Some(curr_node.clone());
513                        }
514                        SparseNode::Extension { hash, .. } => {
515                            if hash.is_some() {
516                                paths_to_reset_hashes
517                                    .push((SparseSubtrieType::from_path(&curr_path), curr_path));
518                            }
519
520                            // We can assume a new branch node will be found after the extension, so
521                            // there's no need to modify branch_parent_path/node even if it's
522                            // already set.
523                            ext_grandparent_path = Some(curr_path);
524                            ext_grandparent_node = Some(curr_node.clone());
525                        }
526                        SparseNode::Empty | SparseNode::Hash(_) | SparseNode::Leaf { .. } => {
527                            unreachable!(
528                                "find_next_to_leaf only continues to a branch or extension"
529                            )
530                        }
531                    }
532
533                    curr_path = next_path;
534
535                    // If we were previously looking at the upper trie, and the new path is in the
536                    // lower trie, we need to pull out a ref to the lower trie.
537                    if curr_subtrie_is_upper {
538                        if let SparseSubtrieType::Lower(idx) =
539                            SparseSubtrieType::from_path(&curr_path)
540                        {
541                            curr_subtrie = self.lower_subtries[idx]
542                                .as_revealed_mut()
543                                .expect("lower subtrie is revealed");
544                            curr_subtrie_is_upper = false;
545                        }
546                    }
547                }
548            };
549        }
550
551        // We've traversed to the leaf and collected its ancestors as necessary. Remove the leaf
552        // from its SparseSubtrie and reset the hashes of the nodes along the path.
553        self.prefix_set.insert(*full_path);
554        leaf_subtrie.inner.values.remove(full_path);
555        for (subtrie_type, path) in paths_to_reset_hashes {
556            let node = match subtrie_type {
557                SparseSubtrieType::Upper => self.upper_subtrie.nodes.get_mut(&path),
558                SparseSubtrieType::Lower(idx) => self.lower_subtries[idx]
559                    .as_revealed_mut()
560                    .expect("lower subtrie is revealed")
561                    .nodes
562                    .get_mut(&path),
563            }
564            .expect("node exists");
565
566            match node {
567                SparseNode::Extension { hash, .. } | SparseNode::Branch { hash, .. } => {
568                    *hash = None
569                }
570                SparseNode::Empty | SparseNode::Hash(_) | SparseNode::Leaf { .. } => {
571                    unreachable!("only branch and extension node hashes can be reset")
572                }
573            }
574        }
575        self.remove_node(&leaf_path);
576
577        // If the leaf was at the root replace its node with the empty value. We can stop execution
578        // here, all remaining logic is related to the ancestors of the leaf.
579        if leaf_path.is_empty() {
580            self.upper_subtrie.nodes.insert(leaf_path, SparseNode::Empty);
581            return Ok(())
582        }
583
584        // If there is a parent branch node (very likely, unless the leaf is at the root) execute
585        // any required changes for that node, relative to the removed leaf.
586        if let (Some(branch_path), Some(SparseNode::Branch { mut state_mask, .. })) =
587            (&branch_parent_path, &branch_parent_node)
588        {
589            let child_nibble = leaf_path.get_unchecked(branch_path.len());
590            state_mask.unset_bit(child_nibble);
591
592            let new_branch_node = if state_mask.count_bits() == 1 {
593                // If only one child is left set in the branch node, we need to collapse it. Get
594                // full path of the only child node left.
595                let remaining_child_path = {
596                    let mut p = *branch_path;
597                    p.push_unchecked(
598                        state_mask.first_set_bit_index().expect("state mask is not empty"),
599                    );
600                    p
601                };
602
603                trace!(
604                    target: "trie::parallel_sparse",
605                    ?leaf_path,
606                    ?branch_path,
607                    ?remaining_child_path,
608                    "Branch node has only one child",
609                );
610
611                let remaining_child_subtrie = self.subtrie_for_path_mut(&remaining_child_path);
612
613                // If the remaining child node is not yet revealed then we have to reveal it here,
614                // otherwise it's not possible to know how to collapse the branch.
615                let remaining_child_node =
616                    match remaining_child_subtrie.nodes.get(&remaining_child_path).unwrap() {
617                        SparseNode::Hash(_) => {
618                            warn!(
619                                target: "trie::parallel_sparse",
620                                child_path = ?remaining_child_path,
621                                leaf_full_path = ?full_path,
622                                "Branch node child not revealed in remove_leaf, falling back to db",
623                            );
624                            if let Some(RevealedNode { node, tree_mask, hash_mask }) =
625                                provider.trie_node(&remaining_child_path)?
626                            {
627                                let decoded = TrieNode::decode(&mut &node[..])?;
628                                trace!(
629                                    target: "trie::parallel_sparse",
630                                    ?remaining_child_path,
631                                    ?decoded,
632                                    ?tree_mask,
633                                    ?hash_mask,
634                                    "Revealing remaining blinded branch child"
635                                );
636                                remaining_child_subtrie.reveal_node(
637                                    remaining_child_path,
638                                    &decoded,
639                                    TrieMasks { hash_mask, tree_mask },
640                                )?;
641                                remaining_child_subtrie.nodes.get(&remaining_child_path).unwrap()
642                            } else {
643                                return Err(SparseTrieErrorKind::NodeNotFoundInProvider {
644                                    path: remaining_child_path,
645                                }
646                                .into())
647                            }
648                        }
649                        node => node,
650                    };
651
652                let (new_branch_node, remove_child) = Self::branch_changes_on_leaf_removal(
653                    branch_path,
654                    &remaining_child_path,
655                    remaining_child_node,
656                );
657
658                if remove_child {
659                    self.move_value_on_leaf_removal(
660                        branch_path,
661                        &new_branch_node,
662                        &remaining_child_path,
663                    );
664                    self.remove_node(&remaining_child_path);
665                }
666
667                if let Some(updates) = self.updates.as_mut() {
668                    updates.updated_nodes.remove(branch_path);
669                    updates.removed_nodes.insert(*branch_path);
670                }
671
672                new_branch_node
673            } else {
674                // If more than one child is left set in the branch, we just re-insert it with the
675                // updated state_mask.
676                SparseNode::new_branch(state_mask)
677            };
678
679            let branch_subtrie = self.subtrie_for_path_mut(branch_path);
680            branch_subtrie.nodes.insert(*branch_path, new_branch_node.clone());
681            branch_parent_node = Some(new_branch_node);
682        };
683
684        // If there is a grandparent extension node then there will necessarily be a parent branch
685        // node. Execute any required changes for the extension node, relative to the (possibly now
686        // replaced with a leaf or extension) branch node.
687        if let (Some(ext_path), Some(SparseNode::Extension { key: shortkey, .. })) =
688            (ext_grandparent_path, &ext_grandparent_node)
689        {
690            let ext_subtrie = self.subtrie_for_path_mut(&ext_path);
691            let branch_path = branch_parent_path.as_ref().unwrap();
692
693            if let Some(new_ext_node) = Self::extension_changes_on_leaf_removal(
694                &ext_path,
695                shortkey,
696                branch_path,
697                branch_parent_node.as_ref().unwrap(),
698            ) {
699                ext_subtrie.nodes.insert(ext_path, new_ext_node.clone());
700                self.move_value_on_leaf_removal(&ext_path, &new_ext_node, branch_path);
701                self.remove_node(branch_path);
702            }
703        }
704
705        Ok(())
706    }
707
708    fn root(&mut self) -> B256 {
709        trace!(target: "trie::parallel_sparse", "Calculating trie root hash");
710
711        // Update all lower subtrie hashes
712        self.update_subtrie_hashes();
713
714        // Update hashes for the upper subtrie using our specialized function
715        // that can access both upper and lower subtrie nodes
716        let mut prefix_set = core::mem::take(&mut self.prefix_set).freeze();
717        let root_rlp = self.update_upper_subtrie_hashes(&mut prefix_set);
718
719        // Return the root hash
720        root_rlp.as_hash().unwrap_or(EMPTY_ROOT_HASH)
721    }
722
723    fn update_subtrie_hashes(&mut self) {
724        trace!(target: "trie::parallel_sparse", "Updating subtrie hashes");
725
726        // Take changed subtries according to the prefix set
727        let mut prefix_set = core::mem::take(&mut self.prefix_set).freeze();
728        let (subtries, unchanged_prefix_set) = self.take_changed_lower_subtries(&mut prefix_set);
729
730        // update metrics
731        #[cfg(feature = "metrics")]
732        self.metrics.subtries_updated.record(subtries.len() as f64);
733
734        // Update the prefix set with the keys that didn't have matching subtries
735        self.prefix_set = unchanged_prefix_set;
736
737        let (tx, rx) = mpsc::channel();
738
739        #[cfg(not(feature = "std"))]
740        // Update subtrie hashes serially if nostd
741        for ChangedSubtrie { index, mut subtrie, mut prefix_set, mut update_actions_buf } in
742            subtries
743        {
744            subtrie.update_hashes(
745                &mut prefix_set,
746                &mut update_actions_buf,
747                &self.branch_node_tree_masks,
748                &self.branch_node_hash_masks,
749            );
750            tx.send((index, subtrie, update_actions_buf)).unwrap();
751        }
752
753        #[cfg(feature = "std")]
754        // Update subtrie hashes in parallel
755        {
756            use rayon::iter::{IntoParallelIterator, ParallelIterator};
757            let branch_node_tree_masks = &self.branch_node_tree_masks;
758            let branch_node_hash_masks = &self.branch_node_hash_masks;
759            subtries
760                .into_par_iter()
761                .map(
762                    |ChangedSubtrie {
763                         index,
764                         mut subtrie,
765                         mut prefix_set,
766                         mut update_actions_buf,
767                     }| {
768                        #[cfg(feature = "metrics")]
769                        let start = std::time::Instant::now();
770                        subtrie.update_hashes(
771                            &mut prefix_set,
772                            &mut update_actions_buf,
773                            branch_node_tree_masks,
774                            branch_node_hash_masks,
775                        );
776                        #[cfg(feature = "metrics")]
777                        self.metrics.subtrie_hash_update_latency.record(start.elapsed());
778                        (index, subtrie, update_actions_buf)
779                    },
780                )
781                .for_each_init(|| tx.clone(), |tx, result| tx.send(result).unwrap());
782        }
783
784        drop(tx);
785
786        // Return updated subtries back to the trie after executing any actions required on the
787        // top-level `SparseTrieUpdates`.
788        for (index, subtrie, update_actions_buf) in rx {
789            if let Some(mut update_actions_buf) = update_actions_buf {
790                self.apply_subtrie_update_actions(
791                    #[allow(clippy::iter_with_drain)]
792                    update_actions_buf.drain(..),
793                );
794                self.update_actions_buffers.push(update_actions_buf);
795            }
796
797            self.lower_subtries[index] = LowerSparseSubtrie::Revealed(subtrie);
798        }
799    }
800
801    fn get_leaf_value(&self, full_path: &Nibbles) -> Option<&Vec<u8>> {
802        self.subtrie_for_path(full_path).and_then(|subtrie| subtrie.inner.values.get(full_path))
803    }
804
805    fn updates_ref(&self) -> Cow<'_, SparseTrieUpdates> {
806        self.updates.as_ref().map_or(Cow::Owned(SparseTrieUpdates::default()), Cow::Borrowed)
807    }
808
809    fn take_updates(&mut self) -> SparseTrieUpdates {
810        self.updates.take().unwrap_or_default()
811    }
812
813    fn wipe(&mut self) {
814        self.upper_subtrie.wipe();
815        self.lower_subtries = [const { LowerSparseSubtrie::Blind(None) }; NUM_LOWER_SUBTRIES];
816        self.prefix_set = PrefixSetMut::all();
817        self.updates = self.updates.is_some().then(SparseTrieUpdates::wiped);
818    }
819
820    fn clear(&mut self) {
821        self.upper_subtrie.clear();
822        self.upper_subtrie.nodes.insert(Nibbles::default(), SparseNode::Empty);
823        for subtrie in &mut self.lower_subtries {
824            subtrie.clear();
825        }
826        self.prefix_set.clear();
827        self.updates = None;
828        self.branch_node_tree_masks.clear();
829        self.branch_node_hash_masks.clear();
830        // `update_actions_buffers` doesn't need to be cleared; we want to reuse the Vecs it has
831        // buffered, and all of those are already inherently cleared when they get used.
832    }
833
834    fn find_leaf(
835        &self,
836        full_path: &Nibbles,
837        expected_value: Option<&Vec<u8>>,
838    ) -> Result<LeafLookup, LeafLookupError> {
839        // Inclusion proof
840        //
841        // First, do a quick check if the value exists in either the upper or lower subtrie's values
842        // map. We assume that if there exists a leaf node, then its value will be in the `values`
843        // map.
844        if let Some(actual_value) = std::iter::once(self.upper_subtrie.as_ref())
845            .chain(self.lower_subtrie_for_path(full_path))
846            .filter_map(|subtrie| subtrie.inner.values.get(full_path))
847            .next()
848        {
849            // We found the leaf, check if the value matches (if expected value was provided)
850            return expected_value
851                .is_none_or(|v| v == actual_value)
852                .then_some(LeafLookup::Exists)
853                .ok_or_else(|| LeafLookupError::ValueMismatch {
854                    path: *full_path,
855                    expected: expected_value.cloned(),
856                    actual: actual_value.clone(),
857                })
858        }
859
860        // If the value does not exist in the `values` map, then this means that the leaf either:
861        // - Does not exist in the trie
862        // - Is missing from the witness
863        // We traverse the trie to find the location where this leaf would have been, showing
864        // that it is not in the trie. Or we find a blinded node, showing that the witness is
865        // not complete.
866        let mut curr_path = Nibbles::new(); // start traversal from root
867        let mut curr_subtrie = self.upper_subtrie.as_ref();
868        let mut curr_subtrie_is_upper = true;
869
870        loop {
871            let curr_node = curr_subtrie.nodes.get(&curr_path).unwrap();
872
873            match Self::find_next_to_leaf(&curr_path, curr_node, full_path) {
874                FindNextToLeafOutcome::NotFound => return Ok(LeafLookup::NonExistent),
875                FindNextToLeafOutcome::BlindedNode(hash) => {
876                    // We hit a blinded node - cannot determine if leaf exists
877                    return Err(LeafLookupError::BlindedNode { path: curr_path, hash });
878                }
879                FindNextToLeafOutcome::Found => {
880                    panic!("target leaf {full_path:?} found at path {curr_path:?}, even though value wasn't in values hashmap");
881                }
882                FindNextToLeafOutcome::ContinueFrom(next_path) => {
883                    curr_path = next_path;
884                    // If we were previously looking at the upper trie, and the new path is in the
885                    // lower trie, we need to pull out a ref to the lower trie.
886                    if curr_subtrie_is_upper {
887                        if let Some(lower_subtrie) = self.lower_subtrie_for_path(&curr_path) {
888                            curr_subtrie = lower_subtrie;
889                            curr_subtrie_is_upper = false;
890                        }
891                    }
892                }
893            }
894        }
895    }
896}
897
898impl ParallelSparseTrie {
899    /// Returns true if retaining updates is enabled for the overall trie.
900    const fn updates_enabled(&self) -> bool {
901        self.updates.is_some()
902    }
903
904    /// Creates a new revealed sparse trie from the given root node.
905    ///
906    /// This function initializes the internal structures and then reveals the root.
907    /// It is a convenient method to create a trie when you already have the root node available.
908    ///
909    /// # Arguments
910    ///
911    /// * `root` - The root node of the trie
912    /// * `masks` - Trie masks for root branch node
913    /// * `retain_updates` - Whether to track updates
914    ///
915    /// # Returns
916    ///
917    /// Self if successful, or an error if revealing fails.
918    pub fn from_root(
919        root: TrieNode,
920        masks: TrieMasks,
921        retain_updates: bool,
922    ) -> SparseTrieResult<Self> {
923        Self::default().with_root(root, masks, retain_updates)
924    }
925
926    /// Returns a reference to the lower `SparseSubtrie` for the given path, or None if the
927    /// path belongs to the upper trie, or if the lower subtrie for the path doesn't exist or is
928    /// blinded.
929    fn lower_subtrie_for_path(&self, path: &Nibbles) -> Option<&SparseSubtrie> {
930        match SparseSubtrieType::from_path(path) {
931            SparseSubtrieType::Upper => None,
932            SparseSubtrieType::Lower(idx) => self.lower_subtries[idx].as_revealed_ref(),
933        }
934    }
935
936    /// Returns a mutable reference to the lower `SparseSubtrie` for the given path, or None if the
937    /// path belongs to the upper trie.
938    ///
939    /// This method will create/reveal a new lower subtrie for the given path if one isn't already.
940    /// If one does exist, but its path field is longer than the given path, then the field will be
941    /// set to the given path.
942    fn lower_subtrie_for_path_mut(&mut self, path: &Nibbles) -> Option<&mut SparseSubtrie> {
943        match SparseSubtrieType::from_path(path) {
944            SparseSubtrieType::Upper => None,
945            SparseSubtrieType::Lower(idx) => {
946                self.lower_subtries[idx].reveal(path);
947                Some(self.lower_subtries[idx].as_revealed_mut().expect("just revealed"))
948            }
949        }
950    }
951
952    /// Returns a reference to either the lower or upper `SparseSubtrie` for the given path,
953    /// depending on the path's length.
954    ///
955    /// Returns `None` if a lower subtrie does not exist for the given path.
956    fn subtrie_for_path(&self, path: &Nibbles) -> Option<&SparseSubtrie> {
957        // We can't just call `lower_subtrie_for_path` and return `upper_subtrie` if it returns
958        // None, because Rust complains about double mutable borrowing `self`.
959        if SparseSubtrieType::path_len_is_upper(path.len()) {
960            Some(&self.upper_subtrie)
961        } else {
962            self.lower_subtrie_for_path(path)
963        }
964    }
965
966    /// Returns a mutable reference to either the lower or upper `SparseSubtrie` for the given path,
967    /// depending on the path's length.
968    ///
969    /// This method will create/reveal a new lower subtrie for the given path if one isn't already.
970    /// If one does exist, but its path field is longer than the given path, then the field will be
971    /// set to the given path.
972    fn subtrie_for_path_mut(&mut self, path: &Nibbles) -> &mut SparseSubtrie {
973        // We can't just call `lower_subtrie_for_path` and return `upper_subtrie` if it returns
974        // None, because Rust complains about double mutable borrowing `self`.
975        if SparseSubtrieType::path_len_is_upper(path.len()) {
976            &mut self.upper_subtrie
977        } else {
978            self.lower_subtrie_for_path_mut(path).unwrap()
979        }
980    }
981
982    /// Returns the next node in the traversal path from the given path towards the leaf for the
983    /// given full leaf path, or an error if any node along the traversal path is not revealed.
984    ///
985    ///
986    /// ## Panics
987    ///
988    /// If `from_path` is not a prefix of `leaf_full_path`.
989    fn find_next_to_leaf(
990        from_path: &Nibbles,
991        from_node: &SparseNode,
992        leaf_full_path: &Nibbles,
993    ) -> FindNextToLeafOutcome {
994        debug_assert!(leaf_full_path.len() >= from_path.len());
995        debug_assert!(leaf_full_path.starts_with(from_path));
996
997        match from_node {
998            // If empty node is found it means the subtrie doesn't have any nodes in it, let alone
999            // the target leaf.
1000            SparseNode::Empty => FindNextToLeafOutcome::NotFound,
1001            SparseNode::Hash(hash) => FindNextToLeafOutcome::BlindedNode(*hash),
1002            SparseNode::Leaf { key, .. } => {
1003                let mut found_full_path = *from_path;
1004                found_full_path.extend(key);
1005
1006                if &found_full_path == leaf_full_path {
1007                    return FindNextToLeafOutcome::Found
1008                }
1009                FindNextToLeafOutcome::NotFound
1010            }
1011            SparseNode::Extension { key, .. } => {
1012                if leaf_full_path.len() == from_path.len() {
1013                    return FindNextToLeafOutcome::NotFound
1014                }
1015
1016                let mut child_path = *from_path;
1017                child_path.extend(key);
1018
1019                if !leaf_full_path.starts_with(&child_path) {
1020                    return FindNextToLeafOutcome::NotFound
1021                }
1022                FindNextToLeafOutcome::ContinueFrom(child_path)
1023            }
1024            SparseNode::Branch { state_mask, .. } => {
1025                if leaf_full_path.len() == from_path.len() {
1026                    return FindNextToLeafOutcome::NotFound
1027                }
1028
1029                let nibble = leaf_full_path.get_unchecked(from_path.len());
1030                if !state_mask.is_bit_set(nibble) {
1031                    return FindNextToLeafOutcome::NotFound
1032                }
1033
1034                let mut child_path = *from_path;
1035                child_path.push_unchecked(nibble);
1036
1037                FindNextToLeafOutcome::ContinueFrom(child_path)
1038            }
1039        }
1040    }
1041
1042    /// Called when a child node has collapsed into its parent as part of `remove_leaf`. If the
1043    /// new parent node is a leaf, then the previous child also was, and if the previous child was
1044    /// on a lower subtrie while the parent is on an upper then the leaf value needs to be moved to
1045    /// the upper.
1046    fn move_value_on_leaf_removal(
1047        &mut self,
1048        parent_path: &Nibbles,
1049        new_parent_node: &SparseNode,
1050        prev_child_path: &Nibbles,
1051    ) {
1052        // If the parent path isn't in the upper then it doesn't matter what the new node is,
1053        // there's no situation where a leaf value needs to be moved.
1054        if SparseSubtrieType::from_path(parent_path).lower_index().is_some() {
1055            return;
1056        }
1057
1058        if let SparseNode::Leaf { key, .. } = new_parent_node {
1059            let Some(prev_child_subtrie) = self.lower_subtrie_for_path_mut(prev_child_path) else {
1060                return;
1061            };
1062
1063            let mut leaf_full_path = *parent_path;
1064            leaf_full_path.extend(key);
1065
1066            let val = prev_child_subtrie.inner.values.remove(&leaf_full_path).expect("ParallelSparseTrie is in an inconsistent state, expected value on subtrie which wasn't found");
1067            self.upper_subtrie.inner.values.insert(leaf_full_path, val);
1068        }
1069    }
1070
1071    /// Used by `remove_leaf` to ensure that when a node is removed from a lower subtrie that any
1072    /// externalities are handled. These can include:
1073    /// - Removing the lower subtrie completely, if it is now empty.
1074    /// - Updating the `path` field of the lower subtrie to indicate that its root node has changed.
1075    ///
1076    /// This method assumes that the caller will deal with putting all other nodes in the trie into
1077    /// a consistent state after the removal of this one.
1078    ///
1079    /// ## Panics
1080    ///
1081    /// - If the removed node was not a leaf or extension.
1082    fn remove_node(&mut self, path: &Nibbles) {
1083        let subtrie = self.subtrie_for_path_mut(path);
1084        let node = subtrie.nodes.remove(path);
1085
1086        let Some(idx) = SparseSubtrieType::from_path(path).lower_index() else {
1087            // When removing a node from the upper trie there's nothing special we need to do to fix
1088            // its path field; the upper trie's path is always empty.
1089            return;
1090        };
1091
1092        match node {
1093            Some(SparseNode::Leaf { .. }) => {
1094                // If the leaf was the final node in its lower subtrie then we can blind the
1095                // subtrie, effectively marking it as empty.
1096                if subtrie.nodes.is_empty() {
1097                    self.lower_subtries[idx].clear();
1098                }
1099            }
1100            Some(SparseNode::Extension { key, .. }) => {
1101                // If the removed extension was the root node of a lower subtrie then the lower
1102                // subtrie's `path` needs to be updated to be whatever node the extension used to
1103                // point to.
1104                if &subtrie.path == path {
1105                    subtrie.path.extend(&key);
1106                }
1107            }
1108            _ => panic!("Expected to remove a leaf or extension, but removed {node:?}"),
1109        }
1110    }
1111
1112    /// Given the path to a parent branch node and a child node which is the sole remaining child on
1113    /// that branch after removing a leaf, returns a node to replace the parent branch node and a
1114    /// boolean indicating if the child should be deleted.
1115    ///
1116    /// ## Panics
1117    ///
1118    /// - If either parent or child node is not already revealed.
1119    /// - If parent's path is not a prefix of the child's path.
1120    fn branch_changes_on_leaf_removal(
1121        parent_path: &Nibbles,
1122        remaining_child_path: &Nibbles,
1123        remaining_child_node: &SparseNode,
1124    ) -> (SparseNode, bool) {
1125        debug_assert!(remaining_child_path.len() > parent_path.len());
1126        debug_assert!(remaining_child_path.starts_with(parent_path));
1127
1128        let remaining_child_nibble = remaining_child_path.get_unchecked(parent_path.len());
1129
1130        // If we swap the branch node out either an extension or leaf, depending on
1131        // what its remaining child is.
1132        match remaining_child_node {
1133            SparseNode::Empty | SparseNode::Hash(_) => {
1134                panic!("remaining child must have been revealed already")
1135            }
1136            // If the only child is a leaf node, we downgrade the branch node into a
1137            // leaf node, prepending the nibble to the key, and delete the old
1138            // child.
1139            SparseNode::Leaf { key, .. } => {
1140                let mut new_key = Nibbles::from_nibbles_unchecked([remaining_child_nibble]);
1141                new_key.extend(key);
1142                (SparseNode::new_leaf(new_key), true)
1143            }
1144            // If the only child node is an extension node, we downgrade the branch
1145            // node into an even longer extension node, prepending the nibble to the
1146            // key, and delete the old child.
1147            SparseNode::Extension { key, .. } => {
1148                let mut new_key = Nibbles::from_nibbles_unchecked([remaining_child_nibble]);
1149                new_key.extend(key);
1150                (SparseNode::new_ext(new_key), true)
1151            }
1152            // If the only child is a branch node, we downgrade the current branch
1153            // node into a one-nibble extension node.
1154            SparseNode::Branch { .. } => (
1155                SparseNode::new_ext(Nibbles::from_nibbles_unchecked([remaining_child_nibble])),
1156                false,
1157            ),
1158        }
1159    }
1160
1161    /// Given the path to a parent extension and its key, and a child node (not necessarily on this
1162    /// subtrie), returns an optional replacement parent node. If a replacement is returned then the
1163    /// child node should be deleted.
1164    ///
1165    /// ## Panics
1166    ///
1167    /// - If either parent or child node is not already revealed.
1168    /// - If parent's path is not a prefix of the child's path.
1169    fn extension_changes_on_leaf_removal(
1170        parent_path: &Nibbles,
1171        parent_key: &Nibbles,
1172        child_path: &Nibbles,
1173        child: &SparseNode,
1174    ) -> Option<SparseNode> {
1175        debug_assert!(child_path.len() > parent_path.len());
1176        debug_assert!(child_path.starts_with(parent_path));
1177
1178        // If the parent node is an extension node, we need to look at its child to see
1179        // if we need to merge it.
1180        match child {
1181            SparseNode::Empty | SparseNode::Hash(_) => {
1182                panic!("child must be revealed")
1183            }
1184            // For a leaf node, we collapse the extension node into a leaf node,
1185            // extending the key. While it's impossible to encounter an extension node
1186            // followed by a leaf node in a complete trie, it's possible here because we
1187            // could have downgraded the extension node's child into a leaf node from a
1188            // branch in a previous call to `branch_changes_on_leaf_removal`.
1189            SparseNode::Leaf { key, .. } => {
1190                let mut new_key = *parent_key;
1191                new_key.extend(key);
1192                Some(SparseNode::new_leaf(new_key))
1193            }
1194            // Similar to the leaf node, for an extension node, we collapse them into one
1195            // extension node, extending the key.
1196            SparseNode::Extension { key, .. } => {
1197                let mut new_key = *parent_key;
1198                new_key.extend(key);
1199                Some(SparseNode::new_ext(new_key))
1200            }
1201            // For a branch node, we just leave the extension node as-is.
1202            SparseNode::Branch { .. } => None,
1203        }
1204    }
1205
1206    /// Drains any [`SparseTrieUpdatesAction`]s from the given subtrie, and applies each action to
1207    /// the given `updates` set. If the given set is None then this is a no-op.
1208    fn apply_subtrie_update_actions(
1209        &mut self,
1210        update_actions: impl Iterator<Item = SparseTrieUpdatesAction>,
1211    ) {
1212        if let Some(updates) = self.updates.as_mut() {
1213            for action in update_actions {
1214                match action {
1215                    SparseTrieUpdatesAction::InsertRemoved(path) => {
1216                        updates.updated_nodes.remove(&path);
1217                        updates.removed_nodes.insert(path);
1218                    }
1219                    SparseTrieUpdatesAction::RemoveUpdated(path) => {
1220                        updates.updated_nodes.remove(&path);
1221                    }
1222                    SparseTrieUpdatesAction::InsertUpdated(path, branch_node) => {
1223                        updates.updated_nodes.insert(path, branch_node);
1224                    }
1225                }
1226            }
1227        };
1228    }
1229
1230    /// Updates hashes for the upper subtrie, using nodes from both upper and lower subtries.
1231    #[instrument(level = "trace", target = "trie::parallel_sparse", skip_all, ret)]
1232    fn update_upper_subtrie_hashes(&mut self, prefix_set: &mut PrefixSet) -> RlpNode {
1233        trace!(target: "trie::parallel_sparse", "Updating upper subtrie hashes");
1234
1235        debug_assert!(self.upper_subtrie.inner.buffers.path_stack.is_empty());
1236        self.upper_subtrie.inner.buffers.path_stack.push(RlpNodePathStackItem {
1237            path: Nibbles::default(), // Start from root
1238            is_in_prefix_set: None,
1239        });
1240
1241        #[cfg(feature = "metrics")]
1242        let start = std::time::Instant::now();
1243
1244        let mut update_actions_buf =
1245            self.updates_enabled().then(|| self.update_actions_buffers.pop().unwrap_or_default());
1246
1247        while let Some(stack_item) = self.upper_subtrie.inner.buffers.path_stack.pop() {
1248            let path = stack_item.path;
1249            let node = if path.len() < UPPER_TRIE_MAX_DEPTH {
1250                self.upper_subtrie.nodes.get_mut(&path).expect("upper subtrie node must exist")
1251            } else {
1252                let index = path_subtrie_index_unchecked(&path);
1253                let node = self.lower_subtries[index]
1254                    .as_revealed_mut()
1255                    .expect("lower subtrie must exist")
1256                    .nodes
1257                    .get_mut(&path)
1258                    .expect("lower subtrie node must exist");
1259                // Lower subtrie root node hashes must be computed before updating upper subtrie
1260                // hashes
1261                debug_assert!(
1262                    node.hash().is_some(),
1263                    "Lower subtrie root node at path {path:?} has no hash"
1264                );
1265                node
1266            };
1267
1268            // Calculate the RLP node for the current node using upper subtrie
1269            self.upper_subtrie.inner.rlp_node(
1270                prefix_set,
1271                &mut update_actions_buf,
1272                stack_item,
1273                node,
1274                &self.branch_node_tree_masks,
1275                &self.branch_node_hash_masks,
1276            );
1277        }
1278
1279        // If there were any branch node updates as a result of calculating the RLP node for the
1280        // upper trie then apply them to the top-level set.
1281        if let Some(mut update_actions_buf) = update_actions_buf {
1282            self.apply_subtrie_update_actions(
1283                #[allow(clippy::iter_with_drain)]
1284                update_actions_buf.drain(..),
1285            );
1286            self.update_actions_buffers.push(update_actions_buf);
1287        }
1288
1289        #[cfg(feature = "metrics")]
1290        self.metrics.subtrie_upper_hash_latency.record(start.elapsed());
1291
1292        debug_assert_eq!(self.upper_subtrie.inner.buffers.rlp_node_stack.len(), 1);
1293        self.upper_subtrie.inner.buffers.rlp_node_stack.pop().unwrap().rlp_node
1294    }
1295
1296    /// Returns:
1297    /// 1. List of lower [subtries](SparseSubtrie) that have changed according to the provided
1298    ///    [prefix set](PrefixSet). See documentation of [`ChangedSubtrie`] for more details. Lower
1299    ///    subtries whose root node is missing a hash will also be returned; this is required to
1300    ///    handle cases where extensions/leafs get shortened and therefore moved from the upper to a
1301    ///    lower subtrie.
1302    /// 2. Prefix set of keys that do not belong to any lower subtrie.
1303    ///
1304    /// This method helps optimize hash recalculations by identifying which specific
1305    /// lower subtries need to be updated. Each lower subtrie can then be updated in parallel.
1306    ///
1307    /// IMPORTANT: The method removes the subtries from `lower_subtries`, and the caller is
1308    /// responsible for returning them back into the array.
1309    fn take_changed_lower_subtries(
1310        &mut self,
1311        prefix_set: &mut PrefixSet,
1312    ) -> (Vec<ChangedSubtrie>, PrefixSetMut) {
1313        // Clone the prefix set to iterate over its keys. Cloning is cheap, it's just an Arc.
1314        let prefix_set_clone = prefix_set.clone();
1315        let mut prefix_set_iter = prefix_set_clone.into_iter().copied().peekable();
1316        let mut changed_subtries = Vec::new();
1317        let mut unchanged_prefix_set = PrefixSetMut::default();
1318        let updates_enabled = self.updates_enabled();
1319
1320        for (index, subtrie) in self.lower_subtries.iter_mut().enumerate() {
1321            if let Some(subtrie) = subtrie.take_revealed_if(|subtrie| {
1322                prefix_set.contains(&subtrie.path) ||
1323                    subtrie.nodes.get(&subtrie.path).is_some_and(|n| n.hash().is_none())
1324            }) {
1325                let prefix_set = if prefix_set.all() {
1326                    unchanged_prefix_set = PrefixSetMut::all();
1327                    PrefixSetMut::all()
1328                } else {
1329                    // Take those keys from the original prefix set that start with the subtrie path
1330                    //
1331                    // Subtries are stored in the order of their paths, so we can use the same
1332                    // prefix set iterator.
1333                    let mut new_prefix_set = Vec::new();
1334                    while let Some(key) = prefix_set_iter.peek() {
1335                        if key.starts_with(&subtrie.path) {
1336                            // If the key starts with the subtrie path, add it to the new prefix set
1337                            new_prefix_set.push(prefix_set_iter.next().unwrap());
1338                        } else if new_prefix_set.is_empty() && key < &subtrie.path {
1339                            // If we didn't yet have any keys that belong to this subtrie, and the
1340                            // current key is still less than the subtrie path, add it to the
1341                            // unchanged prefix set
1342                            unchanged_prefix_set.insert(prefix_set_iter.next().unwrap());
1343                        } else {
1344                            // If we're past the subtrie path, we're done with this subtrie. Do not
1345                            // advance the iterator, the next key will be processed either by the
1346                            // next subtrie or inserted into the unchanged prefix set.
1347                            break
1348                        }
1349                    }
1350                    PrefixSetMut::from(new_prefix_set)
1351                }
1352                .freeze();
1353
1354                // We need the full path of root node of the lower subtrie to the unchanged prefix
1355                // set, so that we don't skip it when calculating hashes for the upper subtrie.
1356                match subtrie.nodes.get(&subtrie.path) {
1357                    Some(SparseNode::Extension { key, .. } | SparseNode::Leaf { key, .. }) => {
1358                        unchanged_prefix_set.insert(subtrie.path.join(key));
1359                    }
1360                    Some(SparseNode::Branch { .. }) => {
1361                        unchanged_prefix_set.insert(subtrie.path);
1362                    }
1363                    _ => {}
1364                }
1365
1366                let update_actions_buf =
1367                    updates_enabled.then(|| self.update_actions_buffers.pop().unwrap_or_default());
1368
1369                changed_subtries.push(ChangedSubtrie {
1370                    index,
1371                    subtrie,
1372                    prefix_set,
1373                    update_actions_buf,
1374                });
1375            }
1376        }
1377
1378        // Extend the unchanged prefix set with the remaining keys that are not part of any subtries
1379        unchanged_prefix_set.extend_keys(prefix_set_iter);
1380
1381        (changed_subtries, unchanged_prefix_set)
1382    }
1383
1384    /// Returns an iterator over all nodes in the trie in no particular order.
1385    #[cfg(test)]
1386    fn all_nodes(&self) -> impl IntoIterator<Item = (&Nibbles, &SparseNode)> {
1387        let mut nodes = vec![];
1388        for subtrie in self.lower_subtries.iter().filter_map(LowerSparseSubtrie::as_revealed_ref) {
1389            nodes.extend(subtrie.nodes.iter())
1390        }
1391        nodes.extend(self.upper_subtrie.nodes.iter());
1392        nodes
1393    }
1394
1395    /// Reveals a trie node in the upper trie if it has not been revealed before. When revealing
1396    /// branch/extension nodes this may recurse into a lower trie to reveal a child.
1397    ///
1398    /// This function decodes a trie node and inserts it into the trie structure. It handles
1399    /// different node types (leaf, extension, branch) by appropriately adding them to the trie and
1400    /// recursively revealing their children.
1401    ///
1402    /// # Arguments
1403    ///
1404    /// * `path` - The path where the node should be revealed
1405    /// * `node` - The trie node to reveal
1406    /// * `masks` - Trie masks for branch nodes
1407    ///
1408    /// # Returns
1409    ///
1410    /// `Ok(())` if successful, or an error if the node was not revealed.
1411    fn reveal_upper_node(
1412        &mut self,
1413        path: Nibbles,
1414        node: &TrieNode,
1415        masks: TrieMasks,
1416    ) -> SparseTrieResult<()> {
1417        // If there is no subtrie for the path it means the path is UPPER_TRIE_MAX_DEPTH or less
1418        // nibbles, and so belongs to the upper trie.
1419        self.upper_subtrie.reveal_node(path, node, masks)?;
1420
1421        // The previous upper_trie.reveal_node call will not have revealed any child nodes via
1422        // reveal_node_or_hash if the child node would be found on a lower subtrie. We handle that
1423        // here by manually checking the specific cases where this could happen, and calling
1424        // reveal_node_or_hash for each.
1425        match node {
1426            TrieNode::Branch(branch) => {
1427                // If a branch is at the cutoff level of the trie then it will be in the upper trie,
1428                // but all of its children will be in a lower trie. Check if a child node would be
1429                // in the lower subtrie, and reveal accordingly.
1430                if !SparseSubtrieType::path_len_is_upper(path.len() + 1) {
1431                    let mut stack_ptr = branch.as_ref().first_child_index();
1432                    for idx in CHILD_INDEX_RANGE {
1433                        if branch.state_mask.is_bit_set(idx) {
1434                            let mut child_path = path;
1435                            child_path.push_unchecked(idx);
1436                            self.lower_subtrie_for_path_mut(&child_path)
1437                                .expect("child_path must have a lower subtrie")
1438                                .reveal_node_or_hash(child_path, &branch.stack[stack_ptr])?;
1439                            stack_ptr += 1;
1440                        }
1441                    }
1442                }
1443            }
1444            TrieNode::Extension(ext) => {
1445                let mut child_path = path;
1446                child_path.extend(&ext.key);
1447                if let Some(subtrie) = self.lower_subtrie_for_path_mut(&child_path) {
1448                    subtrie.reveal_node_or_hash(child_path, &ext.child)?;
1449                }
1450            }
1451            TrieNode::EmptyRoot | TrieNode::Leaf(_) => (),
1452        }
1453
1454        Ok(())
1455    }
1456}
1457
1458/// This is a subtrie of the [`ParallelSparseTrie`] that contains a map from path to sparse trie
1459/// nodes.
1460#[derive(Clone, PartialEq, Eq, Debug, Default)]
1461pub struct SparseSubtrie {
1462    /// The root path of this subtrie.
1463    ///
1464    /// This is the _full_ path to this subtrie, meaning it includes the first
1465    /// [`UPPER_TRIE_MAX_DEPTH`] nibbles that we also use for indexing subtries in the
1466    /// [`ParallelSparseTrie`].
1467    ///
1468    /// There should be a node for this path in `nodes` map.
1469    pub(crate) path: Nibbles,
1470    /// The map from paths to sparse trie nodes within this subtrie.
1471    nodes: HashMap<Nibbles, SparseNode>,
1472    /// Subset of fields for mutable access while `nodes` field is also being mutably borrowed.
1473    inner: SparseSubtrieInner,
1474}
1475
1476/// Returned by the `find_next_to_leaf` method to indicate either that the leaf has been found,
1477/// traversal should be continued from the given path, or the leaf is not in the trie.
1478enum FindNextToLeafOutcome {
1479    /// `Found` indicates that the leaf was found at the given path.
1480    Found,
1481    /// `ContinueFrom` indicates that traversal should continue from the given path.
1482    ContinueFrom(Nibbles),
1483    /// `NotFound` indicates that there is no way to traverse to the leaf, as it is not in the
1484    /// trie.
1485    NotFound,
1486    /// `BlindedNode` indicates that the node is blinded with the contained hash and cannot be
1487    /// traversed.
1488    BlindedNode(B256),
1489}
1490
1491impl SparseSubtrie {
1492    /// Creates a new empty subtrie with the specified root path.
1493    pub(crate) fn new(path: Nibbles) -> Self {
1494        Self { path, ..Default::default() }
1495    }
1496
1497    /// Returns true if this subtrie has any nodes, false otherwise.
1498    pub(crate) fn is_empty(&self) -> bool {
1499        self.nodes.is_empty()
1500    }
1501
1502    /// Returns true if the current path and its child are both found in the same level.
1503    fn is_child_same_level(current_path: &Nibbles, child_path: &Nibbles) -> bool {
1504        let current_level = core::mem::discriminant(&SparseSubtrieType::from_path(current_path));
1505        let child_level = core::mem::discriminant(&SparseSubtrieType::from_path(child_path));
1506        current_level == child_level
1507    }
1508
1509    /// Updates or inserts a leaf node at the specified key path with the provided RLP-encoded
1510    /// value.
1511    ///
1512    /// If the leaf did not previously exist, this method adjusts the trie structure by inserting
1513    /// new leaf nodes, splitting branch nodes, or collapsing extension nodes as needed.
1514    ///
1515    /// # Returns
1516    ///
1517    /// Returns the `Ok` if the update is successful.
1518    ///
1519    /// Note: If an update requires revealing a blinded node, an error is returned if the blinded
1520    /// provider returns an error.
1521    pub fn update_leaf(
1522        &mut self,
1523        full_path: Nibbles,
1524        value: Vec<u8>,
1525        provider: impl TrieNodeProvider,
1526        retain_updates: bool,
1527    ) -> SparseTrieResult<()> {
1528        debug_assert!(full_path.starts_with(&self.path));
1529        let existing = self.inner.values.insert(full_path, value);
1530        if existing.is_some() {
1531            // trie structure unchanged, return immediately
1532            return Ok(())
1533        }
1534
1535        // Here we are starting at the root of the subtrie, and traversing from there.
1536        let mut current = Some(self.path);
1537        while let Some(current_path) = current {
1538            match self.update_next_node(current_path, &full_path, retain_updates)? {
1539                LeafUpdateStep::Continue { next_node } => {
1540                    current = Some(next_node);
1541                }
1542                LeafUpdateStep::Complete { reveal_path, .. } => {
1543                    if let Some(reveal_path) = reveal_path {
1544                        if self.nodes.get(&reveal_path).expect("node must exist").is_hash() {
1545                            warn!(
1546                                target: "trie::parallel_sparse",
1547                                child_path = ?reveal_path,
1548                                leaf_full_path = ?full_path,
1549                                "Extension node child not revealed in update_leaf, falling back to db",
1550                            );
1551                            if let Some(RevealedNode { node, tree_mask, hash_mask }) =
1552                                provider.trie_node(&reveal_path)?
1553                            {
1554                                let decoded = TrieNode::decode(&mut &node[..])?;
1555                                trace!(
1556                                    target: "trie::parallel_sparse",
1557                                    ?reveal_path,
1558                                    ?decoded,
1559                                    ?tree_mask,
1560                                    ?hash_mask,
1561                                    "Revealing child (from lower)",
1562                                );
1563                                self.reveal_node(
1564                                    reveal_path,
1565                                    &decoded,
1566                                    TrieMasks { hash_mask, tree_mask },
1567                                )?;
1568                            } else {
1569                                return Err(SparseTrieErrorKind::NodeNotFoundInProvider {
1570                                    path: reveal_path,
1571                                }
1572                                .into())
1573                            }
1574                        }
1575                    }
1576
1577                    current = None;
1578                }
1579                LeafUpdateStep::NodeNotFound => {
1580                    current = None;
1581                }
1582            }
1583        }
1584
1585        Ok(())
1586    }
1587
1588    /// Processes the current node, returning what to do next in the leaf update process.
1589    ///
1590    /// This will add or update any nodes in the trie as necessary.
1591    ///
1592    /// Returns a `LeafUpdateStep` containing the next node to process (if any) and
1593    /// the paths of nodes that were inserted during this step.
1594    fn update_next_node(
1595        &mut self,
1596        mut current: Nibbles,
1597        path: &Nibbles,
1598        retain_updates: bool,
1599    ) -> SparseTrieResult<LeafUpdateStep> {
1600        debug_assert!(path.starts_with(&self.path));
1601        debug_assert!(current.starts_with(&self.path));
1602        debug_assert!(path.starts_with(&current));
1603        let Some(node) = self.nodes.get_mut(&current) else {
1604            return Ok(LeafUpdateStep::NodeNotFound);
1605        };
1606        match node {
1607            SparseNode::Empty => {
1608                // We need to insert the node with a different path and key depending on the path of
1609                // the subtrie.
1610                let path = path.slice(self.path.len()..);
1611                *node = SparseNode::new_leaf(path);
1612                Ok(LeafUpdateStep::complete_with_insertions(vec![current], None))
1613            }
1614            SparseNode::Hash(hash) => {
1615                Err(SparseTrieErrorKind::BlindedNode { path: current, hash: *hash }.into())
1616            }
1617            SparseNode::Leaf { key: current_key, .. } => {
1618                current.extend(current_key);
1619
1620                // this leaf is being updated
1621                debug_assert!(
1622                    &current != path,
1623                    "we already checked leaf presence in the beginning"
1624                );
1625
1626                // find the common prefix
1627                let common = current.common_prefix_length(path);
1628
1629                // update existing node
1630                let new_ext_key = current.slice(current.len() - current_key.len()..common);
1631                *node = SparseNode::new_ext(new_ext_key);
1632
1633                // create a branch node and corresponding leaves
1634                self.nodes.reserve(3);
1635                let branch_path = current.slice(..common);
1636                let new_leaf_path = path.slice(..=common);
1637                let existing_leaf_path = current.slice(..=common);
1638
1639                self.nodes.insert(
1640                    branch_path,
1641                    SparseNode::new_split_branch(
1642                        current.get_unchecked(common),
1643                        path.get_unchecked(common),
1644                    ),
1645                );
1646                self.nodes.insert(new_leaf_path, SparseNode::new_leaf(path.slice(common + 1..)));
1647                self.nodes
1648                    .insert(existing_leaf_path, SparseNode::new_leaf(current.slice(common + 1..)));
1649
1650                Ok(LeafUpdateStep::complete_with_insertions(
1651                    vec![branch_path, new_leaf_path, existing_leaf_path],
1652                    None,
1653                ))
1654            }
1655            SparseNode::Extension { key, .. } => {
1656                current.extend(key);
1657
1658                if !path.starts_with(&current) {
1659                    // find the common prefix
1660                    let common = current.common_prefix_length(path);
1661                    *key = current.slice(current.len() - key.len()..common);
1662
1663                    // If branch node updates retention is enabled, we need to query the
1664                    // extension node child to later set the hash mask for a parent branch node
1665                    // correctly.
1666                    let reveal_path = retain_updates.then_some(current);
1667
1668                    // create state mask for new branch node
1669                    // NOTE: this might overwrite the current extension node
1670                    self.nodes.reserve(3);
1671                    let branch_path = current.slice(..common);
1672                    let new_leaf_path = path.slice(..=common);
1673                    let branch = SparseNode::new_split_branch(
1674                        current.get_unchecked(common),
1675                        path.get_unchecked(common),
1676                    );
1677
1678                    self.nodes.insert(branch_path, branch);
1679
1680                    // create new leaf
1681                    let new_leaf = SparseNode::new_leaf(path.slice(common + 1..));
1682                    self.nodes.insert(new_leaf_path, new_leaf);
1683
1684                    let mut inserted_nodes = vec![branch_path, new_leaf_path];
1685
1686                    // recreate extension to previous child if needed
1687                    let key = current.slice(common + 1..);
1688                    if !key.is_empty() {
1689                        let ext_path = current.slice(..=common);
1690                        self.nodes.insert(ext_path, SparseNode::new_ext(key));
1691                        inserted_nodes.push(ext_path);
1692                    }
1693
1694                    return Ok(LeafUpdateStep::complete_with_insertions(inserted_nodes, reveal_path))
1695                }
1696
1697                Ok(LeafUpdateStep::continue_with(current))
1698            }
1699            SparseNode::Branch { state_mask, .. } => {
1700                let nibble = path.get_unchecked(current.len());
1701                current.push_unchecked(nibble);
1702                if !state_mask.is_bit_set(nibble) {
1703                    state_mask.set_bit(nibble);
1704                    let new_leaf = SparseNode::new_leaf(path.slice(current.len()..));
1705                    self.nodes.insert(current, new_leaf);
1706                    return Ok(LeafUpdateStep::complete_with_insertions(vec![current], None))
1707                }
1708
1709                // If the nibble is set, we can continue traversing the branch.
1710                Ok(LeafUpdateStep::continue_with(current))
1711            }
1712        }
1713    }
1714
1715    /// Internal implementation of the method of the same name on `ParallelSparseTrie`.
1716    fn reveal_node(
1717        &mut self,
1718        path: Nibbles,
1719        node: &TrieNode,
1720        masks: TrieMasks,
1721    ) -> SparseTrieResult<()> {
1722        debug_assert!(path.starts_with(&self.path));
1723
1724        // If the node is already revealed and it's not a hash node, do nothing.
1725        if self.nodes.get(&path).is_some_and(|node| !node.is_hash()) {
1726            return Ok(())
1727        }
1728
1729        match node {
1730            TrieNode::EmptyRoot => {
1731                // For an empty root, ensure that we are at the root path, and at the upper subtrie.
1732                debug_assert!(path.is_empty());
1733                debug_assert!(self.path.is_empty());
1734                self.nodes.insert(path, SparseNode::Empty);
1735            }
1736            TrieNode::Branch(branch) => {
1737                // For a branch node, iterate over all potential children
1738                let mut stack_ptr = branch.as_ref().first_child_index();
1739                for idx in CHILD_INDEX_RANGE {
1740                    if branch.state_mask.is_bit_set(idx) {
1741                        let mut child_path = path;
1742                        child_path.push_unchecked(idx);
1743                        if Self::is_child_same_level(&path, &child_path) {
1744                            // Reveal each child node or hash it has, but only if the child is on
1745                            // the same level as the parent.
1746                            self.reveal_node_or_hash(child_path, &branch.stack[stack_ptr])?;
1747                        }
1748                        stack_ptr += 1;
1749                    }
1750                }
1751                // Update the branch node entry in the nodes map, handling cases where a blinded
1752                // node is now replaced with a revealed node.
1753                match self.nodes.entry(path) {
1754                    Entry::Occupied(mut entry) => match entry.get() {
1755                        // Replace a hash node with a fully revealed branch node.
1756                        SparseNode::Hash(hash) => {
1757                            entry.insert(SparseNode::Branch {
1758                                state_mask: branch.state_mask,
1759                                // Memoize the hash of a previously blinded node in a new branch
1760                                // node.
1761                                hash: Some(*hash),
1762                                store_in_db_trie: Some(
1763                                    masks.hash_mask.is_some_and(|mask| !mask.is_empty()) ||
1764                                        masks.tree_mask.is_some_and(|mask| !mask.is_empty()),
1765                                ),
1766                            });
1767                        }
1768                        // Branch node already exists, or an extension node was placed where a
1769                        // branch node was before.
1770                        SparseNode::Branch { .. } | SparseNode::Extension { .. } => {}
1771                        // All other node types can't be handled.
1772                        node @ (SparseNode::Empty | SparseNode::Leaf { .. }) => {
1773                            return Err(SparseTrieErrorKind::Reveal {
1774                                path: *entry.key(),
1775                                node: Box::new(node.clone()),
1776                            }
1777                            .into())
1778                        }
1779                    },
1780                    Entry::Vacant(entry) => {
1781                        entry.insert(SparseNode::new_branch(branch.state_mask));
1782                    }
1783                }
1784            }
1785            TrieNode::Extension(ext) => match self.nodes.entry(path) {
1786                Entry::Occupied(mut entry) => match entry.get() {
1787                    // Replace a hash node with a revealed extension node.
1788                    SparseNode::Hash(hash) => {
1789                        let mut child_path = *entry.key();
1790                        child_path.extend(&ext.key);
1791                        entry.insert(SparseNode::Extension {
1792                            key: ext.key,
1793                            // Memoize the hash of a previously blinded node in a new extension
1794                            // node.
1795                            hash: Some(*hash),
1796                            store_in_db_trie: None,
1797                        });
1798                        if Self::is_child_same_level(&path, &child_path) {
1799                            self.reveal_node_or_hash(child_path, &ext.child)?;
1800                        }
1801                    }
1802                    // Extension node already exists, or an extension node was placed where a branch
1803                    // node was before.
1804                    SparseNode::Extension { .. } | SparseNode::Branch { .. } => {}
1805                    // All other node types can't be handled.
1806                    node @ (SparseNode::Empty | SparseNode::Leaf { .. }) => {
1807                        return Err(SparseTrieErrorKind::Reveal {
1808                            path: *entry.key(),
1809                            node: Box::new(node.clone()),
1810                        }
1811                        .into())
1812                    }
1813                },
1814                Entry::Vacant(entry) => {
1815                    let mut child_path = *entry.key();
1816                    child_path.extend(&ext.key);
1817                    entry.insert(SparseNode::new_ext(ext.key));
1818                    if Self::is_child_same_level(&path, &child_path) {
1819                        self.reveal_node_or_hash(child_path, &ext.child)?;
1820                    }
1821                }
1822            },
1823            TrieNode::Leaf(leaf) => match self.nodes.entry(path) {
1824                Entry::Occupied(mut entry) => match entry.get() {
1825                    // Replace a hash node with a revealed leaf node and store leaf node value.
1826                    SparseNode::Hash(hash) => {
1827                        let mut full = *entry.key();
1828                        full.extend(&leaf.key);
1829                        self.inner.values.insert(full, leaf.value.clone());
1830                        entry.insert(SparseNode::Leaf {
1831                            key: leaf.key,
1832                            // Memoize the hash of a previously blinded node in a new leaf
1833                            // node.
1834                            hash: Some(*hash),
1835                        });
1836                    }
1837                    // Leaf node already exists.
1838                    SparseNode::Leaf { .. } => {}
1839                    // All other node types can't be handled.
1840                    node @ (SparseNode::Empty |
1841                    SparseNode::Extension { .. } |
1842                    SparseNode::Branch { .. }) => {
1843                        return Err(SparseTrieErrorKind::Reveal {
1844                            path: *entry.key(),
1845                            node: Box::new(node.clone()),
1846                        }
1847                        .into())
1848                    }
1849                },
1850                Entry::Vacant(entry) => {
1851                    let mut full = *entry.key();
1852                    full.extend(&leaf.key);
1853                    entry.insert(SparseNode::new_leaf(leaf.key));
1854                    self.inner.values.insert(full, leaf.value.clone());
1855                }
1856            },
1857        }
1858
1859        Ok(())
1860    }
1861
1862    /// Reveals either a node or its hash placeholder based on the provided child data.
1863    ///
1864    /// When traversing the trie, we often encounter references to child nodes that
1865    /// are either directly embedded or represented by their hash. This method
1866    /// handles both cases:
1867    ///
1868    /// 1. If the child data represents a hash (32+1=33 bytes), store it as a hash node
1869    /// 2. Otherwise, decode the data as a [`TrieNode`] and recursively reveal it using
1870    ///    `reveal_node`
1871    ///
1872    /// # Returns
1873    ///
1874    /// Returns `Ok(())` if successful, or an error if the node cannot be revealed.
1875    ///
1876    /// # Error Handling
1877    ///
1878    /// Will error if there's a conflict between a new hash node and an existing one
1879    /// at the same path
1880    fn reveal_node_or_hash(&mut self, path: Nibbles, child: &[u8]) -> SparseTrieResult<()> {
1881        if child.len() == B256::len_bytes() + 1 {
1882            let hash = B256::from_slice(&child[1..]);
1883            match self.nodes.entry(path) {
1884                Entry::Occupied(entry) => match entry.get() {
1885                    // Hash node with a different hash can't be handled.
1886                    SparseNode::Hash(previous_hash) if previous_hash != &hash => {
1887                        return Err(SparseTrieErrorKind::Reveal {
1888                            path: *entry.key(),
1889                            node: Box::new(SparseNode::Hash(hash)),
1890                        }
1891                        .into())
1892                    }
1893                    _ => {}
1894                },
1895                Entry::Vacant(entry) => {
1896                    entry.insert(SparseNode::Hash(hash));
1897                }
1898            }
1899            return Ok(())
1900        }
1901
1902        self.reveal_node(path, &TrieNode::decode(&mut &child[..])?, TrieMasks::none())
1903    }
1904
1905    /// Recalculates and updates the RLP hashes for the changed nodes in this subtrie.
1906    ///
1907    /// The function starts from the subtrie root, traverses down to leaves, and then calculates
1908    /// the hashes from leaves back up to the root. It uses a stack from [`SparseSubtrieBuffers`] to
1909    /// track the traversal and accumulate RLP encodings.
1910    ///
1911    /// # Parameters
1912    ///
1913    /// - `prefix_set`: The set of trie paths whose nodes have changed.
1914    /// - `update_actions`: A buffer which `SparseTrieUpdatesAction`s will be written to in the
1915    ///   event that any changes to the top-level updates are required. If None then update
1916    ///   retention is disabled.
1917    /// - `branch_node_tree_masks`: The tree masks for branch nodes
1918    /// - `branch_node_hash_masks`: The hash masks for branch nodes
1919    ///
1920    /// # Returns
1921    ///
1922    /// A tuple containing the root node of the updated subtrie.
1923    ///
1924    /// # Panics
1925    ///
1926    /// If the node at the root path does not exist.
1927    #[instrument(level = "trace", target = "trie::parallel_sparse", skip_all, fields(root = ?self.path), ret)]
1928    fn update_hashes(
1929        &mut self,
1930        prefix_set: &mut PrefixSet,
1931        update_actions: &mut Option<Vec<SparseTrieUpdatesAction>>,
1932        branch_node_tree_masks: &HashMap<Nibbles, TrieMask>,
1933        branch_node_hash_masks: &HashMap<Nibbles, TrieMask>,
1934    ) -> RlpNode {
1935        trace!(target: "trie::parallel_sparse", "Updating subtrie hashes");
1936
1937        debug_assert!(prefix_set.iter().all(|path| path.starts_with(&self.path)));
1938
1939        debug_assert!(self.inner.buffers.path_stack.is_empty());
1940        self.inner
1941            .buffers
1942            .path_stack
1943            .push(RlpNodePathStackItem { path: self.path, is_in_prefix_set: None });
1944
1945        while let Some(stack_item) = self.inner.buffers.path_stack.pop() {
1946            let path = stack_item.path;
1947            let node = self
1948                .nodes
1949                .get_mut(&path)
1950                .unwrap_or_else(|| panic!("node at path {path:?} does not exist"));
1951
1952            self.inner.rlp_node(
1953                prefix_set,
1954                update_actions,
1955                stack_item,
1956                node,
1957                branch_node_tree_masks,
1958                branch_node_hash_masks,
1959            );
1960        }
1961
1962        debug_assert_eq!(self.inner.buffers.rlp_node_stack.len(), 1);
1963        self.inner.buffers.rlp_node_stack.pop().unwrap().rlp_node
1964    }
1965
1966    /// Removes all nodes and values from the subtrie, resetting it to a blank state
1967    /// with only an empty root node. This is used when a storage root is deleted.
1968    fn wipe(&mut self) {
1969        self.nodes = HashMap::from_iter([(Nibbles::default(), SparseNode::Empty)]);
1970        self.inner.clear();
1971    }
1972
1973    /// Clears the subtrie, keeping the data structures allocated.
1974    pub(crate) fn clear(&mut self) {
1975        self.nodes.clear();
1976        self.inner.clear();
1977    }
1978}
1979
1980/// Helper type for [`SparseSubtrie`] to mutably access only a subset of fields from the original
1981/// struct.
1982#[derive(Clone, PartialEq, Eq, Debug, Default)]
1983struct SparseSubtrieInner {
1984    /// Map from leaf key paths to their values.
1985    /// All values are stored here instead of directly in leaf nodes.
1986    values: HashMap<Nibbles, Vec<u8>>,
1987    /// Reusable buffers for [`SparseSubtrie::update_hashes`].
1988    buffers: SparseSubtrieBuffers,
1989}
1990
1991impl SparseSubtrieInner {
1992    /// Computes the RLP encoding and its hash for a single (trie node)[`SparseNode`].
1993    ///
1994    /// # Deferred Processing
1995    ///
1996    /// When an extension or a branch node depends on child nodes that haven't been computed yet,
1997    /// the function pushes the current node back onto the path stack along with its children,
1998    /// then returns early. This allows the iterative algorithm to process children first before
1999    /// retrying the parent.
2000    ///
2001    /// # Parameters
2002    ///
2003    /// - `prefix_set`: Set of prefixes (key paths) that have been marked as updated
2004    /// - `update_actions`: A buffer which `SparseTrieUpdatesAction`s will be written to in the
2005    ///   event that any changes to the top-level updates are required. If None then update
2006    ///   retention is disabled.
2007    /// - `stack_item`: The stack item to process
2008    /// - `node`: The sparse node to process (will be mutated to update hash)
2009    /// - `branch_node_tree_masks`: The tree masks for branch nodes
2010    /// - `branch_node_hash_masks`: The hash masks for branch nodes
2011    ///
2012    /// # Side Effects
2013    ///
2014    /// - Updates the node's hash field after computing RLP
2015    /// - Pushes nodes to [`SparseSubtrieBuffers::path_stack`] to manage traversal
2016    /// - May push items onto the path stack for deferred processing
2017    ///
2018    /// # Exit condition
2019    ///
2020    /// Once all nodes have been processed and all RLPs and hashes calculated, pushes the root node
2021    /// onto the [`SparseSubtrieBuffers::rlp_node_stack`] and exits.
2022    fn rlp_node(
2023        &mut self,
2024        prefix_set: &mut PrefixSet,
2025        update_actions: &mut Option<Vec<SparseTrieUpdatesAction>>,
2026        mut stack_item: RlpNodePathStackItem,
2027        node: &mut SparseNode,
2028        branch_node_tree_masks: &HashMap<Nibbles, TrieMask>,
2029        branch_node_hash_masks: &HashMap<Nibbles, TrieMask>,
2030    ) {
2031        let path = stack_item.path;
2032        trace!(
2033            target: "trie::parallel_sparse",
2034            ?path,
2035            ?node,
2036            "Calculating node RLP"
2037        );
2038
2039        // Check if the path is in the prefix set.
2040        // First, check the cached value. If it's `None`, then check the prefix set, and update
2041        // the cached value.
2042        let mut prefix_set_contains = |path: &Nibbles| {
2043            *stack_item.is_in_prefix_set.get_or_insert_with(|| prefix_set.contains(path))
2044        };
2045
2046        let (rlp_node, node_type) = match node {
2047            SparseNode::Empty => (RlpNode::word_rlp(&EMPTY_ROOT_HASH), SparseNodeType::Empty),
2048            SparseNode::Hash(hash) => {
2049                // Return pre-computed hash of a blinded node immediately
2050                (RlpNode::word_rlp(hash), SparseNodeType::Hash)
2051            }
2052            SparseNode::Leaf { key, hash } => {
2053                let mut path = path;
2054                path.extend(key);
2055                let value = self.values.get(&path);
2056                if let Some(hash) = hash.filter(|_| !prefix_set_contains(&path) || value.is_none())
2057                {
2058                    // If the node hash is already computed, and either the node path is not in
2059                    // the prefix set or the leaf doesn't belong to the current trie (its value is
2060                    // absent), return the pre-computed hash
2061                    (RlpNode::word_rlp(&hash), SparseNodeType::Leaf)
2062                } else {
2063                    // Encode the leaf node and update its hash
2064                    let value = self.values.get(&path).unwrap();
2065                    self.buffers.rlp_buf.clear();
2066                    let rlp_node = LeafNodeRef { key, value }.rlp(&mut self.buffers.rlp_buf);
2067                    *hash = rlp_node.as_hash();
2068                    (rlp_node, SparseNodeType::Leaf)
2069                }
2070            }
2071            SparseNode::Extension { key, hash, store_in_db_trie } => {
2072                let mut child_path = path;
2073                child_path.extend(key);
2074                if let Some((hash, store_in_db_trie)) =
2075                    hash.zip(*store_in_db_trie).filter(|_| !prefix_set_contains(&path))
2076                {
2077                    // If the node hash is already computed, and the node path is not in
2078                    // the prefix set, return the pre-computed hash
2079                    (
2080                        RlpNode::word_rlp(&hash),
2081                        SparseNodeType::Extension { store_in_db_trie: Some(store_in_db_trie) },
2082                    )
2083                } else if self.buffers.rlp_node_stack.last().is_some_and(|e| e.path == child_path) {
2084                    // Top of the stack has the child node, we can encode the extension node and
2085                    // update its hash
2086                    let RlpNodeStackItem { path: _, rlp_node: child, node_type: child_node_type } =
2087                        self.buffers.rlp_node_stack.pop().unwrap();
2088                    self.buffers.rlp_buf.clear();
2089                    let rlp_node =
2090                        ExtensionNodeRef::new(key, &child).rlp(&mut self.buffers.rlp_buf);
2091                    *hash = rlp_node.as_hash();
2092
2093                    let store_in_db_trie_value = child_node_type.store_in_db_trie();
2094
2095                    trace!(
2096                        target: "trie::parallel_sparse",
2097                        ?path,
2098                        ?child_path,
2099                        ?child_node_type,
2100                        "Extension node"
2101                    );
2102
2103                    *store_in_db_trie = store_in_db_trie_value;
2104
2105                    (
2106                        rlp_node,
2107                        SparseNodeType::Extension {
2108                            // Inherit the `store_in_db_trie` flag from the child node, which is
2109                            // always the branch node
2110                            store_in_db_trie: store_in_db_trie_value,
2111                        },
2112                    )
2113                } else {
2114                    // Need to defer processing until child is computed, on the next
2115                    // invocation update the node's hash.
2116                    self.buffers.path_stack.extend([
2117                        RlpNodePathStackItem {
2118                            path,
2119                            is_in_prefix_set: Some(prefix_set_contains(&path)),
2120                        },
2121                        RlpNodePathStackItem { path: child_path, is_in_prefix_set: None },
2122                    ]);
2123                    return
2124                }
2125            }
2126            SparseNode::Branch { state_mask, hash, store_in_db_trie } => {
2127                if let Some((hash, store_in_db_trie)) =
2128                    hash.zip(*store_in_db_trie).filter(|_| !prefix_set_contains(&path))
2129                {
2130                    // If the node hash is already computed, and the node path is not in
2131                    // the prefix set, return the pre-computed hash
2132                    self.buffers.rlp_node_stack.push(RlpNodeStackItem {
2133                        path,
2134                        rlp_node: RlpNode::word_rlp(&hash),
2135                        node_type: SparseNodeType::Branch {
2136                            store_in_db_trie: Some(store_in_db_trie),
2137                        },
2138                    });
2139                    return
2140                }
2141
2142                let retain_updates = update_actions.is_some() && prefix_set_contains(&path);
2143
2144                self.buffers.branch_child_buf.clear();
2145                // Walk children in a reverse order from `f` to `0`, so we pop the `0` first
2146                // from the stack and keep walking in the sorted order.
2147                for bit in CHILD_INDEX_RANGE.rev() {
2148                    if state_mask.is_bit_set(bit) {
2149                        let mut child = path;
2150                        child.push_unchecked(bit);
2151                        self.buffers.branch_child_buf.push(child);
2152                    }
2153                }
2154
2155                self.buffers
2156                    .branch_value_stack_buf
2157                    .resize(self.buffers.branch_child_buf.len(), Default::default());
2158                let mut added_children = false;
2159
2160                let mut tree_mask = TrieMask::default();
2161                let mut hash_mask = TrieMask::default();
2162                let mut hashes = Vec::new();
2163                for (i, child_path) in self.buffers.branch_child_buf.iter().enumerate() {
2164                    if self.buffers.rlp_node_stack.last().is_some_and(|e| &e.path == child_path) {
2165                        let RlpNodeStackItem {
2166                            path: _,
2167                            rlp_node: child,
2168                            node_type: child_node_type,
2169                        } = self.buffers.rlp_node_stack.pop().unwrap();
2170
2171                        // Update the masks only if we need to retain trie updates
2172                        if retain_updates {
2173                            // SAFETY: it's a child, so it's never empty
2174                            let last_child_nibble = child_path.last().unwrap();
2175
2176                            // Determine whether we need to set trie mask bit.
2177                            let should_set_tree_mask_bit = if let Some(store_in_db_trie) =
2178                                child_node_type.store_in_db_trie()
2179                            {
2180                                // A branch or an extension node explicitly set the
2181                                // `store_in_db_trie` flag
2182                                store_in_db_trie
2183                            } else {
2184                                // A blinded node has the tree mask bit set
2185                                child_node_type.is_hash() &&
2186                                    branch_node_tree_masks
2187                                        .get(&path)
2188                                        .is_some_and(|mask| mask.is_bit_set(last_child_nibble))
2189                            };
2190                            if should_set_tree_mask_bit {
2191                                tree_mask.set_bit(last_child_nibble);
2192                            }
2193
2194                            // Set the hash mask. If a child node is a revealed branch node OR
2195                            // is a blinded node that has its hash mask bit set according to the
2196                            // database, set the hash mask bit and save the hash.
2197                            let hash = child.as_hash().filter(|_| {
2198                                child_node_type.is_branch() ||
2199                                    (child_node_type.is_hash() &&
2200                                        branch_node_hash_masks.get(&path).is_some_and(
2201                                            |mask| mask.is_bit_set(last_child_nibble),
2202                                        ))
2203                            });
2204                            if let Some(hash) = hash {
2205                                hash_mask.set_bit(last_child_nibble);
2206                                hashes.push(hash);
2207                            }
2208                        }
2209
2210                        // Insert children in the resulting buffer in a normal order,
2211                        // because initially we iterated in reverse.
2212                        // SAFETY: i < len and len is never 0
2213                        let original_idx = self.buffers.branch_child_buf.len() - i - 1;
2214                        self.buffers.branch_value_stack_buf[original_idx] = child;
2215                        added_children = true;
2216                    } else {
2217                        // Need to defer processing until children are computed, on the next
2218                        // invocation update the node's hash.
2219                        debug_assert!(!added_children);
2220                        self.buffers.path_stack.push(RlpNodePathStackItem {
2221                            path,
2222                            is_in_prefix_set: Some(prefix_set_contains(&path)),
2223                        });
2224                        self.buffers.path_stack.extend(
2225                            self.buffers
2226                                .branch_child_buf
2227                                .drain(..)
2228                                .map(|path| RlpNodePathStackItem { path, is_in_prefix_set: None }),
2229                        );
2230                        return
2231                    }
2232                }
2233
2234                trace!(
2235                    target: "trie::parallel_sparse",
2236                    ?path,
2237                    ?tree_mask,
2238                    ?hash_mask,
2239                    "Branch node masks"
2240                );
2241
2242                // Top of the stack has all children node, we can encode the branch node and
2243                // update its hash
2244                self.buffers.rlp_buf.clear();
2245                let branch_node_ref =
2246                    BranchNodeRef::new(&self.buffers.branch_value_stack_buf, *state_mask);
2247                let rlp_node = branch_node_ref.rlp(&mut self.buffers.rlp_buf);
2248                *hash = rlp_node.as_hash();
2249
2250                // Save a branch node update only if it's not a root node, and we need to
2251                // persist updates.
2252                let store_in_db_trie_value = if let Some(update_actions) =
2253                    update_actions.as_mut().filter(|_| retain_updates && !path.is_empty())
2254                {
2255                    let store_in_db_trie = !tree_mask.is_empty() || !hash_mask.is_empty();
2256                    if store_in_db_trie {
2257                        // Store in DB trie if there are either any children that are stored in
2258                        // the DB trie, or any children represent hashed values
2259                        hashes.reverse();
2260                        let branch_node = BranchNodeCompact::new(
2261                            *state_mask,
2262                            tree_mask,
2263                            hash_mask,
2264                            hashes,
2265                            hash.filter(|_| path.is_empty()),
2266                        );
2267                        update_actions
2268                            .push(SparseTrieUpdatesAction::InsertUpdated(path, branch_node));
2269                    } else if branch_node_tree_masks.get(&path).is_some_and(|mask| !mask.is_empty()) ||
2270                        branch_node_hash_masks.get(&path).is_some_and(|mask| !mask.is_empty())
2271                    {
2272                        // If new tree and hash masks are empty, but previously they weren't, we
2273                        // need to remove the node update and add the node itself to the list of
2274                        // removed nodes.
2275                        update_actions.push(SparseTrieUpdatesAction::InsertRemoved(path));
2276                    } else if branch_node_tree_masks.get(&path).is_none_or(|mask| mask.is_empty()) &&
2277                        branch_node_hash_masks.get(&path).is_none_or(|mask| mask.is_empty())
2278                    {
2279                        // If new tree and hash masks are empty, and they were previously empty
2280                        // as well, we need to remove the node update.
2281                        update_actions.push(SparseTrieUpdatesAction::RemoveUpdated(path));
2282                    }
2283
2284                    store_in_db_trie
2285                } else {
2286                    false
2287                };
2288                *store_in_db_trie = Some(store_in_db_trie_value);
2289
2290                (
2291                    rlp_node,
2292                    SparseNodeType::Branch { store_in_db_trie: Some(store_in_db_trie_value) },
2293                )
2294            }
2295        };
2296
2297        self.buffers.rlp_node_stack.push(RlpNodeStackItem { path, rlp_node, node_type });
2298        trace!(
2299            target: "trie::parallel_sparse",
2300            ?path,
2301            ?node_type,
2302            "Added node to RLP node stack"
2303        );
2304    }
2305
2306    /// Clears the subtrie, keeping the data structures allocated.
2307    fn clear(&mut self) {
2308        self.values.clear();
2309        self.buffers.clear();
2310    }
2311}
2312
2313/// Represents the outcome of processing a node during leaf insertion
2314#[derive(Clone, Debug, PartialEq, Eq, Default)]
2315pub enum LeafUpdateStep {
2316    /// Continue traversing to the next node
2317    Continue {
2318        /// The next node path to process
2319        next_node: Nibbles,
2320    },
2321    /// Update is complete with nodes inserted
2322    Complete {
2323        /// The node paths that were inserted during this step
2324        inserted_nodes: Vec<Nibbles>,
2325        /// Path to a node which may need to be revealed
2326        reveal_path: Option<Nibbles>,
2327    },
2328    /// The node was not found
2329    #[default]
2330    NodeNotFound,
2331}
2332
2333impl LeafUpdateStep {
2334    /// Creates a step to continue with the next node
2335    pub const fn continue_with(next_node: Nibbles) -> Self {
2336        Self::Continue { next_node }
2337    }
2338
2339    /// Creates a step indicating completion with inserted nodes
2340    pub const fn complete_with_insertions(
2341        inserted_nodes: Vec<Nibbles>,
2342        reveal_path: Option<Nibbles>,
2343    ) -> Self {
2344        Self::Complete { inserted_nodes, reveal_path }
2345    }
2346}
2347
2348/// Sparse Subtrie Type.
2349///
2350/// Used to determine the type of subtrie a certain path belongs to:
2351/// - Paths in the range `0x..=0xf` belong to the upper subtrie.
2352/// - Paths in the range `0x00..` belong to one of the lower subtries. The index of the lower
2353///   subtrie is determined by the first [`UPPER_TRIE_MAX_DEPTH`] nibbles of the path.
2354///
2355/// There can be at most [`NUM_LOWER_SUBTRIES`] lower subtries.
2356#[derive(Clone, Copy, PartialEq, Eq, Debug)]
2357pub enum SparseSubtrieType {
2358    /// Upper subtrie with paths in the range `0x..=0xf`
2359    Upper,
2360    /// Lower subtrie with paths in the range `0x00..`. Includes the index of the subtrie,
2361    /// according to the path prefix.
2362    Lower(usize),
2363}
2364
2365impl SparseSubtrieType {
2366    /// Returns true if a node at a path of the given length would be placed in the upper subtrie.
2367    ///
2368    /// Nodes with paths shorter than [`UPPER_TRIE_MAX_DEPTH`] nibbles belong to the upper subtrie,
2369    /// while longer paths belong to the lower subtries.
2370    pub const fn path_len_is_upper(len: usize) -> bool {
2371        len < UPPER_TRIE_MAX_DEPTH
2372    }
2373
2374    /// Returns the type of subtrie based on the given path.
2375    pub fn from_path(path: &Nibbles) -> Self {
2376        if Self::path_len_is_upper(path.len()) {
2377            Self::Upper
2378        } else {
2379            Self::Lower(path_subtrie_index_unchecked(path))
2380        }
2381    }
2382
2383    /// Returns the index of the lower subtrie, if it exists.
2384    pub const fn lower_index(&self) -> Option<usize> {
2385        match self {
2386            Self::Upper => None,
2387            Self::Lower(index) => Some(*index),
2388        }
2389    }
2390}
2391
2392impl Ord for SparseSubtrieType {
2393    /// Orders two [`SparseSubtrieType`]s such that `Upper` is less than `Lower(_)`, and `Lower`s
2394    /// are ordered by their index.
2395    fn cmp(&self, other: &Self) -> Ordering {
2396        match (self, other) {
2397            (Self::Upper, Self::Upper) => Ordering::Equal,
2398            (Self::Upper, Self::Lower(_)) => Ordering::Less,
2399            (Self::Lower(_), Self::Upper) => Ordering::Greater,
2400            (Self::Lower(idx_a), Self::Lower(idx_b)) if idx_a == idx_b => Ordering::Equal,
2401            (Self::Lower(idx_a), Self::Lower(idx_b)) => idx_a.cmp(idx_b),
2402        }
2403    }
2404}
2405
2406impl PartialOrd for SparseSubtrieType {
2407    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
2408        Some(self.cmp(other))
2409    }
2410}
2411
2412/// Collection of reusable buffers for calculating subtrie hashes.
2413///
2414/// These buffers reduce allocations when computing RLP representations during trie updates.
2415#[derive(Clone, PartialEq, Eq, Debug, Default)]
2416pub struct SparseSubtrieBuffers {
2417    /// Stack of RLP node paths
2418    path_stack: Vec<RlpNodePathStackItem>,
2419    /// Stack of RLP nodes
2420    rlp_node_stack: Vec<RlpNodeStackItem>,
2421    /// Reusable branch child path
2422    branch_child_buf: SmallVec<[Nibbles; 16]>,
2423    /// Reusable branch value stack
2424    branch_value_stack_buf: SmallVec<[RlpNode; 16]>,
2425    /// Reusable RLP buffer
2426    rlp_buf: Vec<u8>,
2427}
2428
2429impl SparseSubtrieBuffers {
2430    /// Clears all buffers.
2431    fn clear(&mut self) {
2432        self.path_stack.clear();
2433        self.rlp_node_stack.clear();
2434        self.branch_child_buf.clear();
2435        self.branch_value_stack_buf.clear();
2436        self.rlp_buf.clear();
2437    }
2438}
2439
2440/// RLP node path stack item.
2441#[derive(Clone, PartialEq, Eq, Debug)]
2442pub struct RlpNodePathStackItem {
2443    /// Path to the node.
2444    pub path: Nibbles,
2445    /// Whether the path is in the prefix set. If [`None`], then unknown yet.
2446    pub is_in_prefix_set: Option<bool>,
2447}
2448
2449/// Changed subtrie.
2450#[derive(Debug)]
2451struct ChangedSubtrie {
2452    /// Lower subtrie index in the range [0, [`NUM_LOWER_SUBTRIES`]).
2453    index: usize,
2454    /// Changed subtrie
2455    subtrie: Box<SparseSubtrie>,
2456    /// Prefix set of keys that belong to the subtrie.
2457    prefix_set: PrefixSet,
2458    /// Reusable buffer for collecting [`SparseTrieUpdatesAction`]s during computations. Will be
2459    /// None if update retention is disabled.
2460    update_actions_buf: Option<Vec<SparseTrieUpdatesAction>>,
2461}
2462
2463/// Convert first [`UPPER_TRIE_MAX_DEPTH`] nibbles of the path into a lower subtrie index in the
2464/// range [0, [`NUM_LOWER_SUBTRIES`]).
2465///
2466/// # Panics
2467///
2468/// If the path is shorter than [`UPPER_TRIE_MAX_DEPTH`] nibbles.
2469fn path_subtrie_index_unchecked(path: &Nibbles) -> usize {
2470    debug_assert_eq!(UPPER_TRIE_MAX_DEPTH, 2);
2471    path.get_byte_unchecked(0) as usize
2472}
2473
2474/// Used by lower subtries to communicate updates to the top-level [`SparseTrieUpdates`] set.
2475#[derive(Clone, Debug, Eq, PartialEq)]
2476enum SparseTrieUpdatesAction {
2477    /// Remove the path from the `updated_nodes`, if it was present, and add it to `removed_nodes`.
2478    InsertRemoved(Nibbles),
2479    /// Remove the path from the `updated_nodes`, if it was present, leaving `removed_nodes`
2480    /// unaffected.
2481    RemoveUpdated(Nibbles),
2482    /// Insert the branch node into `updated_nodes`.
2483    InsertUpdated(Nibbles, BranchNodeCompact),
2484}
2485
2486#[cfg(test)]
2487mod tests {
2488    use super::{
2489        path_subtrie_index_unchecked, LowerSparseSubtrie, ParallelSparseTrie, SparseSubtrie,
2490        SparseSubtrieType,
2491    };
2492    use crate::trie::ChangedSubtrie;
2493    use alloy_primitives::{
2494        b256, hex,
2495        map::{B256Set, DefaultHashBuilder, HashMap},
2496        B256, U256,
2497    };
2498    use alloy_rlp::{Decodable, Encodable};
2499    use alloy_trie::{BranchNodeCompact, Nibbles};
2500    use assert_matches::assert_matches;
2501    use itertools::Itertools;
2502    use proptest::{prelude::*, sample::SizeRange};
2503    use proptest_arbitrary_interop::arb;
2504    use reth_execution_errors::{SparseTrieError, SparseTrieErrorKind};
2505    use reth_primitives_traits::Account;
2506    use reth_provider::{test_utils::create_test_provider_factory, TrieWriter};
2507    use reth_trie::{
2508        hashed_cursor::{noop::NoopHashedAccountCursor, HashedPostStateAccountCursor},
2509        node_iter::{TrieElement, TrieNodeIter},
2510        trie_cursor::{noop::NoopAccountTrieCursor, TrieCursor, TrieCursorFactory},
2511        walker::TrieWalker,
2512        HashedPostState,
2513    };
2514    use reth_trie_common::{
2515        prefix_set::PrefixSetMut,
2516        proof::{ProofNodes, ProofRetainer},
2517        updates::TrieUpdates,
2518        BranchNode, ExtensionNode, HashBuilder, LeafNode, RlpNode, TrieMask, TrieNode,
2519        EMPTY_ROOT_HASH,
2520    };
2521    use reth_trie_db::DatabaseTrieCursorFactory;
2522    use reth_trie_sparse::{
2523        provider::{DefaultTrieNodeProvider, RevealedNode, TrieNodeProvider},
2524        LeafLookup, LeafLookupError, RevealedSparseNode, SerialSparseTrie, SparseNode,
2525        SparseTrieInterface, SparseTrieUpdates, TrieMasks,
2526    };
2527    use std::collections::{BTreeMap, BTreeSet};
2528
2529    /// Pad nibbles to the length of a B256 hash with zeros on the right.
2530    fn pad_nibbles_right(mut nibbles: Nibbles) -> Nibbles {
2531        nibbles.extend(&Nibbles::from_nibbles_unchecked(vec![
2532            0;
2533            B256::len_bytes() * 2 - nibbles.len()
2534        ]));
2535        nibbles
2536    }
2537
2538    /// Mock trie node provider for testing that allows pre-setting nodes at specific paths.
2539    ///
2540    /// This provider can be used in tests to simulate trie nodes that need to be revealed
2541    /// during trie operations, particularly when collapsing branch nodes during leaf removal.
2542    #[derive(Debug, Clone)]
2543    struct MockTrieNodeProvider {
2544        /// Mapping from path to revealed node data
2545        nodes: HashMap<Nibbles, RevealedNode, DefaultHashBuilder>,
2546    }
2547
2548    impl MockTrieNodeProvider {
2549        /// Creates a new empty mock provider
2550        fn new() -> Self {
2551            Self { nodes: HashMap::default() }
2552        }
2553
2554        /// Adds a revealed node at the specified path
2555        fn add_revealed_node(&mut self, path: Nibbles, node: RevealedNode) {
2556            self.nodes.insert(path, node);
2557        }
2558    }
2559
2560    impl TrieNodeProvider for MockTrieNodeProvider {
2561        fn trie_node(&self, path: &Nibbles) -> Result<Option<RevealedNode>, SparseTrieError> {
2562            Ok(self.nodes.get(path).cloned())
2563        }
2564    }
2565
2566    fn create_account(nonce: u64) -> Account {
2567        Account { nonce, ..Default::default() }
2568    }
2569
2570    fn encode_account_value(nonce: u64) -> Vec<u8> {
2571        let account = Account { nonce, ..Default::default() };
2572        let trie_account = account.into_trie_account(EMPTY_ROOT_HASH);
2573        let mut buf = Vec::new();
2574        trie_account.encode(&mut buf);
2575        buf
2576    }
2577
2578    /// Test context that provides helper methods for trie testing
2579    #[derive(Default)]
2580    struct ParallelSparseTrieTestContext;
2581
2582    impl ParallelSparseTrieTestContext {
2583        /// Assert that a lower subtrie exists at the given path
2584        fn assert_subtrie_exists(&self, trie: &ParallelSparseTrie, path: &Nibbles) {
2585            let idx = path_subtrie_index_unchecked(path);
2586            assert!(
2587                trie.lower_subtries[idx].as_revealed_ref().is_some(),
2588                "Expected lower subtrie at path {path:?} to exist",
2589            );
2590        }
2591
2592        /// Get a lower subtrie, panicking if it doesn't exist
2593        fn get_subtrie<'a>(
2594            &self,
2595            trie: &'a ParallelSparseTrie,
2596            path: &Nibbles,
2597        ) -> &'a SparseSubtrie {
2598            let idx = path_subtrie_index_unchecked(path);
2599            trie.lower_subtries[idx]
2600                .as_revealed_ref()
2601                .unwrap_or_else(|| panic!("Lower subtrie at path {path:?} should exist"))
2602        }
2603
2604        /// Assert that a lower subtrie has a specific path field value
2605        fn assert_subtrie_path(
2606            &self,
2607            trie: &ParallelSparseTrie,
2608            subtrie_prefix: impl AsRef<[u8]>,
2609            expected_path: impl AsRef<[u8]>,
2610        ) {
2611            let subtrie_prefix = Nibbles::from_nibbles(subtrie_prefix);
2612            let expected_path = Nibbles::from_nibbles(expected_path);
2613            let idx = path_subtrie_index_unchecked(&subtrie_prefix);
2614
2615            let subtrie = trie.lower_subtries[idx].as_revealed_ref().unwrap_or_else(|| {
2616                panic!("Lower subtrie at prefix {subtrie_prefix:?} should exist")
2617            });
2618
2619            assert_eq!(
2620                subtrie.path, expected_path,
2621                "Subtrie at prefix {subtrie_prefix:?} should have path {expected_path:?}, but has {:?}",
2622                subtrie.path
2623            );
2624        }
2625
2626        /// Create test leaves with consecutive account values
2627        fn create_test_leaves(&self, paths: &[&[u8]]) -> Vec<(Nibbles, Vec<u8>)> {
2628            paths
2629                .iter()
2630                .enumerate()
2631                .map(|(i, path)| (Nibbles::from_nibbles(path), encode_account_value(i as u64 + 1)))
2632                .collect()
2633        }
2634
2635        /// Create a single test leaf with the given path and value nonce
2636        fn create_test_leaf(&self, path: impl AsRef<[u8]>, value_nonce: u64) -> (Nibbles, Vec<u8>) {
2637            (Nibbles::from_nibbles(path), encode_account_value(value_nonce))
2638        }
2639
2640        /// Update multiple leaves in the trie
2641        fn update_leaves(
2642            &self,
2643            trie: &mut ParallelSparseTrie,
2644            leaves: impl IntoIterator<Item = (Nibbles, Vec<u8>)>,
2645        ) {
2646            for (path, value) in leaves {
2647                trie.update_leaf(path, value, DefaultTrieNodeProvider).unwrap();
2648            }
2649        }
2650
2651        /// Create an assertion builder for a subtrie
2652        fn assert_subtrie<'a>(
2653            &self,
2654            trie: &'a ParallelSparseTrie,
2655            path: Nibbles,
2656        ) -> SubtrieAssertion<'a> {
2657            self.assert_subtrie_exists(trie, &path);
2658            let subtrie = self.get_subtrie(trie, &path);
2659            SubtrieAssertion::new(subtrie)
2660        }
2661
2662        /// Create an assertion builder for the upper subtrie
2663        fn assert_upper_subtrie<'a>(&self, trie: &'a ParallelSparseTrie) -> SubtrieAssertion<'a> {
2664            SubtrieAssertion::new(&trie.upper_subtrie)
2665        }
2666
2667        /// Assert the root, trie updates, and nodes against the hash builder output.
2668        fn assert_with_hash_builder(
2669            &self,
2670            trie: &mut ParallelSparseTrie,
2671            hash_builder_root: B256,
2672            hash_builder_updates: TrieUpdates,
2673            hash_builder_proof_nodes: ProofNodes,
2674        ) {
2675            assert_eq!(trie.root(), hash_builder_root);
2676            pretty_assertions::assert_eq!(
2677                BTreeMap::from_iter(trie.updates_ref().updated_nodes.clone()),
2678                BTreeMap::from_iter(hash_builder_updates.account_nodes)
2679            );
2680            assert_eq_parallel_sparse_trie_proof_nodes(trie, hash_builder_proof_nodes);
2681        }
2682    }
2683
2684    /// Assertion builder for subtrie structure
2685    struct SubtrieAssertion<'a> {
2686        subtrie: &'a SparseSubtrie,
2687    }
2688
2689    impl<'a> SubtrieAssertion<'a> {
2690        fn new(subtrie: &'a SparseSubtrie) -> Self {
2691            Self { subtrie }
2692        }
2693
2694        fn has_branch(self, path: &Nibbles, expected_mask_bits: &[u8]) -> Self {
2695            match self.subtrie.nodes.get(path) {
2696                Some(SparseNode::Branch { state_mask, .. }) => {
2697                    for bit in expected_mask_bits {
2698                        assert!(
2699                            state_mask.is_bit_set(*bit),
2700                            "Expected branch at {path:?} to have bit {bit} set, instead mask is: {state_mask:?}",
2701                        );
2702                    }
2703                }
2704                node => panic!("Expected branch node at {path:?}, found {node:?}"),
2705            }
2706            self
2707        }
2708
2709        fn has_leaf(self, path: &Nibbles, expected_key: &Nibbles) -> Self {
2710            match self.subtrie.nodes.get(path) {
2711                Some(SparseNode::Leaf { key, .. }) => {
2712                    assert_eq!(
2713                        *key, *expected_key,
2714                        "Expected leaf at {path:?} to have key {expected_key:?}, found {key:?}",
2715                    );
2716                }
2717                node => panic!("Expected leaf node at {path:?}, found {node:?}"),
2718            }
2719            self
2720        }
2721
2722        fn has_extension(self, path: &Nibbles, expected_key: &Nibbles) -> Self {
2723            match self.subtrie.nodes.get(path) {
2724                Some(SparseNode::Extension { key, .. }) => {
2725                    assert_eq!(
2726                        *key, *expected_key,
2727                        "Expected extension at {path:?} to have key {expected_key:?}, found {key:?}",
2728                    );
2729                }
2730                node => panic!("Expected extension node at {path:?}, found {node:?}"),
2731            }
2732            self
2733        }
2734
2735        fn has_hash(self, path: &Nibbles, expected_hash: &B256) -> Self {
2736            match self.subtrie.nodes.get(path) {
2737                Some(SparseNode::Hash(hash)) => {
2738                    assert_eq!(
2739                        *hash, *expected_hash,
2740                        "Expected hash at {path:?} to be {expected_hash:?}, found {hash:?}",
2741                    );
2742                }
2743                node => panic!("Expected hash node at {path:?}, found {node:?}"),
2744            }
2745            self
2746        }
2747
2748        fn has_value(self, path: &Nibbles, expected_value: &[u8]) -> Self {
2749            let actual = self.subtrie.inner.values.get(path);
2750            assert_eq!(
2751                actual.map(|v| v.as_slice()),
2752                Some(expected_value),
2753                "Expected value at {path:?} to be {expected_value:?}, found {actual:?}",
2754            );
2755            self
2756        }
2757
2758        fn has_no_value(self, path: &Nibbles) -> Self {
2759            let actual = self.subtrie.inner.values.get(path);
2760            assert!(actual.is_none(), "Expected no value at {path:?}, but found {actual:?}");
2761            self
2762        }
2763    }
2764
2765    fn create_leaf_node(key: impl AsRef<[u8]>, value_nonce: u64) -> TrieNode {
2766        TrieNode::Leaf(LeafNode::new(Nibbles::from_nibbles(key), encode_account_value(value_nonce)))
2767    }
2768
2769    fn create_extension_node(key: impl AsRef<[u8]>, child_hash: B256) -> TrieNode {
2770        TrieNode::Extension(ExtensionNode::new(
2771            Nibbles::from_nibbles(key),
2772            RlpNode::word_rlp(&child_hash),
2773        ))
2774    }
2775
2776    fn create_branch_node_with_children(
2777        children_indices: &[u8],
2778        child_hashes: impl IntoIterator<Item = RlpNode>,
2779    ) -> TrieNode {
2780        let mut stack = Vec::new();
2781        let mut state_mask = TrieMask::default();
2782
2783        for (&idx, hash) in children_indices.iter().zip(child_hashes.into_iter()) {
2784            state_mask.set_bit(idx);
2785            stack.push(hash);
2786        }
2787
2788        TrieNode::Branch(BranchNode::new(stack, state_mask))
2789    }
2790
2791    /// Calculate the state root by feeding the provided state to the hash builder and retaining the
2792    /// proofs for the provided targets.
2793    ///
2794    /// Returns the state root and the retained proof nodes.
2795    fn run_hash_builder(
2796        state: impl IntoIterator<Item = (Nibbles, Account)> + Clone,
2797        trie_cursor: impl TrieCursor,
2798        destroyed_accounts: B256Set,
2799        proof_targets: impl IntoIterator<Item = Nibbles>,
2800    ) -> (B256, TrieUpdates, ProofNodes, HashMap<Nibbles, TrieMask>, HashMap<Nibbles, TrieMask>)
2801    {
2802        let mut account_rlp = Vec::new();
2803
2804        let mut hash_builder = HashBuilder::default()
2805            .with_updates(true)
2806            .with_proof_retainer(ProofRetainer::from_iter(proof_targets));
2807
2808        let mut prefix_set = PrefixSetMut::default();
2809        prefix_set.extend_keys(state.clone().into_iter().map(|(nibbles, _)| nibbles));
2810        prefix_set.extend_keys(destroyed_accounts.iter().map(Nibbles::unpack));
2811        let walker = TrieWalker::<_>::state_trie(trie_cursor, prefix_set.freeze())
2812            .with_deletions_retained(true);
2813        let hashed_post_state = HashedPostState::default()
2814            .with_accounts(state.into_iter().map(|(nibbles, account)| {
2815                (nibbles.pack().into_inner().unwrap().into(), Some(account))
2816            }))
2817            .into_sorted();
2818        let mut node_iter = TrieNodeIter::state_trie(
2819            walker,
2820            HashedPostStateAccountCursor::new(
2821                NoopHashedAccountCursor::default(),
2822                hashed_post_state.accounts(),
2823            ),
2824        );
2825
2826        while let Some(node) = node_iter.try_next().unwrap() {
2827            match node {
2828                TrieElement::Branch(branch) => {
2829                    hash_builder.add_branch(branch.key, branch.value, branch.children_are_in_trie);
2830                }
2831                TrieElement::Leaf(key, account) => {
2832                    let account = account.into_trie_account(EMPTY_ROOT_HASH);
2833                    account.encode(&mut account_rlp);
2834
2835                    hash_builder.add_leaf(Nibbles::unpack(key), &account_rlp);
2836                    account_rlp.clear();
2837                }
2838            }
2839        }
2840        let root = hash_builder.root();
2841        let proof_nodes = hash_builder.take_proof_nodes();
2842        let branch_node_hash_masks = hash_builder
2843            .updated_branch_nodes
2844            .clone()
2845            .unwrap_or_default()
2846            .iter()
2847            .map(|(path, node)| (*path, node.hash_mask))
2848            .collect();
2849        let branch_node_tree_masks = hash_builder
2850            .updated_branch_nodes
2851            .clone()
2852            .unwrap_or_default()
2853            .iter()
2854            .map(|(path, node)| (*path, node.tree_mask))
2855            .collect();
2856
2857        let mut trie_updates = TrieUpdates::default();
2858        let removed_keys = node_iter.walker.take_removed_keys();
2859        trie_updates.finalize(hash_builder, removed_keys, destroyed_accounts);
2860
2861        (root, trie_updates, proof_nodes, branch_node_hash_masks, branch_node_tree_masks)
2862    }
2863
2864    /// Returns a `ParallelSparseTrie` pre-loaded with the given nodes, as well as leaf values
2865    /// inferred from any provided leaf nodes.
2866    fn new_test_trie<Nodes>(nodes: Nodes) -> ParallelSparseTrie
2867    where
2868        Nodes: Iterator<Item = (Nibbles, SparseNode)>,
2869    {
2870        let mut trie = ParallelSparseTrie::default().with_updates(true);
2871
2872        for (path, node) in nodes {
2873            let subtrie = trie.subtrie_for_path_mut(&path);
2874            if let SparseNode::Leaf { key, .. } = &node {
2875                let mut full_key = path;
2876                full_key.extend(key);
2877                subtrie.inner.values.insert(full_key, "LEAF VALUE".into());
2878            }
2879            subtrie.nodes.insert(path, node);
2880        }
2881        trie
2882    }
2883
2884    fn parallel_sparse_trie_nodes(
2885        sparse_trie: &ParallelSparseTrie,
2886    ) -> impl IntoIterator<Item = (&Nibbles, &SparseNode)> {
2887        let lower_sparse_nodes = sparse_trie
2888            .lower_subtries
2889            .iter()
2890            .filter_map(|subtrie| subtrie.as_revealed_ref())
2891            .flat_map(|subtrie| subtrie.nodes.iter());
2892
2893        let upper_sparse_nodes = sparse_trie.upper_subtrie.nodes.iter();
2894
2895        lower_sparse_nodes.chain(upper_sparse_nodes).sorted_by_key(|(path, _)| *path)
2896    }
2897
2898    /// Assert that the parallel sparse trie nodes and the proof nodes from the hash builder are
2899    /// equal.
2900    fn assert_eq_parallel_sparse_trie_proof_nodes(
2901        sparse_trie: &ParallelSparseTrie,
2902        proof_nodes: ProofNodes,
2903    ) {
2904        let proof_nodes = proof_nodes
2905            .into_nodes_sorted()
2906            .into_iter()
2907            .map(|(path, node)| (path, TrieNode::decode(&mut node.as_ref()).unwrap()));
2908
2909        let all_sparse_nodes = parallel_sparse_trie_nodes(sparse_trie);
2910
2911        for ((proof_node_path, proof_node), (sparse_node_path, sparse_node)) in
2912            proof_nodes.zip(all_sparse_nodes)
2913        {
2914            assert_eq!(&proof_node_path, sparse_node_path);
2915
2916            let equals = match (&proof_node, &sparse_node) {
2917                // Both nodes are empty
2918                (TrieNode::EmptyRoot, SparseNode::Empty) => true,
2919                // Both nodes are branches and have the same state mask
2920                (
2921                    TrieNode::Branch(BranchNode { state_mask: proof_state_mask, .. }),
2922                    SparseNode::Branch { state_mask: sparse_state_mask, .. },
2923                ) => proof_state_mask == sparse_state_mask,
2924                // Both nodes are extensions and have the same key
2925                (
2926                    TrieNode::Extension(ExtensionNode { key: proof_key, .. }),
2927                    SparseNode::Extension { key: sparse_key, .. },
2928                ) |
2929                // Both nodes are leaves and have the same key
2930                (
2931                    TrieNode::Leaf(LeafNode { key: proof_key, .. }),
2932                    SparseNode::Leaf { key: sparse_key, .. },
2933                ) => proof_key == sparse_key,
2934                // Empty and hash nodes are specific to the sparse trie, skip them
2935                (_, SparseNode::Empty | SparseNode::Hash(_)) => continue,
2936                _ => false,
2937            };
2938            assert!(
2939                equals,
2940                "path: {proof_node_path:?}\nproof node: {proof_node:?}\nsparse node: {sparse_node:?}"
2941            );
2942        }
2943    }
2944
2945    #[test]
2946    fn test_get_changed_subtries_empty() {
2947        let mut trie = ParallelSparseTrie::default();
2948        let mut prefix_set = PrefixSetMut::from([Nibbles::default()]).freeze();
2949
2950        let (subtries, unchanged_prefix_set) = trie.take_changed_lower_subtries(&mut prefix_set);
2951        assert!(subtries.is_empty());
2952        assert_eq!(unchanged_prefix_set, PrefixSetMut::from(prefix_set.iter().copied()));
2953    }
2954
2955    #[test]
2956    fn test_get_changed_subtries() {
2957        // Create a trie with three subtries
2958        let mut trie = ParallelSparseTrie::default();
2959        let subtrie_1 = Box::new(SparseSubtrie::new(Nibbles::from_nibbles([0x0, 0x0])));
2960        let subtrie_1_index = path_subtrie_index_unchecked(&subtrie_1.path);
2961        let subtrie_2 = Box::new(SparseSubtrie::new(Nibbles::from_nibbles([0x1, 0x0])));
2962        let subtrie_2_index = path_subtrie_index_unchecked(&subtrie_2.path);
2963        let subtrie_3 = Box::new(SparseSubtrie::new(Nibbles::from_nibbles([0x3, 0x0])));
2964        let subtrie_3_index = path_subtrie_index_unchecked(&subtrie_3.path);
2965
2966        // Add subtries at specific positions
2967        trie.lower_subtries[subtrie_1_index] = LowerSparseSubtrie::Revealed(subtrie_1.clone());
2968        trie.lower_subtries[subtrie_2_index] = LowerSparseSubtrie::Revealed(subtrie_2.clone());
2969        trie.lower_subtries[subtrie_3_index] = LowerSparseSubtrie::Revealed(subtrie_3);
2970
2971        let unchanged_prefix_set = PrefixSetMut::from([
2972            Nibbles::from_nibbles([0x0]),
2973            Nibbles::from_nibbles([0x2, 0x0, 0x0]),
2974        ]);
2975        // Create a prefix set with the keys that match only the second subtrie
2976        let mut prefix_set = PrefixSetMut::from([
2977            // Match second subtrie
2978            Nibbles::from_nibbles([0x1, 0x0, 0x0]),
2979            Nibbles::from_nibbles([0x1, 0x0, 0x1, 0x0]),
2980        ]);
2981        prefix_set.extend(unchanged_prefix_set);
2982        let mut prefix_set = prefix_set.freeze();
2983
2984        // Second subtrie should be removed and returned
2985        let (subtries, unchanged_prefix_set) = trie.take_changed_lower_subtries(&mut prefix_set);
2986        assert_eq!(
2987            subtries
2988                .into_iter()
2989                .map(|ChangedSubtrie { index, subtrie, prefix_set, .. }| {
2990                    (index, subtrie, prefix_set.iter().copied().collect::<Vec<_>>())
2991                })
2992                .collect::<Vec<_>>(),
2993            vec![(
2994                subtrie_2_index,
2995                subtrie_2,
2996                vec![
2997                    Nibbles::from_nibbles([0x1, 0x0, 0x0]),
2998                    Nibbles::from_nibbles([0x1, 0x0, 0x1, 0x0])
2999                ]
3000            )]
3001        );
3002        assert_eq!(unchanged_prefix_set, unchanged_prefix_set);
3003        assert!(trie.lower_subtries[subtrie_2_index].as_revealed_ref().is_none());
3004
3005        // First subtrie should remain unchanged
3006        assert_eq!(trie.lower_subtries[subtrie_1_index], LowerSparseSubtrie::Revealed(subtrie_1));
3007    }
3008
3009    #[test]
3010    fn test_get_changed_subtries_all() {
3011        // Create a trie with three subtries
3012        let mut trie = ParallelSparseTrie::default();
3013        let subtrie_1 = Box::new(SparseSubtrie::new(Nibbles::from_nibbles([0x0, 0x0])));
3014        let subtrie_1_index = path_subtrie_index_unchecked(&subtrie_1.path);
3015        let subtrie_2 = Box::new(SparseSubtrie::new(Nibbles::from_nibbles([0x1, 0x0])));
3016        let subtrie_2_index = path_subtrie_index_unchecked(&subtrie_2.path);
3017        let subtrie_3 = Box::new(SparseSubtrie::new(Nibbles::from_nibbles([0x3, 0x0])));
3018        let subtrie_3_index = path_subtrie_index_unchecked(&subtrie_3.path);
3019
3020        // Add subtries at specific positions
3021        trie.lower_subtries[subtrie_1_index] = LowerSparseSubtrie::Revealed(subtrie_1.clone());
3022        trie.lower_subtries[subtrie_2_index] = LowerSparseSubtrie::Revealed(subtrie_2.clone());
3023        trie.lower_subtries[subtrie_3_index] = LowerSparseSubtrie::Revealed(subtrie_3.clone());
3024
3025        // Create a prefix set that matches any key
3026        let mut prefix_set = PrefixSetMut::all().freeze();
3027
3028        // All subtries should be removed and returned
3029        let (subtries, unchanged_prefix_set) = trie.take_changed_lower_subtries(&mut prefix_set);
3030        assert_eq!(
3031            subtries
3032                .into_iter()
3033                .map(|ChangedSubtrie { index, subtrie, prefix_set, .. }| {
3034                    (index, subtrie, prefix_set.all())
3035                })
3036                .collect::<Vec<_>>(),
3037            vec![
3038                (subtrie_1_index, subtrie_1, true),
3039                (subtrie_2_index, subtrie_2, true),
3040                (subtrie_3_index, subtrie_3, true)
3041            ]
3042        );
3043        assert_eq!(unchanged_prefix_set, PrefixSetMut::all());
3044
3045        assert!(trie.lower_subtries.iter().all(|subtrie| subtrie.as_revealed_ref().is_none()));
3046    }
3047
3048    #[test]
3049    fn test_sparse_subtrie_type() {
3050        assert_eq!(SparseSubtrieType::from_path(&Nibbles::new()), SparseSubtrieType::Upper);
3051        assert_eq!(
3052            SparseSubtrieType::from_path(&Nibbles::from_nibbles([0])),
3053            SparseSubtrieType::Upper
3054        );
3055        assert_eq!(
3056            SparseSubtrieType::from_path(&Nibbles::from_nibbles([15])),
3057            SparseSubtrieType::Upper
3058        );
3059        assert_eq!(
3060            SparseSubtrieType::from_path(&Nibbles::from_nibbles([0, 0])),
3061            SparseSubtrieType::Lower(0)
3062        );
3063        assert_eq!(
3064            SparseSubtrieType::from_path(&Nibbles::from_nibbles([0, 0, 0])),
3065            SparseSubtrieType::Lower(0)
3066        );
3067        assert_eq!(
3068            SparseSubtrieType::from_path(&Nibbles::from_nibbles([0, 1])),
3069            SparseSubtrieType::Lower(1)
3070        );
3071        assert_eq!(
3072            SparseSubtrieType::from_path(&Nibbles::from_nibbles([0, 1, 0])),
3073            SparseSubtrieType::Lower(1)
3074        );
3075        assert_eq!(
3076            SparseSubtrieType::from_path(&Nibbles::from_nibbles([0, 15])),
3077            SparseSubtrieType::Lower(15)
3078        );
3079        assert_eq!(
3080            SparseSubtrieType::from_path(&Nibbles::from_nibbles([15, 0])),
3081            SparseSubtrieType::Lower(240)
3082        );
3083        assert_eq!(
3084            SparseSubtrieType::from_path(&Nibbles::from_nibbles([15, 1])),
3085            SparseSubtrieType::Lower(241)
3086        );
3087        assert_eq!(
3088            SparseSubtrieType::from_path(&Nibbles::from_nibbles([15, 15])),
3089            SparseSubtrieType::Lower(255)
3090        );
3091        assert_eq!(
3092            SparseSubtrieType::from_path(&Nibbles::from_nibbles([15, 15, 15])),
3093            SparseSubtrieType::Lower(255)
3094        );
3095    }
3096
3097    #[test]
3098    fn test_reveal_node_leaves() {
3099        let mut trie = ParallelSparseTrie::default();
3100
3101        // Reveal leaf in the upper trie
3102        {
3103            let path = Nibbles::from_nibbles([0x1]);
3104            let node = create_leaf_node([0x2, 0x3], 42);
3105            let masks = TrieMasks::none();
3106
3107            trie.reveal_nodes(vec![RevealedSparseNode { path, node, masks }]).unwrap();
3108
3109            assert_matches!(
3110                trie.upper_subtrie.nodes.get(&path),
3111                Some(SparseNode::Leaf { key, hash: None })
3112                if key == &Nibbles::from_nibbles([0x2, 0x3])
3113            );
3114
3115            let full_path = Nibbles::from_nibbles([0x1, 0x2, 0x3]);
3116            assert_eq!(
3117                trie.upper_subtrie.inner.values.get(&full_path),
3118                Some(&encode_account_value(42))
3119            );
3120        }
3121
3122        // Reveal leaf in a lower trie
3123        {
3124            let path = Nibbles::from_nibbles([0x1, 0x2]);
3125            let node = create_leaf_node([0x3, 0x4], 42);
3126            let masks = TrieMasks::none();
3127
3128            trie.reveal_nodes(vec![RevealedSparseNode { path, node, masks }]).unwrap();
3129
3130            // Check that the lower subtrie was created
3131            let idx = path_subtrie_index_unchecked(&path);
3132            assert!(trie.lower_subtries[idx].as_revealed_ref().is_some());
3133
3134            // Check that the lower subtrie's path was correctly set
3135            let lower_subtrie = trie.lower_subtries[idx].as_revealed_ref().unwrap();
3136            assert_eq!(lower_subtrie.path, path);
3137
3138            assert_matches!(
3139                lower_subtrie.nodes.get(&path),
3140                Some(SparseNode::Leaf { key, hash: None })
3141                if key == &Nibbles::from_nibbles([0x3, 0x4])
3142            );
3143        }
3144
3145        // Reveal leaf in a lower trie with a longer path, shouldn't result in the subtrie's root
3146        // path changing.
3147        {
3148            let path = Nibbles::from_nibbles([0x1, 0x2, 0x3]);
3149            let node = create_leaf_node([0x4, 0x5], 42);
3150            let masks = TrieMasks::none();
3151
3152            trie.reveal_nodes(vec![RevealedSparseNode { path, node, masks }]).unwrap();
3153
3154            // Check that the lower subtrie's path hasn't changed
3155            let idx = path_subtrie_index_unchecked(&path);
3156            let lower_subtrie = trie.lower_subtries[idx].as_revealed_ref().unwrap();
3157            assert_eq!(lower_subtrie.path, Nibbles::from_nibbles([0x1, 0x2]));
3158        }
3159    }
3160
3161    #[test]
3162    fn test_reveal_node_extension_all_upper() {
3163        let path = Nibbles::new();
3164        let child_hash = B256::repeat_byte(0xab);
3165        let node = create_extension_node([0x1], child_hash);
3166        let masks = TrieMasks::none();
3167        let trie = ParallelSparseTrie::from_root(node, masks, true).unwrap();
3168
3169        assert_matches!(
3170            trie.upper_subtrie.nodes.get(&path),
3171            Some(SparseNode::Extension { key, hash: None, .. })
3172            if key == &Nibbles::from_nibbles([0x1])
3173        );
3174
3175        // Child path should be in upper trie
3176        let child_path = Nibbles::from_nibbles([0x1]);
3177        assert_eq!(trie.upper_subtrie.nodes.get(&child_path), Some(&SparseNode::Hash(child_hash)));
3178    }
3179
3180    #[test]
3181    fn test_reveal_node_extension_cross_level() {
3182        let path = Nibbles::new();
3183        let child_hash = B256::repeat_byte(0xcd);
3184        let node = create_extension_node([0x1, 0x2, 0x3], child_hash);
3185        let masks = TrieMasks::none();
3186        let trie = ParallelSparseTrie::from_root(node, masks, true).unwrap();
3187
3188        // Extension node should be in upper trie
3189        assert_matches!(
3190            trie.upper_subtrie.nodes.get(&path),
3191            Some(SparseNode::Extension { key, hash: None, .. })
3192            if key == &Nibbles::from_nibbles([0x1, 0x2, 0x3])
3193        );
3194
3195        // Child path (0x1, 0x2, 0x3) should be in lower trie
3196        let child_path = Nibbles::from_nibbles([0x1, 0x2, 0x3]);
3197        let idx = path_subtrie_index_unchecked(&child_path);
3198        assert!(trie.lower_subtries[idx].as_revealed_ref().is_some());
3199
3200        let lower_subtrie = trie.lower_subtries[idx].as_revealed_ref().unwrap();
3201        assert_eq!(lower_subtrie.path, child_path);
3202        assert_eq!(lower_subtrie.nodes.get(&child_path), Some(&SparseNode::Hash(child_hash)));
3203    }
3204
3205    #[test]
3206    fn test_reveal_node_extension_cross_level_boundary() {
3207        let mut trie = ParallelSparseTrie::default();
3208        let path = Nibbles::from_nibbles([0x1]);
3209        let child_hash = B256::repeat_byte(0xcd);
3210        let node = create_extension_node([0x2], child_hash);
3211        let masks = TrieMasks::none();
3212
3213        trie.reveal_nodes(vec![RevealedSparseNode { path, node, masks }]).unwrap();
3214
3215        // Extension node should be in upper trie
3216        assert_matches!(
3217            trie.upper_subtrie.nodes.get(&path),
3218            Some(SparseNode::Extension { key, hash: None, .. })
3219            if key == &Nibbles::from_nibbles([0x2])
3220        );
3221
3222        // Child path (0x1, 0x2) should be in lower trie
3223        let child_path = Nibbles::from_nibbles([0x1, 0x2]);
3224        let idx = path_subtrie_index_unchecked(&child_path);
3225        assert!(trie.lower_subtries[idx].as_revealed_ref().is_some());
3226
3227        let lower_subtrie = trie.lower_subtries[idx].as_revealed_ref().unwrap();
3228        assert_eq!(lower_subtrie.path, child_path);
3229        assert_eq!(lower_subtrie.nodes.get(&child_path), Some(&SparseNode::Hash(child_hash)));
3230    }
3231
3232    #[test]
3233    fn test_reveal_node_branch_all_upper() {
3234        let path = Nibbles::new();
3235        let child_hashes = [
3236            RlpNode::word_rlp(&B256::repeat_byte(0x11)),
3237            RlpNode::word_rlp(&B256::repeat_byte(0x22)),
3238        ];
3239        let node = create_branch_node_with_children(&[0x0, 0x5], child_hashes.clone());
3240        let masks = TrieMasks::none();
3241        let trie = ParallelSparseTrie::from_root(node, masks, true).unwrap();
3242
3243        // Branch node should be in upper trie
3244        assert_matches!(
3245            trie.upper_subtrie.nodes.get(&path),
3246            Some(SparseNode::Branch { state_mask, hash: None, .. })
3247            if *state_mask == 0b0000000000100001.into()
3248        );
3249
3250        // Children should be in upper trie (paths of length 2)
3251        let child_path_0 = Nibbles::from_nibbles([0x0]);
3252        let child_path_5 = Nibbles::from_nibbles([0x5]);
3253        assert_eq!(
3254            trie.upper_subtrie.nodes.get(&child_path_0),
3255            Some(&SparseNode::Hash(child_hashes[0].as_hash().unwrap()))
3256        );
3257        assert_eq!(
3258            trie.upper_subtrie.nodes.get(&child_path_5),
3259            Some(&SparseNode::Hash(child_hashes[1].as_hash().unwrap()))
3260        );
3261    }
3262
3263    #[test]
3264    fn test_reveal_node_branch_cross_level() {
3265        let mut trie = ParallelSparseTrie::default();
3266        let path = Nibbles::from_nibbles([0x1]); // Exactly 1 nibbles - boundary case
3267        let child_hashes = [
3268            RlpNode::word_rlp(&B256::repeat_byte(0x33)),
3269            RlpNode::word_rlp(&B256::repeat_byte(0x44)),
3270            RlpNode::word_rlp(&B256::repeat_byte(0x55)),
3271        ];
3272        let node = create_branch_node_with_children(&[0x0, 0x7, 0xf], child_hashes.clone());
3273        let masks = TrieMasks::none();
3274
3275        trie.reveal_nodes(vec![RevealedSparseNode { path, node, masks }]).unwrap();
3276
3277        // Branch node should be in upper trie
3278        assert_matches!(
3279            trie.upper_subtrie.nodes.get(&path),
3280            Some(SparseNode::Branch { state_mask, hash: None, .. })
3281            if *state_mask == 0b1000000010000001.into()
3282        );
3283
3284        // All children should be in lower tries since they have paths of length 3
3285        let child_paths = [
3286            Nibbles::from_nibbles([0x1, 0x0]),
3287            Nibbles::from_nibbles([0x1, 0x7]),
3288            Nibbles::from_nibbles([0x1, 0xf]),
3289        ];
3290
3291        for (i, child_path) in child_paths.iter().enumerate() {
3292            let idx = path_subtrie_index_unchecked(child_path);
3293            let lower_subtrie = trie.lower_subtries[idx].as_revealed_ref().unwrap();
3294            assert_eq!(&lower_subtrie.path, child_path);
3295            assert_eq!(
3296                lower_subtrie.nodes.get(child_path),
3297                Some(&SparseNode::Hash(child_hashes[i].as_hash().unwrap())),
3298            );
3299        }
3300    }
3301
3302    #[test]
3303    fn test_update_subtrie_hashes_prefix_set_matching() {
3304        // Create a trie and reveal leaf nodes using reveal_nodes
3305        let mut trie = ParallelSparseTrie::default();
3306
3307        // Create dummy leaf nodes.
3308        let leaf_1_full_path = Nibbles::from_nibbles([0; 64]);
3309        let leaf_1_path = leaf_1_full_path.slice(..2);
3310        let leaf_1_key = leaf_1_full_path.slice(2..);
3311        let leaf_2_full_path = Nibbles::from_nibbles([vec![0, 1], vec![0; 62]].concat());
3312        let leaf_2_path = leaf_2_full_path.slice(..2);
3313        let leaf_2_key = leaf_2_full_path.slice(2..);
3314        let leaf_3_full_path = Nibbles::from_nibbles([vec![0, 2], vec![0; 62]].concat());
3315        let leaf_3_path = leaf_3_full_path.slice(..2);
3316        let leaf_3_key = leaf_3_full_path.slice(2..);
3317        let leaf_1 = create_leaf_node(leaf_1_key.to_vec(), 1);
3318        let leaf_2 = create_leaf_node(leaf_2_key.to_vec(), 2);
3319        let leaf_3 = create_leaf_node(leaf_3_key.to_vec(), 3);
3320
3321        // Create branch node with hashes for each leaf.
3322        let child_hashes = [
3323            RlpNode::word_rlp(&B256::repeat_byte(0x00)),
3324            RlpNode::word_rlp(&B256::repeat_byte(0x11)),
3325            // deliberately omit hash for leaf_3
3326        ];
3327        let branch_path = Nibbles::from_nibbles([0x0]);
3328        let branch_node = create_branch_node_with_children(&[0x0, 0x1, 0x2], child_hashes);
3329
3330        // Reveal nodes using reveal_nodes
3331        trie.reveal_nodes(vec![
3332            RevealedSparseNode { path: branch_path, node: branch_node, masks: TrieMasks::none() },
3333            RevealedSparseNode { path: leaf_1_path, node: leaf_1, masks: TrieMasks::none() },
3334            RevealedSparseNode { path: leaf_2_path, node: leaf_2, masks: TrieMasks::none() },
3335            RevealedSparseNode { path: leaf_3_path, node: leaf_3, masks: TrieMasks::none() },
3336        ])
3337        .unwrap();
3338
3339        // Calculate subtrie indexes
3340        let subtrie_1_index = SparseSubtrieType::from_path(&leaf_1_path).lower_index().unwrap();
3341        let subtrie_2_index = SparseSubtrieType::from_path(&leaf_2_path).lower_index().unwrap();
3342        let subtrie_3_index = SparseSubtrieType::from_path(&leaf_3_path).lower_index().unwrap();
3343
3344        let mut unchanged_prefix_set = PrefixSetMut::from([
3345            Nibbles::from_nibbles([0x0]),
3346            leaf_2_full_path,
3347            Nibbles::from_nibbles([0x3, 0x0, 0x0]),
3348        ]);
3349        // Create a prefix set with the keys that match only the second subtrie
3350        let mut prefix_set = PrefixSetMut::from([
3351            // Match second subtrie
3352            Nibbles::from_nibbles([0x0, 0x1, 0x0]),
3353            Nibbles::from_nibbles([0x0, 0x1, 0x1, 0x0]),
3354        ]);
3355        prefix_set.extend(unchanged_prefix_set.clone());
3356        trie.prefix_set = prefix_set;
3357
3358        // Update subtrie hashes
3359        trie.update_subtrie_hashes();
3360
3361        // We expect that leaf 3 (0x02) should have been added to the prefix set, because it is
3362        // missing a hash and is the root node of a lower subtrie, and therefore would need to have
3363        // that hash calculated by `update_upper_subtrie_hashes`.
3364        unchanged_prefix_set.insert(leaf_3_full_path);
3365
3366        // Check that the prefix set was updated
3367        assert_eq!(
3368            trie.prefix_set.clone().freeze().into_iter().collect::<Vec<_>>(),
3369            unchanged_prefix_set.freeze().into_iter().collect::<Vec<_>>()
3370        );
3371        // Check that subtries were returned back to the array
3372        assert!(trie.lower_subtries[subtrie_1_index].as_revealed_ref().is_some());
3373        assert!(trie.lower_subtries[subtrie_2_index].as_revealed_ref().is_some());
3374        assert!(trie.lower_subtries[subtrie_3_index].as_revealed_ref().is_some());
3375    }
3376
3377    #[test]
3378    fn test_subtrie_update_hashes() {
3379        let mut subtrie = Box::new(SparseSubtrie::new(Nibbles::from_nibbles([0x0, 0x0])));
3380
3381        // Create leaf nodes with paths 0x0...0, 0x00001...0, 0x0010...0
3382        let leaf_1_full_path = Nibbles::from_nibbles([0; 64]);
3383        let leaf_1_path = leaf_1_full_path.slice(..5);
3384        let leaf_1_key = leaf_1_full_path.slice(5..);
3385        let leaf_2_full_path = Nibbles::from_nibbles([vec![0, 0, 0, 0, 1], vec![0; 59]].concat());
3386        let leaf_2_path = leaf_2_full_path.slice(..5);
3387        let leaf_2_key = leaf_2_full_path.slice(5..);
3388        let leaf_3_full_path = Nibbles::from_nibbles([vec![0, 0, 1], vec![0; 61]].concat());
3389        let leaf_3_path = leaf_3_full_path.slice(..3);
3390        let leaf_3_key = leaf_3_full_path.slice(3..);
3391
3392        let account_1 = create_account(1);
3393        let account_2 = create_account(2);
3394        let account_3 = create_account(3);
3395        let leaf_1 = create_leaf_node(leaf_1_key.to_vec(), account_1.nonce);
3396        let leaf_2 = create_leaf_node(leaf_2_key.to_vec(), account_2.nonce);
3397        let leaf_3 = create_leaf_node(leaf_3_key.to_vec(), account_3.nonce);
3398
3399        // Create bottom branch node
3400        let branch_1_path = Nibbles::from_nibbles([0, 0, 0, 0]);
3401        let branch_1 = create_branch_node_with_children(
3402            &[0, 1],
3403            vec![
3404                RlpNode::from_rlp(&alloy_rlp::encode(&leaf_1)),
3405                RlpNode::from_rlp(&alloy_rlp::encode(&leaf_2)),
3406            ],
3407        );
3408
3409        // Create an extension node
3410        let extension_path = Nibbles::from_nibbles([0, 0, 0]);
3411        let extension_key = Nibbles::from_nibbles([0]);
3412        let extension = create_extension_node(
3413            extension_key.to_vec(),
3414            RlpNode::from_rlp(&alloy_rlp::encode(&branch_1)).as_hash().unwrap(),
3415        );
3416
3417        // Create top branch node
3418        let branch_2_path = Nibbles::from_nibbles([0, 0]);
3419        let branch_2 = create_branch_node_with_children(
3420            &[0, 1],
3421            vec![
3422                RlpNode::from_rlp(&alloy_rlp::encode(&extension)),
3423                RlpNode::from_rlp(&alloy_rlp::encode(&leaf_3)),
3424            ],
3425        );
3426
3427        // Reveal nodes
3428        subtrie.reveal_node(branch_2_path, &branch_2, TrieMasks::none()).unwrap();
3429        subtrie.reveal_node(leaf_1_path, &leaf_1, TrieMasks::none()).unwrap();
3430        subtrie.reveal_node(extension_path, &extension, TrieMasks::none()).unwrap();
3431        subtrie.reveal_node(branch_1_path, &branch_1, TrieMasks::none()).unwrap();
3432        subtrie.reveal_node(leaf_2_path, &leaf_2, TrieMasks::none()).unwrap();
3433        subtrie.reveal_node(leaf_3_path, &leaf_3, TrieMasks::none()).unwrap();
3434
3435        // Run hash builder for two leaf nodes
3436        let (_, _, proof_nodes, _, _) = run_hash_builder(
3437            [
3438                (leaf_1_full_path, account_1),
3439                (leaf_2_full_path, account_2),
3440                (leaf_3_full_path, account_3),
3441            ],
3442            NoopAccountTrieCursor::default(),
3443            Default::default(),
3444            [
3445                branch_1_path,
3446                extension_path,
3447                branch_2_path,
3448                leaf_1_full_path,
3449                leaf_2_full_path,
3450                leaf_3_full_path,
3451            ],
3452        );
3453
3454        // Update hashes for the subtrie
3455        subtrie.update_hashes(
3456            &mut PrefixSetMut::from([leaf_1_full_path, leaf_2_full_path, leaf_3_full_path])
3457                .freeze(),
3458            &mut None,
3459            &HashMap::default(),
3460            &HashMap::default(),
3461        );
3462
3463        // Compare hashes between hash builder and subtrie
3464        let hash_builder_branch_1_hash =
3465            RlpNode::from_rlp(proof_nodes.get(&branch_1_path).unwrap().as_ref()).as_hash().unwrap();
3466        let subtrie_branch_1_hash = subtrie.nodes.get(&branch_1_path).unwrap().hash().unwrap();
3467        assert_eq!(hash_builder_branch_1_hash, subtrie_branch_1_hash);
3468
3469        let hash_builder_extension_hash =
3470            RlpNode::from_rlp(proof_nodes.get(&extension_path).unwrap().as_ref())
3471                .as_hash()
3472                .unwrap();
3473        let subtrie_extension_hash = subtrie.nodes.get(&extension_path).unwrap().hash().unwrap();
3474        assert_eq!(hash_builder_extension_hash, subtrie_extension_hash);
3475
3476        let hash_builder_branch_2_hash =
3477            RlpNode::from_rlp(proof_nodes.get(&branch_2_path).unwrap().as_ref()).as_hash().unwrap();
3478        let subtrie_branch_2_hash = subtrie.nodes.get(&branch_2_path).unwrap().hash().unwrap();
3479        assert_eq!(hash_builder_branch_2_hash, subtrie_branch_2_hash);
3480
3481        let subtrie_leaf_1_hash = subtrie.nodes.get(&leaf_1_path).unwrap().hash().unwrap();
3482        let hash_builder_leaf_1_hash =
3483            RlpNode::from_rlp(proof_nodes.get(&leaf_1_path).unwrap().as_ref()).as_hash().unwrap();
3484        assert_eq!(hash_builder_leaf_1_hash, subtrie_leaf_1_hash);
3485
3486        let hash_builder_leaf_2_hash =
3487            RlpNode::from_rlp(proof_nodes.get(&leaf_2_path).unwrap().as_ref()).as_hash().unwrap();
3488        let subtrie_leaf_2_hash = subtrie.nodes.get(&leaf_2_path).unwrap().hash().unwrap();
3489        assert_eq!(hash_builder_leaf_2_hash, subtrie_leaf_2_hash);
3490
3491        let hash_builder_leaf_3_hash =
3492            RlpNode::from_rlp(proof_nodes.get(&leaf_3_path).unwrap().as_ref()).as_hash().unwrap();
3493        let subtrie_leaf_3_hash = subtrie.nodes.get(&leaf_3_path).unwrap().hash().unwrap();
3494        assert_eq!(hash_builder_leaf_3_hash, subtrie_leaf_3_hash);
3495    }
3496
3497    #[test]
3498    fn test_remove_leaf_branch_becomes_extension() {
3499        //
3500        // 0x:      Extension (Key = 5)
3501        // 0x5:     └── Branch (Mask = 1001)
3502        // 0x50:        ├── 0 -> Extension (Key = 23)
3503        // 0x5023:      │        └── Branch (Mask = 0101)
3504        // 0x50231:     │            ├── 1 -> Leaf
3505        // 0x50233:     │            └── 3 -> Leaf
3506        // 0x53:        └── 3 -> Leaf (Key = 7)
3507        //
3508        // After removing 0x53, extension+branch+extension become a single extension
3509        //
3510        let mut trie = new_test_trie(
3511            [
3512                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
3513                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(TrieMask::new(0b1001))),
3514                (
3515                    Nibbles::from_nibbles([0x5, 0x0]),
3516                    SparseNode::new_ext(Nibbles::from_nibbles([0x2, 0x3])),
3517                ),
3518                (
3519                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3]),
3520                    SparseNode::new_branch(TrieMask::new(0b0101)),
3521                ),
3522                (
3523                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]),
3524                    SparseNode::new_leaf(Nibbles::new()),
3525                ),
3526                (
3527                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
3528                    SparseNode::new_leaf(Nibbles::new()),
3529                ),
3530                (
3531                    Nibbles::from_nibbles([0x5, 0x3]),
3532                    SparseNode::new_leaf(Nibbles::from_nibbles([0x7])),
3533                ),
3534            ]
3535            .into_iter(),
3536        );
3537
3538        let provider = MockTrieNodeProvider::new();
3539
3540        // Remove the leaf with a full path of 0x537
3541        let leaf_full_path = Nibbles::from_nibbles([0x5, 0x3, 0x7]);
3542        trie.remove_leaf(&leaf_full_path, provider).unwrap();
3543
3544        let upper_subtrie = &trie.upper_subtrie;
3545        let lower_subtrie_50 = trie.lower_subtries[0x50].as_revealed_ref().unwrap();
3546
3547        // Check that the `SparseSubtrie` the leaf was removed from was itself removed, as it is now
3548        // empty.
3549        assert_matches!(trie.lower_subtries[0x53].as_revealed_ref(), None);
3550
3551        // Check that the leaf node was removed, and that its parent/grandparent were modified
3552        // appropriately.
3553        assert_matches!(
3554            upper_subtrie.nodes.get(&Nibbles::from_nibbles([])),
3555            Some(SparseNode::Extension{ key, ..})
3556            if key == &Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3])
3557        );
3558        assert_matches!(upper_subtrie.nodes.get(&Nibbles::from_nibbles([0x5])), None);
3559        assert_matches!(lower_subtrie_50.nodes.get(&Nibbles::from_nibbles([0x5, 0x0])), None);
3560        assert_matches!(
3561            lower_subtrie_50.nodes.get(&Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3])),
3562            Some(SparseNode::Branch{ state_mask, .. })
3563            if *state_mask == 0b0101.into()
3564        );
3565    }
3566
3567    #[test]
3568    fn test_remove_leaf_branch_becomes_leaf() {
3569        //
3570        // 0x:      Branch (Mask = 0011)
3571        // 0x0:     ├── 0 -> Leaf (Key = 12)
3572        // 0x1:     └── 1 -> Leaf (Key = 34)
3573        //
3574        // After removing 0x012, branch becomes a leaf
3575        //
3576        let mut trie = new_test_trie(
3577            [
3578                (Nibbles::default(), SparseNode::new_branch(TrieMask::new(0b0011))),
3579                (
3580                    Nibbles::from_nibbles([0x0]),
3581                    SparseNode::new_leaf(Nibbles::from_nibbles([0x1, 0x2])),
3582                ),
3583                (
3584                    Nibbles::from_nibbles([0x1]),
3585                    SparseNode::new_leaf(Nibbles::from_nibbles([0x3, 0x4])),
3586                ),
3587            ]
3588            .into_iter(),
3589        );
3590
3591        // Add the branch node to updated_nodes to simulate it being modified earlier
3592        if let Some(updates) = trie.updates.as_mut() {
3593            updates
3594                .updated_nodes
3595                .insert(Nibbles::default(), BranchNodeCompact::new(0b11, 0, 0, vec![], None));
3596        }
3597
3598        let provider = MockTrieNodeProvider::new();
3599
3600        // Remove the leaf with a full path of 0x012
3601        let leaf_full_path = Nibbles::from_nibbles([0x0, 0x1, 0x2]);
3602        trie.remove_leaf(&leaf_full_path, provider).unwrap();
3603
3604        let upper_subtrie = &trie.upper_subtrie;
3605
3606        // Check that the leaf's value was removed
3607        assert_matches!(upper_subtrie.inner.values.get(&leaf_full_path), None);
3608
3609        // Check that the branch node collapsed into a leaf node with the remaining child's key
3610        assert_matches!(
3611            upper_subtrie.nodes.get(&Nibbles::default()),
3612            Some(SparseNode::Leaf{ key, ..})
3613            if key == &Nibbles::from_nibbles([0x1, 0x3, 0x4])
3614        );
3615
3616        // Check that the remaining child node was removed
3617        assert_matches!(upper_subtrie.nodes.get(&Nibbles::from_nibbles([0x1])), None);
3618        // Check that the removed child node was also removed
3619        assert_matches!(upper_subtrie.nodes.get(&Nibbles::from_nibbles([0x0])), None);
3620
3621        // Check that updates were tracked correctly when branch collapsed
3622        let updates = trie.updates.as_ref().unwrap();
3623
3624        // The branch at root should be marked as removed since it collapsed
3625        assert!(updates.removed_nodes.contains(&Nibbles::default()));
3626
3627        // The branch should no longer be in updated_nodes
3628        assert!(!updates.updated_nodes.contains_key(&Nibbles::default()));
3629    }
3630
3631    #[test]
3632    fn test_remove_leaf_extension_becomes_leaf() {
3633        //
3634        // 0x:      Extension (Key = 5)
3635        // 0x5:     └── Branch (Mask = 0011)
3636        // 0x50:        ├── 0 -> Leaf (Key = 12)
3637        // 0x51:        └── 1 -> Leaf (Key = 34)
3638        //
3639        // After removing 0x5012, extension+branch becomes a leaf
3640        //
3641        let mut trie = new_test_trie(
3642            [
3643                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
3644                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(TrieMask::new(0b0011))),
3645                (
3646                    Nibbles::from_nibbles([0x5, 0x0]),
3647                    SparseNode::new_leaf(Nibbles::from_nibbles([0x1, 0x2])),
3648                ),
3649                (
3650                    Nibbles::from_nibbles([0x5, 0x1]),
3651                    SparseNode::new_leaf(Nibbles::from_nibbles([0x3, 0x4])),
3652                ),
3653            ]
3654            .into_iter(),
3655        );
3656
3657        let provider = MockTrieNodeProvider::new();
3658
3659        // Remove the leaf with a full path of 0x5012
3660        let leaf_full_path = Nibbles::from_nibbles([0x5, 0x0, 0x1, 0x2]);
3661        trie.remove_leaf(&leaf_full_path, provider).unwrap();
3662
3663        let upper_subtrie = &trie.upper_subtrie;
3664
3665        // Check that both lower subtries were removed. 0x50 should have been removed because
3666        // removing its leaf made it empty. 0x51 should have been removed after its own leaf was
3667        // collapsed into the upper trie, leaving it also empty.
3668        assert_matches!(trie.lower_subtries[0x50].as_revealed_ref(), None);
3669        assert_matches!(trie.lower_subtries[0x51].as_revealed_ref(), None);
3670
3671        // Check that the other leaf's value was moved to the upper trie
3672        let other_leaf_full_value = Nibbles::from_nibbles([0x5, 0x1, 0x3, 0x4]);
3673        assert_matches!(upper_subtrie.inner.values.get(&other_leaf_full_value), Some(_));
3674
3675        // Check that the extension node collapsed into a leaf node
3676        assert_matches!(
3677            upper_subtrie.nodes.get(&Nibbles::default()),
3678            Some(SparseNode::Leaf{ key, ..})
3679            if key == &Nibbles::from_nibbles([0x5, 0x1, 0x3, 0x4])
3680        );
3681
3682        // Check that intermediate nodes were removed
3683        assert_matches!(upper_subtrie.nodes.get(&Nibbles::from_nibbles([0x5])), None);
3684    }
3685
3686    #[test]
3687    fn test_remove_leaf_branch_on_branch() {
3688        //
3689        // 0x:      Branch (Mask = 0101)
3690        // 0x0:     ├── 0 -> Leaf (Key = 12)
3691        // 0x2:     └── 2 -> Branch (Mask = 0011)
3692        // 0x20:        ├── 0 -> Leaf (Key = 34)
3693        // 0x21:        └── 1 -> Leaf (Key = 56)
3694        //
3695        // After removing 0x2034, the inner branch becomes a leaf
3696        //
3697        let mut trie = new_test_trie(
3698            [
3699                (Nibbles::default(), SparseNode::new_branch(TrieMask::new(0b0101))),
3700                (
3701                    Nibbles::from_nibbles([0x0]),
3702                    SparseNode::new_leaf(Nibbles::from_nibbles([0x1, 0x2])),
3703                ),
3704                (Nibbles::from_nibbles([0x2]), SparseNode::new_branch(TrieMask::new(0b0011))),
3705                (
3706                    Nibbles::from_nibbles([0x2, 0x0]),
3707                    SparseNode::new_leaf(Nibbles::from_nibbles([0x3, 0x4])),
3708                ),
3709                (
3710                    Nibbles::from_nibbles([0x2, 0x1]),
3711                    SparseNode::new_leaf(Nibbles::from_nibbles([0x5, 0x6])),
3712                ),
3713            ]
3714            .into_iter(),
3715        );
3716
3717        let provider = MockTrieNodeProvider::new();
3718
3719        // Remove the leaf with a full path of 0x2034
3720        let leaf_full_path = Nibbles::from_nibbles([0x2, 0x0, 0x3, 0x4]);
3721        trie.remove_leaf(&leaf_full_path, provider).unwrap();
3722
3723        let upper_subtrie = &trie.upper_subtrie;
3724
3725        // Check that both lower subtries were removed. 0x20 should have been removed because
3726        // removing its leaf made it empty. 0x21 should have been removed after its own leaf was
3727        // collapsed into the upper trie, leaving it also empty.
3728        assert_matches!(trie.lower_subtries[0x20].as_revealed_ref(), None);
3729        assert_matches!(trie.lower_subtries[0x21].as_revealed_ref(), None);
3730
3731        // Check that the other leaf's value was moved to the upper trie
3732        let other_leaf_full_value = Nibbles::from_nibbles([0x2, 0x1, 0x5, 0x6]);
3733        assert_matches!(upper_subtrie.inner.values.get(&other_leaf_full_value), Some(_));
3734
3735        // Check that the root branch still exists unchanged
3736        assert_matches!(
3737            upper_subtrie.nodes.get(&Nibbles::default()),
3738            Some(SparseNode::Branch{ state_mask, .. })
3739            if *state_mask == 0b0101.into()
3740        );
3741
3742        // Check that the inner branch became an extension
3743        assert_matches!(
3744            upper_subtrie.nodes.get(&Nibbles::from_nibbles([0x2])),
3745            Some(SparseNode::Leaf{ key, ..})
3746            if key == &Nibbles::from_nibbles([0x1, 0x5, 0x6])
3747        );
3748    }
3749
3750    #[test]
3751    fn test_remove_leaf_lower_subtrie_root_path_update() {
3752        //
3753        // 0x:        Extension (Key = 123, root of lower subtrie)
3754        // 0x123:     └── Branch (Mask = 0011000)
3755        // 0x1233:        ├── 3 -> Leaf (Key = [])
3756        // 0x1234:        └── 4 -> Extension (Key = 5)
3757        // 0x12345:           └── Branch (Mask = 0011)
3758        // 0x123450:              ├── 0 -> Leaf (Key = [])
3759        // 0x123451:              └── 1 -> Leaf (Key = [])
3760        //
3761        // After removing leaf at 0x1233, the branch at 0x123 becomes an extension to 0x12345, which
3762        // then gets merged with the root extension at 0x. The lower subtrie's `path` field should
3763        // be updated from 0x123 to 0x12345.
3764        //
3765        let mut trie = new_test_trie(
3766            [
3767                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x1, 0x2, 0x3]))),
3768                (
3769                    Nibbles::from_nibbles([0x1, 0x2, 0x3]),
3770                    SparseNode::new_branch(TrieMask::new(0b0011000)),
3771                ),
3772                (
3773                    Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x3]),
3774                    SparseNode::new_leaf(Nibbles::default()),
3775                ),
3776                (
3777                    Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4]),
3778                    SparseNode::new_ext(Nibbles::from_nibbles([0x5])),
3779                ),
3780                (
3781                    Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4, 0x5]),
3782                    SparseNode::new_branch(TrieMask::new(0b0011)),
3783                ),
3784                (
3785                    Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4, 0x5, 0x0]),
3786                    SparseNode::new_leaf(Nibbles::default()),
3787                ),
3788                (
3789                    Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4, 0x5, 0x1]),
3790                    SparseNode::new_leaf(Nibbles::default()),
3791                ),
3792            ]
3793            .into_iter(),
3794        );
3795
3796        let provider = MockTrieNodeProvider::new();
3797
3798        // Verify initial state - the lower subtrie's path should be 0x123
3799        let lower_subtrie_root_path = Nibbles::from_nibbles([0x1, 0x2, 0x3]);
3800        assert_matches!(
3801            trie.lower_subtrie_for_path_mut(&lower_subtrie_root_path),
3802            Some(subtrie)
3803            if subtrie.path == lower_subtrie_root_path
3804        );
3805
3806        // Remove the leaf at 0x1233
3807        let leaf_full_path = Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x3]);
3808        trie.remove_leaf(&leaf_full_path, provider).unwrap();
3809
3810        // After removal:
3811        // 1. The branch at 0x123 should become an extension to 0x12345
3812        // 2. That extension should merge with the root extension at 0x
3813        // 3. The lower subtrie's path should be updated to 0x12345
3814        let lower_subtrie = trie.lower_subtries[0x12].as_revealed_ref().unwrap();
3815        assert_eq!(lower_subtrie.path, Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4, 0x5]));
3816
3817        // Verify the root extension now points all the way to 0x12345
3818        assert_matches!(
3819            trie.upper_subtrie.nodes.get(&Nibbles::default()),
3820            Some(SparseNode::Extension { key, .. })
3821            if key == &Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4, 0x5])
3822        );
3823
3824        // Verify the branch at 0x12345 hasn't been modified
3825        assert_matches!(
3826            lower_subtrie.nodes.get(&Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4, 0x5])),
3827            Some(SparseNode::Branch { state_mask, .. })
3828            if state_mask == &TrieMask::new(0b0011)
3829        );
3830    }
3831
3832    #[test]
3833    fn test_remove_leaf_remaining_child_needs_reveal() {
3834        //
3835        // 0x:      Branch (Mask = 0011)
3836        // 0x0:     ├── 0 -> Leaf (Key = 12)
3837        // 0x1:     └── 1 -> Hash (blinded leaf)
3838        //
3839        // After removing 0x012, the hash node needs to be revealed to collapse the branch
3840        //
3841        let mut trie = new_test_trie(
3842            [
3843                (Nibbles::default(), SparseNode::new_branch(TrieMask::new(0b0011))),
3844                (
3845                    Nibbles::from_nibbles([0x0]),
3846                    SparseNode::new_leaf(Nibbles::from_nibbles([0x1, 0x2])),
3847                ),
3848                (Nibbles::from_nibbles([0x1]), SparseNode::Hash(B256::repeat_byte(0xab))),
3849            ]
3850            .into_iter(),
3851        );
3852
3853        // Create a mock provider that will reveal the blinded leaf
3854        let mut provider = MockTrieNodeProvider::new();
3855        let revealed_leaf = create_leaf_node([0x3, 0x4], 42);
3856        let mut encoded = Vec::new();
3857        revealed_leaf.encode(&mut encoded);
3858        provider.add_revealed_node(
3859            Nibbles::from_nibbles([0x1]),
3860            RevealedNode { node: encoded.into(), tree_mask: None, hash_mask: None },
3861        );
3862
3863        // Remove the leaf with a full path of 0x012
3864        let leaf_full_path = Nibbles::from_nibbles([0x0, 0x1, 0x2]);
3865        trie.remove_leaf(&leaf_full_path, provider).unwrap();
3866
3867        let upper_subtrie = &trie.upper_subtrie;
3868
3869        // Check that the leaf value was removed
3870        assert_matches!(upper_subtrie.inner.values.get(&leaf_full_path), None);
3871
3872        // Check that the branch node collapsed into a leaf node with the revealed child's key
3873        assert_matches!(
3874            upper_subtrie.nodes.get(&Nibbles::default()),
3875            Some(SparseNode::Leaf{ key, ..})
3876            if key == &Nibbles::from_nibbles([0x1, 0x3, 0x4])
3877        );
3878
3879        // Check that the remaining child node was removed (since it was merged)
3880        assert_matches!(upper_subtrie.nodes.get(&Nibbles::from_nibbles([0x1])), None);
3881    }
3882
3883    #[test]
3884    fn test_remove_leaf_root() {
3885        //
3886        // 0x:      Leaf (Key = 123)
3887        //
3888        // After removing 0x123, the trie becomes empty
3889        //
3890        let mut trie = new_test_trie(std::iter::once((
3891            Nibbles::default(),
3892            SparseNode::new_leaf(Nibbles::from_nibbles([0x1, 0x2, 0x3])),
3893        )));
3894
3895        let provider = MockTrieNodeProvider::new();
3896
3897        // Remove the leaf with a full key of 0x123
3898        let leaf_full_path = Nibbles::from_nibbles([0x1, 0x2, 0x3]);
3899        trie.remove_leaf(&leaf_full_path, provider).unwrap();
3900
3901        let upper_subtrie = &trie.upper_subtrie;
3902
3903        // Check that the leaf value was removed
3904        assert_matches!(upper_subtrie.inner.values.get(&leaf_full_path), None);
3905
3906        // Check that the root node was changed to Empty
3907        assert_matches!(upper_subtrie.nodes.get(&Nibbles::default()), Some(SparseNode::Empty));
3908    }
3909
3910    #[test]
3911    fn test_remove_leaf_unsets_hash_along_path() {
3912        //
3913        // Creates a trie structure:
3914        // 0x:      Branch (with hash set)
3915        // 0x0:     ├── Extension (with hash set)
3916        // 0x01:    │   └── Branch (with hash set)
3917        // 0x012:   │       ├── Leaf (Key = 34, with hash set)
3918        // 0x013:   │       ├── Leaf (Key = 56, with hash set)
3919        // 0x014:   │       └── Leaf (Key = 78, with hash set)
3920        // 0x1:     └── Leaf (Key = 78, with hash set)
3921        //
3922        // When removing leaf at 0x01234, all nodes along the path (root branch,
3923        // extension at 0x0, branch at 0x01) should have their hash field unset
3924        //
3925
3926        let mut trie = new_test_trie(
3927            [
3928                (
3929                    Nibbles::default(),
3930                    SparseNode::Branch {
3931                        state_mask: TrieMask::new(0b0011),
3932                        hash: Some(B256::repeat_byte(0x10)),
3933                        store_in_db_trie: None,
3934                    },
3935                ),
3936                (
3937                    Nibbles::from_nibbles([0x0]),
3938                    SparseNode::Extension {
3939                        key: Nibbles::from_nibbles([0x1]),
3940                        hash: Some(B256::repeat_byte(0x20)),
3941                        store_in_db_trie: None,
3942                    },
3943                ),
3944                (
3945                    Nibbles::from_nibbles([0x0, 0x1]),
3946                    SparseNode::Branch {
3947                        state_mask: TrieMask::new(0b11100),
3948                        hash: Some(B256::repeat_byte(0x30)),
3949                        store_in_db_trie: None,
3950                    },
3951                ),
3952                (
3953                    Nibbles::from_nibbles([0x0, 0x1, 0x2]),
3954                    SparseNode::Leaf {
3955                        key: Nibbles::from_nibbles([0x3, 0x4]),
3956                        hash: Some(B256::repeat_byte(0x40)),
3957                    },
3958                ),
3959                (
3960                    Nibbles::from_nibbles([0x0, 0x1, 0x3]),
3961                    SparseNode::Leaf {
3962                        key: Nibbles::from_nibbles([0x5, 0x6]),
3963                        hash: Some(B256::repeat_byte(0x50)),
3964                    },
3965                ),
3966                (
3967                    Nibbles::from_nibbles([0x0, 0x1, 0x4]),
3968                    SparseNode::Leaf {
3969                        key: Nibbles::from_nibbles([0x6, 0x7]),
3970                        hash: Some(B256::repeat_byte(0x60)),
3971                    },
3972                ),
3973                (
3974                    Nibbles::from_nibbles([0x1]),
3975                    SparseNode::Leaf {
3976                        key: Nibbles::from_nibbles([0x7, 0x8]),
3977                        hash: Some(B256::repeat_byte(0x70)),
3978                    },
3979                ),
3980            ]
3981            .into_iter(),
3982        );
3983
3984        let provider = MockTrieNodeProvider::new();
3985
3986        // Remove a leaf which does not exist; this should have no effect.
3987        trie.remove_leaf(&Nibbles::from_nibbles([0x0, 0x1, 0x2, 0x3, 0x4, 0xF]), &provider)
3988            .unwrap();
3989        for (path, node) in trie.all_nodes() {
3990            assert!(node.hash().is_some(), "path {path:?} should still have a hash");
3991        }
3992
3993        // Remove the leaf at path 0x01234
3994        let leaf_full_path = Nibbles::from_nibbles([0x0, 0x1, 0x2, 0x3, 0x4]);
3995        trie.remove_leaf(&leaf_full_path, &provider).unwrap();
3996
3997        let upper_subtrie = &trie.upper_subtrie;
3998        let lower_subtrie_10 = trie.lower_subtries[0x01].as_revealed_ref().unwrap();
3999
4000        // Verify that hash fields are unset for all nodes along the path to the removed leaf
4001        assert_matches!(
4002            upper_subtrie.nodes.get(&Nibbles::default()),
4003            Some(SparseNode::Branch { hash: None, .. })
4004        );
4005        assert_matches!(
4006            upper_subtrie.nodes.get(&Nibbles::from_nibbles([0x0])),
4007            Some(SparseNode::Extension { hash: None, .. })
4008        );
4009        assert_matches!(
4010            lower_subtrie_10.nodes.get(&Nibbles::from_nibbles([0x0, 0x1])),
4011            Some(SparseNode::Branch { hash: None, .. })
4012        );
4013
4014        // Verify that nodes not on the path still have their hashes
4015        assert_matches!(
4016            upper_subtrie.nodes.get(&Nibbles::from_nibbles([0x1])),
4017            Some(SparseNode::Leaf { hash: Some(_), .. })
4018        );
4019        assert_matches!(
4020            lower_subtrie_10.nodes.get(&Nibbles::from_nibbles([0x0, 0x1, 0x3])),
4021            Some(SparseNode::Leaf { hash: Some(_), .. })
4022        );
4023        assert_matches!(
4024            lower_subtrie_10.nodes.get(&Nibbles::from_nibbles([0x0, 0x1, 0x4])),
4025            Some(SparseNode::Leaf { hash: Some(_), .. })
4026        );
4027    }
4028
4029    #[test]
4030    fn test_parallel_sparse_trie_root() {
4031        // Step 1: Create the trie structure
4032        // Extension node at 0x with key 0x2 (goes to upper subtrie)
4033        let extension_path = Nibbles::new();
4034        let extension_key = Nibbles::from_nibbles([0x2]);
4035
4036        // Branch node at 0x2 with children 0 and 1 (goes to upper subtrie)
4037        let branch_path = Nibbles::from_nibbles([0x2]);
4038
4039        // Leaf nodes at 0x20 and 0x21 (go to lower subtries)
4040        let leaf_1_path = Nibbles::from_nibbles([0x2, 0x0]);
4041        let leaf_1_key = Nibbles::from_nibbles(vec![0; 62]); // Remaining key
4042        let leaf_1_full_path = Nibbles::from_nibbles([vec![0x2, 0x0], vec![0; 62]].concat());
4043
4044        let leaf_2_path = Nibbles::from_nibbles([0x2, 0x1]);
4045        let leaf_2_key = Nibbles::from_nibbles(vec![0; 62]); // Remaining key
4046        let leaf_2_full_path = Nibbles::from_nibbles([vec![0x2, 0x1], vec![0; 62]].concat());
4047
4048        // Create accounts
4049        let account_1 = create_account(1);
4050        let account_2 = create_account(2);
4051
4052        // Create leaf nodes
4053        let leaf_1 = create_leaf_node(leaf_1_key.to_vec(), account_1.nonce);
4054        let leaf_2 = create_leaf_node(leaf_2_key.to_vec(), account_2.nonce);
4055
4056        // Create branch node with children at indices 0 and 1
4057        let branch = create_branch_node_with_children(
4058            &[0, 1],
4059            vec![
4060                RlpNode::from_rlp(&alloy_rlp::encode(&leaf_1)),
4061                RlpNode::from_rlp(&alloy_rlp::encode(&leaf_2)),
4062            ],
4063        );
4064
4065        // Create extension node pointing to branch
4066        let extension = create_extension_node(
4067            extension_key.to_vec(),
4068            RlpNode::from_rlp(&alloy_rlp::encode(&branch)).as_hash().unwrap(),
4069        );
4070
4071        // Step 2: Reveal nodes in the trie
4072        let mut trie = ParallelSparseTrie::from_root(extension, TrieMasks::none(), true).unwrap();
4073        trie.reveal_nodes(vec![
4074            RevealedSparseNode { path: branch_path, node: branch, masks: TrieMasks::none() },
4075            RevealedSparseNode { path: leaf_1_path, node: leaf_1, masks: TrieMasks::none() },
4076            RevealedSparseNode { path: leaf_2_path, node: leaf_2, masks: TrieMasks::none() },
4077        ])
4078        .unwrap();
4079
4080        // Step 3: Reset hashes for all revealed nodes to test actual hash calculation
4081        // Reset upper subtrie node hashes
4082        trie.upper_subtrie.nodes.get_mut(&extension_path).unwrap().set_hash(None);
4083        trie.upper_subtrie.nodes.get_mut(&branch_path).unwrap().set_hash(None);
4084
4085        // Reset lower subtrie node hashes
4086        let leaf_1_subtrie_idx = path_subtrie_index_unchecked(&leaf_1_path);
4087        let leaf_2_subtrie_idx = path_subtrie_index_unchecked(&leaf_2_path);
4088
4089        trie.lower_subtries[leaf_1_subtrie_idx]
4090            .as_revealed_mut()
4091            .unwrap()
4092            .nodes
4093            .get_mut(&leaf_1_path)
4094            .unwrap()
4095            .set_hash(None);
4096        trie.lower_subtries[leaf_2_subtrie_idx]
4097            .as_revealed_mut()
4098            .unwrap()
4099            .nodes
4100            .get_mut(&leaf_2_path)
4101            .unwrap()
4102            .set_hash(None);
4103
4104        // Step 4: Add changed leaf node paths to prefix set
4105        trie.prefix_set.insert(leaf_1_full_path);
4106        trie.prefix_set.insert(leaf_2_full_path);
4107
4108        // Step 5: Calculate root using our implementation
4109        let root = trie.root();
4110
4111        // Step 6: Calculate root using HashBuilder for comparison
4112        let (hash_builder_root, _, _proof_nodes, _, _) = run_hash_builder(
4113            [(leaf_1_full_path, account_1), (leaf_2_full_path, account_2)],
4114            NoopAccountTrieCursor::default(),
4115            Default::default(),
4116            [extension_path, branch_path, leaf_1_full_path, leaf_2_full_path],
4117        );
4118
4119        // Step 7: Verify the roots match
4120        assert_eq!(root, hash_builder_root);
4121
4122        // Verify hashes were computed
4123        let leaf_1_subtrie = trie.lower_subtries[leaf_1_subtrie_idx].as_revealed_ref().unwrap();
4124        let leaf_2_subtrie = trie.lower_subtries[leaf_2_subtrie_idx].as_revealed_ref().unwrap();
4125        assert!(trie.upper_subtrie.nodes.get(&extension_path).unwrap().hash().is_some());
4126        assert!(trie.upper_subtrie.nodes.get(&branch_path).unwrap().hash().is_some());
4127        assert!(leaf_1_subtrie.nodes.get(&leaf_1_path).unwrap().hash().is_some());
4128        assert!(leaf_2_subtrie.nodes.get(&leaf_2_path).unwrap().hash().is_some());
4129    }
4130
4131    #[test]
4132    fn sparse_trie_empty_update_one() {
4133        let ctx = ParallelSparseTrieTestContext;
4134
4135        let key = Nibbles::unpack(B256::with_last_byte(42));
4136        let value = || Account::default();
4137        let value_encoded = || {
4138            let mut account_rlp = Vec::new();
4139            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
4140            account_rlp
4141        };
4142
4143        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
4144            run_hash_builder(
4145                [(key, value())],
4146                NoopAccountTrieCursor::default(),
4147                Default::default(),
4148                [key],
4149            );
4150
4151        let mut sparse = ParallelSparseTrie::default().with_updates(true);
4152        ctx.update_leaves(&mut sparse, [(key, value_encoded())]);
4153        ctx.assert_with_hash_builder(
4154            &mut sparse,
4155            hash_builder_root,
4156            hash_builder_updates,
4157            hash_builder_proof_nodes,
4158        );
4159    }
4160
4161    #[test]
4162    fn sparse_trie_empty_update_multiple_lower_nibbles() {
4163        let ctx = ParallelSparseTrieTestContext;
4164
4165        let paths = (0..=16).map(|b| Nibbles::unpack(B256::with_last_byte(b))).collect::<Vec<_>>();
4166        let value = || Account::default();
4167        let value_encoded = || {
4168            let mut account_rlp = Vec::new();
4169            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
4170            account_rlp
4171        };
4172
4173        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
4174            run_hash_builder(
4175                paths.iter().copied().zip(std::iter::repeat_with(value)),
4176                NoopAccountTrieCursor::default(),
4177                Default::default(),
4178                paths.clone(),
4179            );
4180
4181        let mut sparse = ParallelSparseTrie::default().with_updates(true);
4182        ctx.update_leaves(
4183            &mut sparse,
4184            paths.into_iter().zip(std::iter::repeat_with(value_encoded)),
4185        );
4186
4187        ctx.assert_with_hash_builder(
4188            &mut sparse,
4189            hash_builder_root,
4190            hash_builder_updates,
4191            hash_builder_proof_nodes,
4192        );
4193    }
4194
4195    #[test]
4196    fn sparse_trie_empty_update_multiple_upper_nibbles() {
4197        let paths = (239..=255).map(|b| Nibbles::unpack(B256::repeat_byte(b))).collect::<Vec<_>>();
4198        let value = || Account::default();
4199        let value_encoded = || {
4200            let mut account_rlp = Vec::new();
4201            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
4202            account_rlp
4203        };
4204
4205        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
4206            run_hash_builder(
4207                paths.iter().copied().zip(std::iter::repeat_with(value)),
4208                NoopAccountTrieCursor::default(),
4209                Default::default(),
4210                paths.clone(),
4211            );
4212
4213        let provider = DefaultTrieNodeProvider;
4214        let mut sparse = ParallelSparseTrie::default().with_updates(true);
4215        for path in &paths {
4216            sparse.update_leaf(*path, value_encoded(), &provider).unwrap();
4217        }
4218        let sparse_root = sparse.root();
4219        let sparse_updates = sparse.take_updates();
4220
4221        assert_eq!(sparse_root, hash_builder_root);
4222        assert_eq!(sparse_updates.updated_nodes, hash_builder_updates.account_nodes);
4223        assert_eq_parallel_sparse_trie_proof_nodes(&sparse, hash_builder_proof_nodes);
4224    }
4225
4226    #[test]
4227    fn sparse_trie_empty_update_multiple() {
4228        let ctx = ParallelSparseTrieTestContext;
4229
4230        let paths = (0..=255)
4231            .map(|b| {
4232                Nibbles::unpack(if b % 2 == 0 {
4233                    B256::repeat_byte(b)
4234                } else {
4235                    B256::with_last_byte(b)
4236                })
4237            })
4238            .collect::<Vec<_>>();
4239        let value = || Account::default();
4240        let value_encoded = || {
4241            let mut account_rlp = Vec::new();
4242            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
4243            account_rlp
4244        };
4245
4246        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
4247            run_hash_builder(
4248                paths.iter().sorted_unstable().copied().zip(std::iter::repeat_with(value)),
4249                NoopAccountTrieCursor::default(),
4250                Default::default(),
4251                paths.clone(),
4252            );
4253
4254        let mut sparse = ParallelSparseTrie::default().with_updates(true);
4255        ctx.update_leaves(
4256            &mut sparse,
4257            paths.iter().copied().zip(std::iter::repeat_with(value_encoded)),
4258        );
4259        ctx.assert_with_hash_builder(
4260            &mut sparse,
4261            hash_builder_root,
4262            hash_builder_updates,
4263            hash_builder_proof_nodes,
4264        );
4265    }
4266
4267    #[test]
4268    fn sparse_trie_empty_update_repeated() {
4269        let ctx = ParallelSparseTrieTestContext;
4270
4271        let paths = (0..=255).map(|b| Nibbles::unpack(B256::repeat_byte(b))).collect::<Vec<_>>();
4272        let old_value = Account { nonce: 1, ..Default::default() };
4273        let old_value_encoded = {
4274            let mut account_rlp = Vec::new();
4275            old_value.into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
4276            account_rlp
4277        };
4278        let new_value = Account { nonce: 2, ..Default::default() };
4279        let new_value_encoded = {
4280            let mut account_rlp = Vec::new();
4281            new_value.into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
4282            account_rlp
4283        };
4284
4285        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
4286            run_hash_builder(
4287                paths.iter().copied().zip(std::iter::repeat_with(|| old_value)),
4288                NoopAccountTrieCursor::default(),
4289                Default::default(),
4290                paths.clone(),
4291            );
4292
4293        let mut sparse = ParallelSparseTrie::default().with_updates(true);
4294        ctx.update_leaves(
4295            &mut sparse,
4296            paths.iter().copied().zip(std::iter::repeat(old_value_encoded)),
4297        );
4298        ctx.assert_with_hash_builder(
4299            &mut sparse,
4300            hash_builder_root,
4301            hash_builder_updates,
4302            hash_builder_proof_nodes,
4303        );
4304
4305        let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
4306            run_hash_builder(
4307                paths.iter().copied().zip(std::iter::repeat(new_value)),
4308                NoopAccountTrieCursor::default(),
4309                Default::default(),
4310                paths.clone(),
4311            );
4312
4313        ctx.update_leaves(
4314            &mut sparse,
4315            paths.iter().copied().zip(std::iter::repeat(new_value_encoded)),
4316        );
4317        ctx.assert_with_hash_builder(
4318            &mut sparse,
4319            hash_builder_root,
4320            hash_builder_updates,
4321            hash_builder_proof_nodes,
4322        );
4323    }
4324
4325    #[test]
4326    fn sparse_trie_remove_leaf() {
4327        let ctx = ParallelSparseTrieTestContext;
4328        let provider = DefaultTrieNodeProvider;
4329        let mut sparse = ParallelSparseTrie::default();
4330
4331        let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec();
4332
4333        ctx.update_leaves(
4334            &mut sparse,
4335            [
4336                (Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]), value.clone()),
4337                (Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]), value.clone()),
4338                (Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3]), value.clone()),
4339                (Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2]), value.clone()),
4340                (Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]), value.clone()),
4341                (Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]), value),
4342            ],
4343        );
4344
4345        // Extension (Key = 5)
4346        // └── Branch (Mask = 1011)
4347        //     ├── 0 -> Extension (Key = 23)
4348        //     │        └── Branch (Mask = 0101)
4349        //     │              ├── 1 -> Leaf (Key = 1, Path = 50231)
4350        //     │              └── 3 -> Leaf (Key = 3, Path = 50233)
4351        //     ├── 2 -> Leaf (Key = 013, Path = 52013)
4352        //     └── 3 -> Branch (Mask = 0101)
4353        //                ├── 1 -> Leaf (Key = 3102, Path = 53102)
4354        //                └── 3 -> Branch (Mask = 1010)
4355        //                       ├── 0 -> Leaf (Key = 3302, Path = 53302)
4356        //                       └── 2 -> Leaf (Key = 3320, Path = 53320)
4357        pretty_assertions::assert_eq!(
4358            parallel_sparse_trie_nodes(&sparse)
4359                .into_iter()
4360                .map(|(k, v)| (*k, v.clone()))
4361                .collect::<BTreeMap<_, _>>(),
4362            BTreeMap::from_iter([
4363                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
4364                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1101.into())),
4365                (
4366                    Nibbles::from_nibbles([0x5, 0x0]),
4367                    SparseNode::new_ext(Nibbles::from_nibbles([0x2, 0x3]))
4368                ),
4369                (
4370                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3]),
4371                    SparseNode::new_branch(0b1010.into())
4372                ),
4373                (
4374                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]),
4375                    SparseNode::new_leaf(Nibbles::default())
4376                ),
4377                (
4378                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
4379                    SparseNode::new_leaf(Nibbles::default())
4380                ),
4381                (
4382                    Nibbles::from_nibbles([0x5, 0x2]),
4383                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0, 0x1, 0x3]))
4384                ),
4385                (Nibbles::from_nibbles([0x5, 0x3]), SparseNode::new_branch(0b1010.into())),
4386                (
4387                    Nibbles::from_nibbles([0x5, 0x3, 0x1]),
4388                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0, 0x2]))
4389                ),
4390                (Nibbles::from_nibbles([0x5, 0x3, 0x3]), SparseNode::new_branch(0b0101.into())),
4391                (
4392                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0]),
4393                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2]))
4394                ),
4395                (
4396                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2]),
4397                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0]))
4398                )
4399            ])
4400        );
4401
4402        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3]), &provider).unwrap();
4403
4404        // Extension (Key = 5)
4405        // └── Branch (Mask = 1001)
4406        //     ├── 0 -> Extension (Key = 23)
4407        //     │        └── Branch (Mask = 0101)
4408        //     │              ├── 1 -> Leaf (Key = 0231, Path = 50231)
4409        //     │              └── 3 -> Leaf (Key = 0233, Path = 50233)
4410        //     └── 3 -> Branch (Mask = 0101)
4411        //                ├── 1 -> Leaf (Key = 3102, Path = 53102)
4412        //                └── 3 -> Branch (Mask = 1010)
4413        //                       ├── 0 -> Leaf (Key = 3302, Path = 53302)
4414        //                       └── 2 -> Leaf (Key = 3320, Path = 53320)
4415        pretty_assertions::assert_eq!(
4416            parallel_sparse_trie_nodes(&sparse)
4417                .into_iter()
4418                .map(|(k, v)| (*k, v.clone()))
4419                .collect::<BTreeMap<_, _>>(),
4420            BTreeMap::from_iter([
4421                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
4422                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
4423                (
4424                    Nibbles::from_nibbles([0x5, 0x0]),
4425                    SparseNode::new_ext(Nibbles::from_nibbles([0x2, 0x3]))
4426                ),
4427                (
4428                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3]),
4429                    SparseNode::new_branch(0b1010.into())
4430                ),
4431                (
4432                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]),
4433                    SparseNode::new_leaf(Nibbles::default())
4434                ),
4435                (
4436                    Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]),
4437                    SparseNode::new_leaf(Nibbles::default())
4438                ),
4439                (Nibbles::from_nibbles([0x5, 0x3]), SparseNode::new_branch(0b1010.into())),
4440                (
4441                    Nibbles::from_nibbles([0x5, 0x3, 0x1]),
4442                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0, 0x2]))
4443                ),
4444                (Nibbles::from_nibbles([0x5, 0x3, 0x3]), SparseNode::new_branch(0b0101.into())),
4445                (
4446                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0]),
4447                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2]))
4448                ),
4449                (
4450                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2]),
4451                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0]))
4452                )
4453            ])
4454        );
4455
4456        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]), &provider).unwrap();
4457
4458        // Extension (Key = 5)
4459        // └── Branch (Mask = 1001)
4460        //     ├── 0 -> Leaf (Key = 0233, Path = 50233)
4461        //     └── 3 -> Branch (Mask = 0101)
4462        //                ├── 1 -> Leaf (Key = 3102, Path = 53102)
4463        //                └── 3 -> Branch (Mask = 1010)
4464        //                       ├── 0 -> Leaf (Key = 3302, Path = 53302)
4465        //                       └── 2 -> Leaf (Key = 3320, Path = 53320)
4466        pretty_assertions::assert_eq!(
4467            parallel_sparse_trie_nodes(&sparse)
4468                .into_iter()
4469                .map(|(k, v)| (*k, v.clone()))
4470                .collect::<BTreeMap<_, _>>(),
4471            BTreeMap::from_iter([
4472                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
4473                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
4474                (
4475                    Nibbles::from_nibbles([0x5, 0x0]),
4476                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2, 0x3, 0x3]))
4477                ),
4478                (Nibbles::from_nibbles([0x5, 0x3]), SparseNode::new_branch(0b1010.into())),
4479                (
4480                    Nibbles::from_nibbles([0x5, 0x3, 0x1]),
4481                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0, 0x2]))
4482                ),
4483                (Nibbles::from_nibbles([0x5, 0x3, 0x3]), SparseNode::new_branch(0b0101.into())),
4484                (
4485                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0]),
4486                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2]))
4487                ),
4488                (
4489                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2]),
4490                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0]))
4491                )
4492            ])
4493        );
4494
4495        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2]), &provider).unwrap();
4496
4497        // Extension (Key = 5)
4498        // └── Branch (Mask = 1001)
4499        //     ├── 0 -> Leaf (Key = 0233, Path = 50233)
4500        //     └── 3 -> Branch (Mask = 1010)
4501        //                ├── 0 -> Leaf (Key = 3302, Path = 53302)
4502        //                └── 2 -> Leaf (Key = 3320, Path = 53320)
4503        pretty_assertions::assert_eq!(
4504            parallel_sparse_trie_nodes(&sparse)
4505                .into_iter()
4506                .map(|(k, v)| (*k, v.clone()))
4507                .collect::<BTreeMap<_, _>>(),
4508            BTreeMap::from_iter([
4509                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
4510                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
4511                (
4512                    Nibbles::from_nibbles([0x5, 0x0]),
4513                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2, 0x3, 0x3]))
4514                ),
4515                (
4516                    Nibbles::from_nibbles([0x5, 0x3]),
4517                    SparseNode::new_ext(Nibbles::from_nibbles([0x3]))
4518                ),
4519                (Nibbles::from_nibbles([0x5, 0x3, 0x3]), SparseNode::new_branch(0b0101.into())),
4520                (
4521                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0]),
4522                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2]))
4523                ),
4524                (
4525                    Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2]),
4526                    SparseNode::new_leaf(Nibbles::from_nibbles([0x0]))
4527                )
4528            ])
4529        );
4530
4531        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]), &provider).unwrap();
4532
4533        // Extension (Key = 5)
4534        // └── Branch (Mask = 1001)
4535        //     ├── 0 -> Leaf (Key = 0233, Path = 50233)
4536        //     └── 3 -> Leaf (Key = 3302, Path = 53302)
4537        pretty_assertions::assert_eq!(
4538            parallel_sparse_trie_nodes(&sparse)
4539                .into_iter()
4540                .map(|(k, v)| (*k, v.clone()))
4541                .collect::<BTreeMap<_, _>>(),
4542            BTreeMap::from_iter([
4543                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0x5]))),
4544                (Nibbles::from_nibbles([0x5]), SparseNode::new_branch(0b1001.into())),
4545                (
4546                    Nibbles::from_nibbles([0x5, 0x0]),
4547                    SparseNode::new_leaf(Nibbles::from_nibbles([0x2, 0x3, 0x3]))
4548                ),
4549                (
4550                    Nibbles::from_nibbles([0x5, 0x3]),
4551                    SparseNode::new_leaf(Nibbles::from_nibbles([0x3, 0x0, 0x2]))
4552                ),
4553            ])
4554        );
4555
4556        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]), &provider).unwrap();
4557
4558        // Leaf (Key = 53302)
4559        pretty_assertions::assert_eq!(
4560            parallel_sparse_trie_nodes(&sparse)
4561                .into_iter()
4562                .map(|(k, v)| (*k, v.clone()))
4563                .collect::<BTreeMap<_, _>>(),
4564            BTreeMap::from_iter([(
4565                Nibbles::default(),
4566                SparseNode::new_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]))
4567            ),])
4568        );
4569
4570        sparse.remove_leaf(&Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x0, 0x2]), &provider).unwrap();
4571
4572        // Empty
4573        pretty_assertions::assert_eq!(
4574            parallel_sparse_trie_nodes(&sparse)
4575                .into_iter()
4576                .map(|(k, v)| (*k, v.clone()))
4577                .collect::<BTreeMap<_, _>>(),
4578            BTreeMap::from_iter([(Nibbles::default(), SparseNode::Empty)])
4579        );
4580    }
4581
4582    #[test]
4583    fn sparse_trie_remove_leaf_blinded() {
4584        let leaf = LeafNode::new(
4585            Nibbles::default(),
4586            alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec(),
4587        );
4588        let branch = TrieNode::Branch(BranchNode::new(
4589            vec![
4590                RlpNode::word_rlp(&B256::repeat_byte(1)),
4591                RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(),
4592            ],
4593            TrieMask::new(0b11),
4594        ));
4595
4596        let provider = DefaultTrieNodeProvider;
4597        let mut sparse = ParallelSparseTrie::from_root(
4598            branch.clone(),
4599            TrieMasks { hash_mask: Some(TrieMask::new(0b01)), tree_mask: None },
4600            false,
4601        )
4602        .unwrap();
4603
4604        // Reveal a branch node and one of its children
4605        //
4606        // Branch (Mask = 11)
4607        // ├── 0 -> Hash (Path = 0)
4608        // └── 1 -> Leaf (Path = 1)
4609        sparse
4610            .reveal_nodes(vec![
4611                RevealedSparseNode {
4612                    path: Nibbles::default(),
4613                    node: branch,
4614                    masks: TrieMasks { hash_mask: None, tree_mask: Some(TrieMask::new(0b01)) },
4615                },
4616                RevealedSparseNode {
4617                    path: Nibbles::from_nibbles([0x1]),
4618                    node: TrieNode::Leaf(leaf),
4619                    masks: TrieMasks::none(),
4620                },
4621            ])
4622            .unwrap();
4623
4624        // Removing a blinded leaf should result in an error
4625        assert_matches!(
4626            sparse.remove_leaf(&Nibbles::from_nibbles([0x0]), &provider).map_err(|e| e.into_kind()),
4627            Err(SparseTrieErrorKind::BlindedNode { path, hash }) if path == Nibbles::from_nibbles([0x0]) && hash == B256::repeat_byte(1)
4628        );
4629    }
4630
4631    #[test]
4632    fn sparse_trie_remove_leaf_non_existent() {
4633        let leaf = LeafNode::new(
4634            Nibbles::default(),
4635            alloy_rlp::encode_fixed_size(&U256::from(1)).to_vec(),
4636        );
4637        let branch = TrieNode::Branch(BranchNode::new(
4638            vec![
4639                RlpNode::word_rlp(&B256::repeat_byte(1)),
4640                RlpNode::from_raw_rlp(&alloy_rlp::encode(leaf.clone())).unwrap(),
4641            ],
4642            TrieMask::new(0b11),
4643        ));
4644
4645        let provider = DefaultTrieNodeProvider;
4646        let mut sparse = ParallelSparseTrie::from_root(
4647            branch.clone(),
4648            TrieMasks { hash_mask: Some(TrieMask::new(0b01)), tree_mask: None },
4649            false,
4650        )
4651        .unwrap();
4652
4653        // Reveal a branch node and one of its children
4654        //
4655        // Branch (Mask = 11)
4656        // ├── 0 -> Hash (Path = 0)
4657        // └── 1 -> Leaf (Path = 1)
4658        sparse
4659            .reveal_nodes(vec![
4660                RevealedSparseNode {
4661                    path: Nibbles::default(),
4662                    node: branch,
4663                    masks: TrieMasks { hash_mask: None, tree_mask: Some(TrieMask::new(0b01)) },
4664                },
4665                RevealedSparseNode {
4666                    path: Nibbles::from_nibbles([0x1]),
4667                    node: TrieNode::Leaf(leaf),
4668                    masks: TrieMasks::none(),
4669                },
4670            ])
4671            .unwrap();
4672
4673        // Removing a non-existent leaf should be a noop
4674        let sparse_old = sparse.clone();
4675        assert_matches!(sparse.remove_leaf(&Nibbles::from_nibbles([0x2]), &provider), Ok(()));
4676        assert_eq!(sparse, sparse_old);
4677    }
4678
4679    #[test]
4680    fn sparse_trie_fuzz() {
4681        // Having only the first 3 nibbles set, we narrow down the range of keys
4682        // to 4096 different hashes. It allows us to generate collisions more likely
4683        // to test the sparse trie updates.
4684        const KEY_NIBBLES_LEN: usize = 3;
4685
4686        fn test(updates: Vec<(BTreeMap<Nibbles, Account>, BTreeSet<Nibbles>)>) {
4687            {
4688                let mut state = BTreeMap::default();
4689                let default_provider = DefaultTrieNodeProvider;
4690                let provider_factory = create_test_provider_factory();
4691                let mut sparse = ParallelSparseTrie::default().with_updates(true);
4692
4693                for (update, keys_to_delete) in updates {
4694                    // Insert state updates into the sparse trie and calculate the root
4695                    for (key, account) in update.clone() {
4696                        let account = account.into_trie_account(EMPTY_ROOT_HASH);
4697                        let mut account_rlp = Vec::new();
4698                        account.encode(&mut account_rlp);
4699                        sparse.update_leaf(key, account_rlp, &default_provider).unwrap();
4700                    }
4701                    // We need to clone the sparse trie, so that all updated branch nodes are
4702                    // preserved, and not only those that were changed after the last call to
4703                    // `root()`.
4704                    let mut updated_sparse = sparse.clone();
4705                    let sparse_root = updated_sparse.root();
4706                    let sparse_updates = updated_sparse.take_updates();
4707
4708                    // Insert state updates into the hash builder and calculate the root
4709                    state.extend(update);
4710                    let provider = provider_factory.provider().unwrap();
4711                    let trie_cursor = DatabaseTrieCursorFactory::new(provider.tx_ref());
4712                    let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
4713                        run_hash_builder(
4714                            state.clone(),
4715                            trie_cursor.account_trie_cursor().unwrap(),
4716                            Default::default(),
4717                            state.keys().copied().collect::<Vec<_>>(),
4718                        );
4719
4720                    // Write trie updates to the database
4721                    let provider_rw = provider_factory.provider_rw().unwrap();
4722                    provider_rw.write_trie_updates(&hash_builder_updates).unwrap();
4723                    provider_rw.commit().unwrap();
4724
4725                    // Assert that the sparse trie root matches the hash builder root
4726                    assert_eq!(sparse_root, hash_builder_root);
4727                    // Assert that the sparse trie updates match the hash builder updates
4728                    pretty_assertions::assert_eq!(
4729                        BTreeMap::from_iter(sparse_updates.updated_nodes),
4730                        BTreeMap::from_iter(hash_builder_updates.account_nodes)
4731                    );
4732                    // Assert that the sparse trie nodes match the hash builder proof nodes
4733                    assert_eq_parallel_sparse_trie_proof_nodes(
4734                        &updated_sparse,
4735                        hash_builder_proof_nodes,
4736                    );
4737
4738                    // Delete some keys from both the hash builder and the sparse trie and check
4739                    // that the sparse trie root still matches the hash builder root
4740                    for key in &keys_to_delete {
4741                        state.remove(key).unwrap();
4742                        sparse.remove_leaf(key, &default_provider).unwrap();
4743                    }
4744
4745                    // We need to clone the sparse trie, so that all updated branch nodes are
4746                    // preserved, and not only those that were changed after the last call to
4747                    // `root()`.
4748                    let mut updated_sparse = sparse.clone();
4749                    let sparse_root = updated_sparse.root();
4750                    let sparse_updates = updated_sparse.take_updates();
4751
4752                    let provider = provider_factory.provider().unwrap();
4753                    let trie_cursor = DatabaseTrieCursorFactory::new(provider.tx_ref());
4754                    let (hash_builder_root, hash_builder_updates, hash_builder_proof_nodes, _, _) =
4755                        run_hash_builder(
4756                            state.clone(),
4757                            trie_cursor.account_trie_cursor().unwrap(),
4758                            keys_to_delete
4759                                .iter()
4760                                .map(|nibbles| B256::from_slice(&nibbles.pack()))
4761                                .collect(),
4762                            state.keys().copied().collect::<Vec<_>>(),
4763                        );
4764
4765                    // Write trie updates to the database
4766                    let provider_rw = provider_factory.provider_rw().unwrap();
4767                    provider_rw.write_trie_updates(&hash_builder_updates).unwrap();
4768                    provider_rw.commit().unwrap();
4769
4770                    // Assert that the sparse trie root matches the hash builder root
4771                    assert_eq!(sparse_root, hash_builder_root);
4772                    // Assert that the sparse trie updates match the hash builder updates
4773                    pretty_assertions::assert_eq!(
4774                        BTreeMap::from_iter(sparse_updates.updated_nodes),
4775                        BTreeMap::from_iter(hash_builder_updates.account_nodes)
4776                    );
4777                    // Assert that the sparse trie nodes match the hash builder proof nodes
4778                    assert_eq_parallel_sparse_trie_proof_nodes(
4779                        &updated_sparse,
4780                        hash_builder_proof_nodes,
4781                    );
4782                }
4783            }
4784        }
4785
4786        fn transform_updates(
4787            updates: Vec<BTreeMap<Nibbles, Account>>,
4788            mut rng: impl rand::Rng,
4789        ) -> Vec<(BTreeMap<Nibbles, Account>, BTreeSet<Nibbles>)> {
4790            let mut keys = BTreeSet::new();
4791            updates
4792                .into_iter()
4793                .map(|update| {
4794                    keys.extend(update.keys().copied());
4795
4796                    let keys_to_delete_len = update.len() / 2;
4797                    let keys_to_delete = (0..keys_to_delete_len)
4798                        .map(|_| {
4799                            let key =
4800                                *rand::seq::IteratorRandom::choose(keys.iter(), &mut rng).unwrap();
4801                            keys.take(&key).unwrap()
4802                        })
4803                        .collect();
4804
4805                    (update, keys_to_delete)
4806                })
4807                .collect::<Vec<_>>()
4808        }
4809
4810        proptest!(ProptestConfig::with_cases(10), |(
4811            updates in proptest::collection::vec(
4812                proptest::collection::btree_map(
4813                    any_with::<Nibbles>(SizeRange::new(KEY_NIBBLES_LEN..=KEY_NIBBLES_LEN)).prop_map(pad_nibbles_right),
4814                    arb::<Account>(),
4815                    1..50,
4816                ),
4817                1..50,
4818            ).prop_perturb(transform_updates)
4819        )| {
4820            test(updates)
4821        });
4822    }
4823
4824    #[test]
4825    fn sparse_trie_fuzz_vs_serial() {
4826        // Having only the first 3 nibbles set, we narrow down the range of keys
4827        // to 4096 different hashes. It allows us to generate collisions more likely
4828        // to test the sparse trie updates.
4829        const KEY_NIBBLES_LEN: usize = 3;
4830
4831        fn test(updates: Vec<(BTreeMap<Nibbles, Account>, BTreeSet<Nibbles>)>) {
4832            let default_provider = DefaultTrieNodeProvider;
4833            let mut serial = SerialSparseTrie::default().with_updates(true);
4834            let mut parallel = ParallelSparseTrie::default().with_updates(true);
4835
4836            for (update, keys_to_delete) in updates {
4837                // Perform leaf updates on both tries
4838                for (key, account) in update.clone() {
4839                    let account = account.into_trie_account(EMPTY_ROOT_HASH);
4840                    let mut account_rlp = Vec::new();
4841                    account.encode(&mut account_rlp);
4842                    serial.update_leaf(key, account_rlp.clone(), &default_provider).unwrap();
4843                    parallel.update_leaf(key, account_rlp, &default_provider).unwrap();
4844                }
4845
4846                // Calculate roots and assert their equality
4847                let serial_root = serial.root();
4848                let parallel_root = parallel.root();
4849                assert_eq!(parallel_root, serial_root);
4850
4851                // Assert that both tries produce the same updates
4852                let serial_updates = serial.take_updates();
4853                let parallel_updates = parallel.take_updates();
4854                pretty_assertions::assert_eq!(
4855                    BTreeMap::from_iter(parallel_updates.updated_nodes),
4856                    BTreeMap::from_iter(serial_updates.updated_nodes),
4857                );
4858                pretty_assertions::assert_eq!(
4859                    BTreeSet::from_iter(parallel_updates.removed_nodes),
4860                    BTreeSet::from_iter(serial_updates.removed_nodes),
4861                );
4862
4863                // Perform leaf removals on both tries
4864                for key in &keys_to_delete {
4865                    parallel.remove_leaf(key, &default_provider).unwrap();
4866                    serial.remove_leaf(key, &default_provider).unwrap();
4867                }
4868
4869                // Calculate roots and assert their equality
4870                let serial_root = serial.root();
4871                let parallel_root = parallel.root();
4872                assert_eq!(parallel_root, serial_root);
4873
4874                // Assert that both tries produce the same updates
4875                let serial_updates = serial.take_updates();
4876                let parallel_updates = parallel.take_updates();
4877                pretty_assertions::assert_eq!(
4878                    BTreeMap::from_iter(parallel_updates.updated_nodes),
4879                    BTreeMap::from_iter(serial_updates.updated_nodes),
4880                );
4881                pretty_assertions::assert_eq!(
4882                    BTreeSet::from_iter(parallel_updates.removed_nodes),
4883                    BTreeSet::from_iter(serial_updates.removed_nodes),
4884                );
4885            }
4886        }
4887
4888        fn transform_updates(
4889            updates: Vec<BTreeMap<Nibbles, Account>>,
4890            mut rng: impl rand::Rng,
4891        ) -> Vec<(BTreeMap<Nibbles, Account>, BTreeSet<Nibbles>)> {
4892            let mut keys = BTreeSet::new();
4893            updates
4894                .into_iter()
4895                .map(|update| {
4896                    keys.extend(update.keys().copied());
4897
4898                    let keys_to_delete_len = update.len() / 2;
4899                    let keys_to_delete = (0..keys_to_delete_len)
4900                        .map(|_| {
4901                            let key =
4902                                *rand::seq::IteratorRandom::choose(keys.iter(), &mut rng).unwrap();
4903                            keys.take(&key).unwrap()
4904                        })
4905                        .collect();
4906
4907                    (update, keys_to_delete)
4908                })
4909                .collect::<Vec<_>>()
4910        }
4911
4912        proptest!(ProptestConfig::with_cases(10), |(
4913            updates in proptest::collection::vec(
4914                proptest::collection::btree_map(
4915                    any_with::<Nibbles>(SizeRange::new(KEY_NIBBLES_LEN..=KEY_NIBBLES_LEN)).prop_map(pad_nibbles_right),
4916                    arb::<Account>(),
4917                    1..50,
4918                ),
4919                1..50,
4920            ).prop_perturb(transform_updates)
4921        )| {
4922            test(updates)
4923        });
4924    }
4925
4926    #[test]
4927    fn sparse_trie_two_leaves_at_lower_roots() {
4928        let provider = DefaultTrieNodeProvider;
4929        let mut trie = ParallelSparseTrie::default().with_updates(true);
4930        let key_50 = Nibbles::unpack(hex!(
4931            "0x5000000000000000000000000000000000000000000000000000000000000000"
4932        ));
4933        let key_51 = Nibbles::unpack(hex!(
4934            "0x5100000000000000000000000000000000000000000000000000000000000000"
4935        ));
4936
4937        let account = Account::default().into_trie_account(EMPTY_ROOT_HASH);
4938        let mut account_rlp = Vec::new();
4939        account.encode(&mut account_rlp);
4940
4941        // Add a leaf and calculate the root.
4942        trie.update_leaf(key_50, account_rlp.clone(), &provider).unwrap();
4943        trie.root();
4944
4945        // Add a second leaf and assert that the root is the expected value.
4946        trie.update_leaf(key_51, account_rlp.clone(), &provider).unwrap();
4947
4948        let expected_root =
4949            hex!("0xdaf0ef9f91a2f179bb74501209effdb5301db1697bcab041eca2234b126e25de");
4950        let root = trie.root();
4951        assert_eq!(root, expected_root);
4952        assert_eq!(SparseTrieUpdates::default(), trie.take_updates());
4953    }
4954
4955    /// We have three leaves that share the same prefix: 0x00, 0x01 and 0x02. Hash builder trie has
4956    /// only nodes 0x00 and 0x01, and we have proofs for them. Node B is new and inserted in the
4957    /// sparse trie first.
4958    ///
4959    /// 1. Reveal the hash builder proof to leaf 0x00 in the sparse trie.
4960    /// 2. Insert leaf 0x01 into the sparse trie.
4961    /// 3. Reveal the hash builder proof to leaf 0x02 in the sparse trie.
4962    ///
4963    /// The hash builder proof to the leaf 0x02 didn't have the leaf 0x01 at the corresponding
4964    /// nibble of the branch node, so we need to adjust the branch node instead of fully
4965    /// replacing it.
4966    #[test]
4967    fn sparse_trie_reveal_node_1() {
4968        let key1 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x00]));
4969        let key2 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x01]));
4970        let key3 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x02]));
4971        let value = || Account::default();
4972        let value_encoded = || {
4973            let mut account_rlp = Vec::new();
4974            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
4975            account_rlp
4976        };
4977
4978        // Generate the proof for the root node and initialize the sparse trie with it
4979        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
4980            run_hash_builder(
4981                [(key1(), value()), (key3(), value())],
4982                NoopAccountTrieCursor::default(),
4983                Default::default(),
4984                [Nibbles::default()],
4985            );
4986
4987        let provider = DefaultTrieNodeProvider;
4988        let mut sparse = ParallelSparseTrie::from_root(
4989            TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(),
4990            TrieMasks {
4991                hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(),
4992                tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(),
4993            },
4994            false,
4995        )
4996        .unwrap();
4997
4998        // Generate the proof for the first key and reveal it in the sparse trie
4999        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
5000            run_hash_builder(
5001                [(key1(), value()), (key3(), value())],
5002                NoopAccountTrieCursor::default(),
5003                Default::default(),
5004                [key1()],
5005            );
5006        let revealed_nodes: Vec<RevealedSparseNode> = hash_builder_proof_nodes
5007            .nodes_sorted()
5008            .into_iter()
5009            .map(|(path, node)| {
5010                let hash_mask = branch_node_hash_masks.get(&path).copied();
5011                let tree_mask = branch_node_tree_masks.get(&path).copied();
5012                RevealedSparseNode {
5013                    path,
5014                    node: TrieNode::decode(&mut &node[..]).unwrap(),
5015                    masks: TrieMasks { hash_mask, tree_mask },
5016                }
5017            })
5018            .collect();
5019        sparse.reveal_nodes(revealed_nodes).unwrap();
5020
5021        // Check that the branch node exists with only two nibbles set
5022        assert_eq!(
5023            sparse.upper_subtrie.nodes.get(&Nibbles::default()),
5024            Some(&SparseNode::new_branch(0b101.into()))
5025        );
5026
5027        // Insert the leaf for the second key
5028        sparse.update_leaf(key2(), value_encoded(), &provider).unwrap();
5029
5030        // Check that the branch node was updated and another nibble was set
5031        assert_eq!(
5032            sparse.upper_subtrie.nodes.get(&Nibbles::default()),
5033            Some(&SparseNode::new_branch(0b111.into()))
5034        );
5035
5036        // Generate the proof for the third key and reveal it in the sparse trie
5037        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
5038            run_hash_builder(
5039                [(key1(), value()), (key3(), value())],
5040                NoopAccountTrieCursor::default(),
5041                Default::default(),
5042                [key3()],
5043            );
5044        let revealed_nodes: Vec<RevealedSparseNode> = hash_builder_proof_nodes
5045            .nodes_sorted()
5046            .into_iter()
5047            .map(|(path, node)| {
5048                let hash_mask = branch_node_hash_masks.get(&path).copied();
5049                let tree_mask = branch_node_tree_masks.get(&path).copied();
5050                RevealedSparseNode {
5051                    path,
5052                    node: TrieNode::decode(&mut &node[..]).unwrap(),
5053                    masks: TrieMasks { hash_mask, tree_mask },
5054                }
5055            })
5056            .collect();
5057        sparse.reveal_nodes(revealed_nodes).unwrap();
5058
5059        // Check that nothing changed in the branch node
5060        assert_eq!(
5061            sparse.upper_subtrie.nodes.get(&Nibbles::default()),
5062            Some(&SparseNode::new_branch(0b111.into()))
5063        );
5064
5065        // Generate the nodes for the full trie with all three key using the hash builder, and
5066        // compare them to the sparse trie
5067        let (_, _, hash_builder_proof_nodes, _, _) = run_hash_builder(
5068            [(key1(), value()), (key2(), value()), (key3(), value())],
5069            NoopAccountTrieCursor::default(),
5070            Default::default(),
5071            [key1(), key2(), key3()],
5072        );
5073
5074        assert_eq_parallel_sparse_trie_proof_nodes(&sparse, hash_builder_proof_nodes);
5075    }
5076
5077    /// We have three leaves: 0x0000, 0x0101, and 0x0102. Hash builder trie has all nodes, and we
5078    /// have proofs for them.
5079    ///
5080    /// 1. Reveal the hash builder proof to leaf 0x00 in the sparse trie.
5081    /// 2. Remove leaf 0x00 from the sparse trie (that will remove the branch node and create an
5082    ///    extension node with the key 0x0000).
5083    /// 3. Reveal the hash builder proof to leaf 0x0101 in the sparse trie.
5084    ///
5085    /// The hash builder proof to the leaf 0x0101 had a branch node in the path, but we turned it
5086    /// into an extension node, so it should ignore this node.
5087    #[test]
5088    fn sparse_trie_reveal_node_2() {
5089        let key1 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x00, 0x00]));
5090        let key2 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x01, 0x01]));
5091        let key3 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x01, 0x02]));
5092        let value = || Account::default();
5093
5094        // Generate the proof for the root node and initialize the sparse trie with it
5095        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
5096            run_hash_builder(
5097                [(key1(), value()), (key2(), value()), (key3(), value())],
5098                NoopAccountTrieCursor::default(),
5099                Default::default(),
5100                [Nibbles::default()],
5101            );
5102
5103        let provider = DefaultTrieNodeProvider;
5104        let mut sparse = ParallelSparseTrie::from_root(
5105            TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(),
5106            TrieMasks {
5107                hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(),
5108                tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(),
5109            },
5110            false,
5111        )
5112        .unwrap();
5113
5114        // Generate the proof for the children of the root branch node and reveal it in the sparse
5115        // trie
5116        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
5117            run_hash_builder(
5118                [(key1(), value()), (key2(), value()), (key3(), value())],
5119                NoopAccountTrieCursor::default(),
5120                Default::default(),
5121                [key1(), Nibbles::from_nibbles_unchecked([0x01])],
5122            );
5123        let revealed_nodes: Vec<RevealedSparseNode> = hash_builder_proof_nodes
5124            .nodes_sorted()
5125            .into_iter()
5126            .map(|(path, node)| {
5127                let hash_mask = branch_node_hash_masks.get(&path).copied();
5128                let tree_mask = branch_node_tree_masks.get(&path).copied();
5129                RevealedSparseNode {
5130                    path,
5131                    node: TrieNode::decode(&mut &node[..]).unwrap(),
5132                    masks: TrieMasks { hash_mask, tree_mask },
5133                }
5134            })
5135            .collect();
5136        sparse.reveal_nodes(revealed_nodes).unwrap();
5137
5138        // Check that the branch node exists
5139        assert_eq!(
5140            sparse.upper_subtrie.nodes.get(&Nibbles::default()),
5141            Some(&SparseNode::new_branch(0b11.into()))
5142        );
5143
5144        // Remove the leaf for the first key
5145        sparse.remove_leaf(&key1(), &provider).unwrap();
5146
5147        // Check that the branch node was turned into an extension node
5148        assert_eq!(
5149            sparse.upper_subtrie.nodes.get(&Nibbles::default()),
5150            Some(&SparseNode::new_ext(Nibbles::from_nibbles_unchecked([0x01])))
5151        );
5152
5153        // Generate the proof for the third key and reveal it in the sparse trie
5154        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
5155            run_hash_builder(
5156                [(key1(), value()), (key2(), value()), (key3(), value())],
5157                NoopAccountTrieCursor::default(),
5158                Default::default(),
5159                [key2()],
5160            );
5161        let revealed_nodes: Vec<RevealedSparseNode> = hash_builder_proof_nodes
5162            .nodes_sorted()
5163            .into_iter()
5164            .map(|(path, node)| {
5165                let hash_mask = branch_node_hash_masks.get(&path).copied();
5166                let tree_mask = branch_node_tree_masks.get(&path).copied();
5167                RevealedSparseNode {
5168                    path,
5169                    node: TrieNode::decode(&mut &node[..]).unwrap(),
5170                    masks: TrieMasks { hash_mask, tree_mask },
5171                }
5172            })
5173            .collect();
5174        sparse.reveal_nodes(revealed_nodes).unwrap();
5175
5176        // Check that nothing changed in the extension node
5177        assert_eq!(
5178            sparse.upper_subtrie.nodes.get(&Nibbles::default()),
5179            Some(&SparseNode::new_ext(Nibbles::from_nibbles_unchecked([0x01])))
5180        );
5181    }
5182
5183    /// We have two leaves that share the same prefix: 0x0001 and 0x0002, and a leaf with a
5184    /// different prefix: 0x0100. Hash builder trie has only the first two leaves, and we have
5185    /// proofs for them.
5186    ///
5187    /// 1. Insert the leaf 0x0100 into the sparse trie, and check that the root extension node was
5188    ///    turned into a branch node.
5189    /// 2. Reveal the leaf 0x0001 in the sparse trie, and check that the root branch node wasn't
5190    ///    overwritten with the extension node from the proof.
5191    #[test]
5192    fn sparse_trie_reveal_node_3() {
5193        let key1 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x00, 0x01]));
5194        let key2 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x00, 0x02]));
5195        let key3 = || pad_nibbles_right(Nibbles::from_nibbles_unchecked([0x01, 0x00]));
5196        let value = || Account::default();
5197        let value_encoded = || {
5198            let mut account_rlp = Vec::new();
5199            value().into_trie_account(EMPTY_ROOT_HASH).encode(&mut account_rlp);
5200            account_rlp
5201        };
5202
5203        // Generate the proof for the root node and initialize the sparse trie with it
5204        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
5205            run_hash_builder(
5206                [(key1(), value()), (key2(), value())],
5207                NoopAccountTrieCursor::default(),
5208                Default::default(),
5209                [Nibbles::default()],
5210            );
5211
5212        let provider = DefaultTrieNodeProvider;
5213        let mut sparse = ParallelSparseTrie::from_root(
5214            TrieNode::decode(&mut &hash_builder_proof_nodes.nodes_sorted()[0].1[..]).unwrap(),
5215            TrieMasks {
5216                hash_mask: branch_node_hash_masks.get(&Nibbles::default()).copied(),
5217                tree_mask: branch_node_tree_masks.get(&Nibbles::default()).copied(),
5218            },
5219            false,
5220        )
5221        .unwrap();
5222
5223        // Check that the root extension node exists
5224        assert_matches!(
5225            sparse.upper_subtrie.nodes.get(&Nibbles::default()),
5226            Some(SparseNode::Extension { key, hash: None, store_in_db_trie: None }) if *key == Nibbles::from_nibbles([0x00])
5227        );
5228
5229        // Insert the leaf with a different prefix
5230        sparse.update_leaf(key3(), value_encoded(), &provider).unwrap();
5231
5232        // Check that the extension node was turned into a branch node
5233        assert_matches!(
5234            sparse.upper_subtrie.nodes.get(&Nibbles::default()),
5235            Some(SparseNode::Branch { state_mask, hash: None, store_in_db_trie: None }) if *state_mask == TrieMask::new(0b11)
5236        );
5237
5238        // Generate the proof for the first key and reveal it in the sparse trie
5239        let (_, _, hash_builder_proof_nodes, branch_node_hash_masks, branch_node_tree_masks) =
5240            run_hash_builder(
5241                [(key1(), value()), (key2(), value())],
5242                NoopAccountTrieCursor::default(),
5243                Default::default(),
5244                [key1()],
5245            );
5246        let revealed_nodes: Vec<RevealedSparseNode> = hash_builder_proof_nodes
5247            .nodes_sorted()
5248            .into_iter()
5249            .map(|(path, node)| {
5250                let hash_mask = branch_node_hash_masks.get(&path).copied();
5251                let tree_mask = branch_node_tree_masks.get(&path).copied();
5252                RevealedSparseNode {
5253                    path,
5254                    node: TrieNode::decode(&mut &node[..]).unwrap(),
5255                    masks: TrieMasks { hash_mask, tree_mask },
5256                }
5257            })
5258            .collect();
5259        sparse.reveal_nodes(revealed_nodes).unwrap();
5260
5261        // Check that the branch node wasn't overwritten by the extension node in the proof
5262        assert_matches!(
5263            sparse.upper_subtrie.nodes.get(&Nibbles::default()),
5264            Some(SparseNode::Branch { state_mask, hash: None, store_in_db_trie: None }) if *state_mask == TrieMask::new(0b11)
5265        );
5266    }
5267
5268    #[test]
5269    fn test_update_leaf_cross_level() {
5270        let ctx = ParallelSparseTrieTestContext;
5271        let mut trie =
5272            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5273
5274        // Test adding leaves that demonstrate the cross-level behavior
5275        // Based on the example: leaves 0x1234, 0x1245, 0x1334, 0x1345
5276        //
5277        // Final trie structure:
5278        // Upper trie:
5279        //   0x: Extension { key: 0x1 }
5280        //   └── 0x1: Branch { state_mask: 0x1100 }
5281        //       └── Subtrie (0x12): pointer to lower subtrie
5282        //       └── Subtrie (0x13): pointer to lower subtrie
5283        //
5284        // Lower subtrie (0x12):
5285        //   0x12: Branch { state_mask: 0x8 | 0x10 }
5286        //   ├── 0x123: Leaf { key: 0x4 }
5287        //   └── 0x124: Leaf { key: 0x5 }
5288        //
5289        // Lower subtrie (0x13):
5290        //   0x13: Branch { state_mask: 0x8 | 0x10 }
5291        //   ├── 0x133: Leaf { key: 0x4 }
5292        //   └── 0x134: Leaf { key: 0x5 }
5293
5294        // First add leaf 0x1345 - this should create a leaf in upper trie at 0x
5295        let (leaf1_path, value1) = ctx.create_test_leaf([0x1, 0x3, 0x4, 0x5], 1);
5296        trie.update_leaf(leaf1_path, value1.clone(), DefaultTrieNodeProvider).unwrap();
5297
5298        // Verify upper trie has a leaf at the root with key 1345
5299        ctx.assert_upper_subtrie(&trie)
5300            .has_leaf(&Nibbles::default(), &Nibbles::from_nibbles([0x1, 0x3, 0x4, 0x5]))
5301            .has_value(&leaf1_path, &value1);
5302
5303        // Add leaf 0x1234 - this should go first in the upper subtrie
5304        let (leaf2_path, value2) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4], 2);
5305        trie.update_leaf(leaf2_path, value2.clone(), DefaultTrieNodeProvider).unwrap();
5306
5307        // Upper trie should now have a branch at 0x1
5308        ctx.assert_upper_subtrie(&trie)
5309            .has_branch(&Nibbles::from_nibbles([0x1]), &[0x2, 0x3])
5310            .has_no_value(&leaf1_path)
5311            .has_no_value(&leaf2_path);
5312
5313        // Add leaf 0x1245 - this should cause a branch and create the 0x12 subtrie
5314        let (leaf3_path, value3) = ctx.create_test_leaf([0x1, 0x2, 0x4, 0x5], 3);
5315        trie.update_leaf(leaf3_path, value3.clone(), DefaultTrieNodeProvider).unwrap();
5316
5317        // Verify lower subtrie at 0x12 exists with correct structure
5318        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
5319            .has_branch(&Nibbles::from_nibbles([0x1, 0x2]), &[0x3, 0x4])
5320            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x3]), &Nibbles::from_nibbles([0x4]))
5321            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x4]), &Nibbles::from_nibbles([0x5]))
5322            .has_value(&leaf2_path, &value2)
5323            .has_value(&leaf3_path, &value3);
5324
5325        // Add leaf 0x1334 - this should create another lower subtrie
5326        let (leaf4_path, value4) = ctx.create_test_leaf([0x1, 0x3, 0x3, 0x4], 4);
5327        trie.update_leaf(leaf4_path, value4.clone(), DefaultTrieNodeProvider).unwrap();
5328
5329        // Verify lower subtrie at 0x13 exists with correct values
5330        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x3]))
5331            .has_value(&leaf1_path, &value1)
5332            .has_value(&leaf4_path, &value4);
5333
5334        // Verify the 0x12 subtrie still has its values
5335        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
5336            .has_value(&leaf2_path, &value2)
5337            .has_value(&leaf3_path, &value3);
5338
5339        // Upper trie has no values
5340        ctx.assert_upper_subtrie(&trie)
5341            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0x1]))
5342            .has_branch(&Nibbles::from_nibbles([0x1]), &[0x2, 0x3])
5343            .has_no_value(&leaf1_path)
5344            .has_no_value(&leaf2_path)
5345            .has_no_value(&leaf3_path)
5346            .has_no_value(&leaf4_path);
5347    }
5348
5349    #[test]
5350    fn test_update_leaf_split_at_level_boundary() {
5351        let ctx = ParallelSparseTrieTestContext;
5352        let mut trie =
5353            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5354
5355        // This test demonstrates what happens when we insert leaves that cause
5356        // splitting exactly at the upper/lower trie boundary (2 nibbles).
5357        //
5358        // Final trie structure:
5359        // Upper trie:
5360        //   0x: Extension { key: 0x12 }
5361        //       └── Subtrie (0x12): pointer to lower subtrie
5362        //
5363        // Lower subtrie (0x12):
5364        //   0x12: Branch { state_mask: 0x4 | 0x8 }
5365        //   ├── 0x122: Leaf { key: 0x4 }
5366        //   └── 0x123: Leaf { key: 0x4 }
5367
5368        // First insert a leaf that ends exactly at the boundary (2 nibbles)
5369        let (first_leaf_path, first_value) = ctx.create_test_leaf([0x1, 0x2, 0x2, 0x4], 1);
5370
5371        trie.update_leaf(first_leaf_path, first_value.clone(), DefaultTrieNodeProvider).unwrap();
5372
5373        // In an empty trie, the first leaf becomes the root, regardless of path length
5374        ctx.assert_upper_subtrie(&trie)
5375            .has_leaf(&Nibbles::default(), &Nibbles::from_nibbles([0x1, 0x2, 0x2, 0x4]))
5376            .has_value(&first_leaf_path, &first_value);
5377
5378        // Now insert another leaf that shares the same 2-nibble prefix
5379        let (second_leaf_path, second_value) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4], 2);
5380
5381        trie.update_leaf(second_leaf_path, second_value.clone(), DefaultTrieNodeProvider).unwrap();
5382
5383        // Now both leaves should be in a lower subtrie at index [0x1, 0x2]
5384        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
5385            .has_branch(&Nibbles::from_nibbles([0x1, 0x2]), &[0x2, 0x3])
5386            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x2]), &Nibbles::from_nibbles([0x4]))
5387            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x3]), &Nibbles::from_nibbles([0x4]))
5388            .has_value(&first_leaf_path, &first_value)
5389            .has_value(&second_leaf_path, &second_value);
5390
5391        // Upper subtrie should no longer have these values
5392        ctx.assert_upper_subtrie(&trie)
5393            .has_no_value(&first_leaf_path)
5394            .has_no_value(&second_leaf_path);
5395    }
5396
5397    #[test]
5398    fn test_update_subtrie_with_multiple_leaves() {
5399        let ctx = ParallelSparseTrieTestContext;
5400        let mut trie =
5401            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5402
5403        // First, add multiple leaves that will create a subtrie structure
5404        // All leaves share the prefix [0x1, 0x2] to ensure they create a subtrie
5405        //
5406        // This should result in a trie with the following structure:
5407        // 0x: Extension { key: 0x12 }
5408        //  └── Subtrie (0x12):
5409        //      0x12: Branch { state_mask: 0x3 | 0x4 }
5410        //      ├── 0x123: Branch { state_mask: 0x4 | 0x5 }
5411        //      │   ├── 0x1234: Leaf { key: 0x }
5412        //      │   └── 0x1235: Leaf { key: 0x }
5413        //      └── 0x124: Branch { state_mask: 0x6 | 0x7 }
5414        //          ├── 0x1246: Leaf { key: 0x }
5415        //          └── 0x1247: Leaf { key: 0x }
5416        let leaves = ctx.create_test_leaves(&[
5417            &[0x1, 0x2, 0x3, 0x4],
5418            &[0x1, 0x2, 0x3, 0x5],
5419            &[0x1, 0x2, 0x4, 0x6],
5420            &[0x1, 0x2, 0x4, 0x7],
5421        ]);
5422
5423        // Insert all leaves
5424        ctx.update_leaves(&mut trie, leaves.clone());
5425
5426        // Verify the upper subtrie has an extension node at the root with key 0x12
5427        ctx.assert_upper_subtrie(&trie)
5428            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0x1, 0x2]));
5429
5430        // Verify the subtrie structure using fluent assertions
5431        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
5432            .has_branch(&Nibbles::from_nibbles([0x1, 0x2]), &[0x3, 0x4])
5433            .has_branch(&Nibbles::from_nibbles([0x1, 0x2, 0x3]), &[0x4, 0x5])
5434            .has_branch(&Nibbles::from_nibbles([0x1, 0x2, 0x4]), &[0x6, 0x7])
5435            .has_value(&leaves[0].0, &leaves[0].1)
5436            .has_value(&leaves[1].0, &leaves[1].1)
5437            .has_value(&leaves[2].0, &leaves[2].1)
5438            .has_value(&leaves[3].0, &leaves[3].1);
5439
5440        // Now update one of the leaves with a new value
5441        let updated_path = Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4]);
5442        let (_, updated_value) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4], 100);
5443
5444        trie.update_leaf(updated_path, updated_value.clone(), DefaultTrieNodeProvider).unwrap();
5445
5446        // Verify the subtrie structure is maintained and value is updated
5447        // The branch structure should remain the same and all values should be present
5448        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
5449            .has_branch(&Nibbles::from_nibbles([0x1, 0x2]), &[0x3, 0x4])
5450            .has_value(&updated_path, &updated_value)
5451            .has_value(&leaves[1].0, &leaves[1].1)
5452            .has_value(&leaves[2].0, &leaves[2].1)
5453            .has_value(&leaves[3].0, &leaves[3].1);
5454
5455        // Add a new leaf that extends an existing branch
5456        let (new_leaf_path, new_leaf_value) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x6], 200);
5457
5458        trie.update_leaf(new_leaf_path, new_leaf_value.clone(), DefaultTrieNodeProvider).unwrap();
5459
5460        // Verify the branch at [0x1, 0x2, 0x3] now has an additional child
5461        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
5462            .has_branch(&Nibbles::from_nibbles([0x1, 0x2, 0x3]), &[0x4, 0x5, 0x6])
5463            .has_value(&new_leaf_path, &new_leaf_value);
5464    }
5465
5466    #[test]
5467    fn test_update_subtrie_extension_node_subtrie() {
5468        let ctx = ParallelSparseTrieTestContext;
5469        let mut trie =
5470            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5471
5472        // All leaves share the prefix [0x1, 0x2] to ensure they create a subtrie
5473        //
5474        // This should result in a trie with the following structure
5475        // 0x: Extension { key: 0x123 }
5476        //  └── Subtrie (0x12):
5477        //      0x123: Branch { state_mask: 0x3 | 0x4 }
5478        //      ├── 0x123: Leaf { key: 0x4 }
5479        //      └── 0x124: Leaf { key: 0x5 }
5480        let leaves = ctx.create_test_leaves(&[&[0x1, 0x2, 0x3, 0x4], &[0x1, 0x2, 0x3, 0x5]]);
5481
5482        // Insert all leaves
5483        ctx.update_leaves(&mut trie, leaves.clone());
5484
5485        // Verify the upper subtrie has an extension node at the root with key 0x123
5486        ctx.assert_upper_subtrie(&trie)
5487            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0x1, 0x2, 0x3]));
5488
5489        // Verify the lower subtrie structure
5490        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
5491            .has_branch(&Nibbles::from_nibbles([0x1, 0x2, 0x3]), &[0x4, 0x5])
5492            .has_value(&leaves[0].0, &leaves[0].1)
5493            .has_value(&leaves[1].0, &leaves[1].1);
5494    }
5495
5496    #[test]
5497    fn update_subtrie_extension_node_cross_level() {
5498        let ctx = ParallelSparseTrieTestContext;
5499        let mut trie =
5500            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5501
5502        // First, add multiple leaves that will create a subtrie structure
5503        // All leaves share the prefix [0x1, 0x2] to ensure they create a branch node and subtrie
5504        //
5505        // This should result in a trie with the following structure
5506        // 0x: Extension { key: 0x12 }
5507        //  └── Subtrie (0x12):
5508        //      0x12: Branch { state_mask: 0x3 | 0x4 }
5509        //      ├── 0x123: Leaf { key: 0x4 }
5510        //      └── 0x124: Leaf { key: 0x5 }
5511        let leaves = ctx.create_test_leaves(&[&[0x1, 0x2, 0x3, 0x4], &[0x1, 0x2, 0x4, 0x5]]);
5512
5513        // Insert all leaves
5514        ctx.update_leaves(&mut trie, leaves.clone());
5515
5516        // Verify the upper subtrie has an extension node at the root with key 0x12
5517        ctx.assert_upper_subtrie(&trie)
5518            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0x1, 0x2]));
5519
5520        // Verify the lower subtrie structure
5521        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
5522            .has_branch(&Nibbles::from_nibbles([0x1, 0x2]), &[0x3, 0x4])
5523            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x3]), &Nibbles::from_nibbles([0x4]))
5524            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x4]), &Nibbles::from_nibbles([0x5]))
5525            .has_value(&leaves[0].0, &leaves[0].1)
5526            .has_value(&leaves[1].0, &leaves[1].1);
5527    }
5528
5529    #[test]
5530    fn test_update_single_nibble_paths() {
5531        let ctx = ParallelSparseTrieTestContext;
5532        let mut trie =
5533            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5534
5535        // Test edge case: single nibble paths that create branches in upper trie
5536        //
5537        // Final trie structure:
5538        // Upper trie:
5539        //   0x: Branch { state_mask: 0x1 | 0x2 | 0x4 | 0x8 }
5540        //   ├── 0x0: Leaf { key: 0x }
5541        //   ├── 0x1: Leaf { key: 0x }
5542        //   ├── 0x2: Leaf { key: 0x }
5543        //   └── 0x3: Leaf { key: 0x }
5544
5545        // Insert leaves with single nibble paths
5546        let (leaf1_path, value1) = ctx.create_test_leaf([0x0], 1);
5547        let (leaf2_path, value2) = ctx.create_test_leaf([0x1], 2);
5548        let (leaf3_path, value3) = ctx.create_test_leaf([0x2], 3);
5549        let (leaf4_path, value4) = ctx.create_test_leaf([0x3], 4);
5550
5551        ctx.update_leaves(
5552            &mut trie,
5553            [
5554                (leaf1_path, value1.clone()),
5555                (leaf2_path, value2.clone()),
5556                (leaf3_path, value3.clone()),
5557                (leaf4_path, value4.clone()),
5558            ],
5559        );
5560
5561        // Verify upper trie has a branch at root with 4 children
5562        ctx.assert_upper_subtrie(&trie)
5563            .has_branch(&Nibbles::default(), &[0x0, 0x1, 0x2, 0x3])
5564            .has_leaf(&Nibbles::from_nibbles([0x0]), &Nibbles::default())
5565            .has_leaf(&Nibbles::from_nibbles([0x1]), &Nibbles::default())
5566            .has_leaf(&Nibbles::from_nibbles([0x2]), &Nibbles::default())
5567            .has_leaf(&Nibbles::from_nibbles([0x3]), &Nibbles::default())
5568            .has_value(&leaf1_path, &value1)
5569            .has_value(&leaf2_path, &value2)
5570            .has_value(&leaf3_path, &value3)
5571            .has_value(&leaf4_path, &value4);
5572    }
5573
5574    #[test]
5575    fn test_update_deep_extension_chain() {
5576        let ctx = ParallelSparseTrieTestContext;
5577        let mut trie =
5578            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5579
5580        // Test edge case: deep extension chains that span multiple levels
5581        //
5582        // Final trie structure:
5583        // Upper trie:
5584        //   0x: Extension { key: 0x111111 }
5585        //       └── Subtrie (0x11): pointer to lower subtrie
5586        //
5587        // Lower subtrie (0x11):
5588        //   0x111111: Branch { state_mask: 0x1 | 0x2 }
5589        //   ├── 0x1111110: Leaf { key: 0x }
5590        //   └── 0x1111111: Leaf { key: 0x }
5591
5592        // Create leaves with a long common prefix
5593        let (leaf1_path, value1) = ctx.create_test_leaf([0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x0], 1);
5594        let (leaf2_path, value2) = ctx.create_test_leaf([0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1], 2);
5595
5596        ctx.update_leaves(&mut trie, [(leaf1_path, value1.clone()), (leaf2_path, value2.clone())]);
5597
5598        // Verify upper trie has extension with the full common prefix
5599        ctx.assert_upper_subtrie(&trie).has_extension(
5600            &Nibbles::default(),
5601            &Nibbles::from_nibbles([0x1, 0x1, 0x1, 0x1, 0x1, 0x1]),
5602        );
5603
5604        // Verify lower subtrie has branch structure
5605        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x1]))
5606            .has_branch(&Nibbles::from_nibbles([0x1, 0x1, 0x1, 0x1, 0x1, 0x1]), &[0x0, 0x1])
5607            .has_leaf(
5608                &Nibbles::from_nibbles([0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x0]),
5609                &Nibbles::default(),
5610            )
5611            .has_leaf(
5612                &Nibbles::from_nibbles([0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1]),
5613                &Nibbles::default(),
5614            )
5615            .has_value(&leaf1_path, &value1)
5616            .has_value(&leaf2_path, &value2);
5617    }
5618
5619    #[test]
5620    fn test_update_branch_with_all_nibbles() {
5621        let ctx = ParallelSparseTrieTestContext;
5622        let mut trie =
5623            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5624
5625        // Test edge case: branch node with all 16 possible nibble children
5626        //
5627        // Final trie structure:
5628        // Upper trie:
5629        //   0x: Extension { key: 0xA }
5630        //       └── Subtrie (0xA0): pointer to lower subtrie
5631        //
5632        // Lower subtrie (0xA0):
5633        //   0xA0: Branch { state_mask: 0xFFFF } (all 16 children)
5634        //   ├── 0xA00: Leaf { key: 0x }
5635        //   ├── 0xA01: Leaf { key: 0x }
5636        //   ├── 0xA02: Leaf { key: 0x }
5637        //   ... (all nibbles 0x0 through 0xF)
5638        //   └── 0xA0F: Leaf { key: 0x }
5639
5640        // Create leaves for all 16 possible nibbles
5641        let mut leaves = Vec::new();
5642        for nibble in 0x0..=0xF {
5643            let (path, value) = ctx.create_test_leaf([0xA, 0x0, nibble], nibble as u64 + 1);
5644            leaves.push((path, value));
5645        }
5646
5647        // Insert all leaves
5648        ctx.update_leaves(&mut trie, leaves.iter().cloned());
5649
5650        // Verify upper trie structure
5651        ctx.assert_upper_subtrie(&trie)
5652            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0xA, 0x0]));
5653
5654        // Verify lower subtrie has branch with all 16 children
5655        let mut subtrie_assert =
5656            ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0xA, 0x0])).has_branch(
5657                &Nibbles::from_nibbles([0xA, 0x0]),
5658                &[0x0, 0x1, 0x2, 0x3, 0x4, 0x5, 0x6, 0x7, 0x8, 0x9, 0xA, 0xB, 0xC, 0xD, 0xE, 0xF],
5659            );
5660
5661        // Verify all leaves exist
5662        for (i, (path, value)) in leaves.iter().enumerate() {
5663            subtrie_assert = subtrie_assert
5664                .has_leaf(&Nibbles::from_nibbles([0xA, 0x0, i as u8]), &Nibbles::default())
5665                .has_value(path, value);
5666        }
5667    }
5668
5669    #[test]
5670    fn test_update_creates_multiple_subtries() {
5671        let ctx = ParallelSparseTrieTestContext;
5672        let mut trie =
5673            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5674
5675        // Test edge case: updates that create multiple subtries at once
5676        //
5677        // Final trie structure:
5678        // Upper trie:
5679        //   0x: Extension { key: 0x0 }
5680        //       └── 0x0: Branch { state_mask: 0xF }
5681        //           ├── Subtrie (0x00): pointer
5682        //           ├── Subtrie (0x01): pointer
5683        //           ├── Subtrie (0x02): pointer
5684        //           └── Subtrie (0x03): pointer
5685        //
5686        // Each lower subtrie has leaves:
5687        //   0xXY: Leaf { key: 0xZ... }
5688
5689        // Create leaves that will force multiple subtries
5690        let leaves = vec![
5691            ctx.create_test_leaf([0x0, 0x0, 0x1, 0x2], 1),
5692            ctx.create_test_leaf([0x0, 0x1, 0x3, 0x4], 2),
5693            ctx.create_test_leaf([0x0, 0x2, 0x5, 0x6], 3),
5694            ctx.create_test_leaf([0x0, 0x3, 0x7, 0x8], 4),
5695        ];
5696
5697        // Insert all leaves
5698        ctx.update_leaves(&mut trie, leaves.iter().cloned());
5699
5700        // Verify upper trie has extension then branch
5701        ctx.assert_upper_subtrie(&trie)
5702            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0x0]))
5703            .has_branch(&Nibbles::from_nibbles([0x0]), &[0x0, 0x1, 0x2, 0x3]);
5704
5705        // Verify each subtrie exists and contains its leaf
5706        for (i, (leaf_path, leaf_value)) in leaves.iter().enumerate() {
5707            let subtrie_path = Nibbles::from_nibbles([0x0, i as u8]);
5708            ctx.assert_subtrie(&trie, subtrie_path)
5709                .has_leaf(
5710                    &subtrie_path,
5711                    &Nibbles::from_nibbles(match i {
5712                        0 => vec![0x1, 0x2],
5713                        1 => vec![0x3, 0x4],
5714                        2 => vec![0x5, 0x6],
5715                        3 => vec![0x7, 0x8],
5716                        _ => unreachable!(),
5717                    }),
5718                )
5719                .has_value(leaf_path, leaf_value);
5720        }
5721    }
5722
5723    #[test]
5724    fn test_update_extension_to_branch_transformation() {
5725        let ctx = ParallelSparseTrieTestContext;
5726        let mut trie =
5727            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5728
5729        // Test edge case: extension node transforms to branch when split
5730        //
5731        // Initial state after first two leaves:
5732        // Upper trie:
5733        //   0x: Extension { key: 0xFF0 }
5734        //       └── Subtrie (0xFF): pointer
5735        //
5736        // After third leaf (0xF0...):
5737        // Upper trie:
5738        //   0x: Extension { key: 0xF }
5739        //       └── 0xF: Branch { state_mask: 0x10 | 0x8000 }
5740        //           ├── Subtrie (0xF0): pointer
5741        //           └── Subtrie (0xFF): pointer
5742
5743        // First two leaves share prefix 0xFF0
5744        let (leaf1_path, value1) = ctx.create_test_leaf([0xF, 0xF, 0x0, 0x1], 1);
5745        let (leaf2_path, value2) = ctx.create_test_leaf([0xF, 0xF, 0x0, 0x2], 2);
5746        let (leaf3_path, value3) = ctx.create_test_leaf([0xF, 0x0, 0x0, 0x3], 3);
5747
5748        ctx.update_leaves(&mut trie, [(leaf1_path, value1.clone()), (leaf2_path, value2.clone())]);
5749
5750        // Verify initial extension structure
5751        ctx.assert_upper_subtrie(&trie)
5752            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0xF, 0xF, 0x0]));
5753
5754        // Add leaf that splits the extension
5755        ctx.update_leaves(&mut trie, [(leaf3_path, value3.clone())]);
5756
5757        // Verify transformed structure
5758        ctx.assert_upper_subtrie(&trie)
5759            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0xF]))
5760            .has_branch(&Nibbles::from_nibbles([0xF]), &[0x0, 0xF]);
5761
5762        // Verify subtries
5763        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0xF, 0xF]))
5764            .has_branch(&Nibbles::from_nibbles([0xF, 0xF, 0x0]), &[0x1, 0x2])
5765            .has_leaf(&Nibbles::from_nibbles([0xF, 0xF, 0x0, 0x1]), &Nibbles::default())
5766            .has_leaf(&Nibbles::from_nibbles([0xF, 0xF, 0x0, 0x2]), &Nibbles::default())
5767            .has_value(&leaf1_path, &value1)
5768            .has_value(&leaf2_path, &value2);
5769
5770        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0xF, 0x0]))
5771            .has_leaf(&Nibbles::from_nibbles([0xF, 0x0]), &Nibbles::from_nibbles([0x0, 0x3]))
5772            .has_value(&leaf3_path, &value3);
5773    }
5774
5775    #[test]
5776    fn test_update_upper_extension_reveal_lower_hash_node() {
5777        let ctx = ParallelSparseTrieTestContext;
5778
5779        // Test edge case: extension pointing to hash node that gets updated to branch
5780        // and reveals the hash node from lower trie
5781        //
5782        // Setup:
5783        // Upper trie:
5784        //   0x: Extension { key: 0xAB }
5785        //       └── Subtrie (0xAB): pointer
5786        // Lower trie (0xAB):
5787        //   0xAB: Hash
5788        //
5789        // After update:
5790        // Upper trie:
5791        //   0x: Extension { key: 0xA }
5792        //       └── 0xA: Branch { state_mask: 0b100000000001 }
5793        //                ├── 0xA0: Leaf { value: ... }
5794        //                └── 0xAB: pointer
5795        // Lower trie (0xAB):
5796        //   0xAB: Branch { state_mask: 0b11 }
5797        //         ├── 0xAB1: Hash
5798        //         └── 0xAB2: Hash
5799
5800        // Create a mock provider that will provide the hash node
5801        let mut provider = MockTrieNodeProvider::new();
5802
5803        // Create revealed branch which will get revealed and add it to the mock provider
5804        let child_hashes = [
5805            RlpNode::word_rlp(&B256::repeat_byte(0x11)),
5806            RlpNode::word_rlp(&B256::repeat_byte(0x22)),
5807        ];
5808        let revealed_branch = create_branch_node_with_children(&[0x1, 0x2], child_hashes);
5809        let mut encoded = Vec::new();
5810        revealed_branch.encode(&mut encoded);
5811        provider.add_revealed_node(
5812            Nibbles::from_nibbles([0xA, 0xB]),
5813            RevealedNode { node: encoded.into(), tree_mask: None, hash_mask: None },
5814        );
5815
5816        let mut trie = new_test_trie(
5817            [
5818                (Nibbles::default(), SparseNode::new_ext(Nibbles::from_nibbles([0xA, 0xB]))),
5819                (Nibbles::from_nibbles([0xA, 0xB]), SparseNode::Hash(B256::repeat_byte(0x42))),
5820            ]
5821            .into_iter(),
5822        );
5823
5824        // Now add a leaf that will force the hash node to become a branch
5825        let (leaf_path, value) = ctx.create_test_leaf([0xA, 0x0], 1);
5826        trie.update_leaf(leaf_path, value, provider).unwrap();
5827
5828        // Verify the structure: extension should now terminate in a branch on the upper trie
5829        ctx.assert_upper_subtrie(&trie)
5830            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0xA]))
5831            .has_branch(&Nibbles::from_nibbles([0xA]), &[0x0, 0xB]);
5832
5833        // Verify the lower trie now has a branch structure
5834        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0xA, 0xB]))
5835            .has_branch(&Nibbles::from_nibbles([0xA, 0xB]), &[0x1, 0x2])
5836            .has_hash(&Nibbles::from_nibbles([0xA, 0xB, 0x1]), &B256::repeat_byte(0x11))
5837            .has_hash(&Nibbles::from_nibbles([0xA, 0xB, 0x2]), &B256::repeat_byte(0x22));
5838    }
5839
5840    #[test]
5841    fn test_update_long_shared_prefix_at_boundary() {
5842        let ctx = ParallelSparseTrieTestContext;
5843        let mut trie =
5844            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5845
5846        // Test edge case: leaves with long shared prefix that ends exactly at 2-nibble boundary
5847        //
5848        // Final trie structure:
5849        // Upper trie:
5850        //   0x: Extension { key: 0xAB }
5851        //       └── Subtrie (0xAB): pointer to lower subtrie
5852        //
5853        // Lower subtrie (0xAB):
5854        //   0xAB: Branch { state_mask: 0x1000 | 0x2000 }
5855        //   ├── 0xABC: Leaf { key: 0xDEF }
5856        //   └── 0xABD: Leaf { key: 0xEF0 }
5857
5858        // Create leaves that share exactly 2 nibbles
5859        let (leaf1_path, value1) = ctx.create_test_leaf([0xA, 0xB, 0xC, 0xD, 0xE, 0xF], 1);
5860        let (leaf2_path, value2) = ctx.create_test_leaf([0xA, 0xB, 0xD, 0xE, 0xF, 0x0], 2);
5861
5862        trie.update_leaf(leaf1_path, value1.clone(), DefaultTrieNodeProvider).unwrap();
5863        trie.update_leaf(leaf2_path, value2.clone(), DefaultTrieNodeProvider).unwrap();
5864
5865        // Verify upper trie structure
5866        ctx.assert_upper_subtrie(&trie)
5867            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0xA, 0xB]));
5868
5869        // Verify lower subtrie structure
5870        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0xA, 0xB]))
5871            .has_branch(&Nibbles::from_nibbles([0xA, 0xB]), &[0xC, 0xD])
5872            .has_leaf(
5873                &Nibbles::from_nibbles([0xA, 0xB, 0xC]),
5874                &Nibbles::from_nibbles([0xD, 0xE, 0xF]),
5875            )
5876            .has_leaf(
5877                &Nibbles::from_nibbles([0xA, 0xB, 0xD]),
5878                &Nibbles::from_nibbles([0xE, 0xF, 0x0]),
5879            )
5880            .has_value(&leaf1_path, &value1)
5881            .has_value(&leaf2_path, &value2);
5882    }
5883
5884    #[test]
5885    fn test_update_branch_to_extension_collapse() {
5886        let ctx = ParallelSparseTrieTestContext;
5887        let mut trie =
5888            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5889
5890        // Test creating a trie with leaves that share a long common prefix
5891        //
5892        // Initial state with 3 leaves (0x1234, 0x2345, 0x2356):
5893        // Upper trie:
5894        //   0x: Branch { state_mask: 0x6 }
5895        //       ├── 0x1: Leaf { key: 0x234 }
5896        //       └── 0x2: Extension { key: 0x3 }
5897        //           └── Subtrie (0x23): pointer
5898        // Lower subtrie (0x23):
5899        //   0x23: Branch { state_mask: 0x30 }
5900        //       ├── 0x234: Leaf { key: 0x5 }
5901        //       └── 0x235: Leaf { key: 0x6 }
5902        //
5903        // Then we create a new trie with leaves (0x1234, 0x1235, 0x1236):
5904        // Expected structure:
5905        // Upper trie:
5906        //   0x: Extension { key: 0x123 }
5907        //       └── Subtrie (0x12): pointer
5908        // Lower subtrie (0x12):
5909        //   0x123: Branch { state_mask: 0x70 } // bits 4, 5, 6 set
5910        //       ├── 0x1234: Leaf { key: 0x }
5911        //       ├── 0x1235: Leaf { key: 0x }
5912        //       └── 0x1236: Leaf { key: 0x }
5913
5914        // Create initial leaves
5915        let (leaf1_path, value1) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4], 1);
5916        let (leaf2_path, value2) = ctx.create_test_leaf([0x2, 0x3, 0x4, 0x5], 2);
5917        let (leaf3_path, value3) = ctx.create_test_leaf([0x2, 0x3, 0x5, 0x6], 3);
5918
5919        trie.update_leaf(leaf1_path, value1, DefaultTrieNodeProvider).unwrap();
5920        trie.update_leaf(leaf2_path, value2, DefaultTrieNodeProvider).unwrap();
5921        trie.update_leaf(leaf3_path, value3, DefaultTrieNodeProvider).unwrap();
5922
5923        // Verify initial structure has branch at root
5924        ctx.assert_upper_subtrie(&trie).has_branch(&Nibbles::default(), &[0x1, 0x2]);
5925
5926        // Now update to create a pattern where extension is more efficient
5927        // Replace leaves to all share prefix 0x123
5928        let (new_leaf1_path, new_value1) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4], 10);
5929        let (new_leaf2_path, new_value2) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x5], 11);
5930        let (new_leaf3_path, new_value3) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x6], 12);
5931
5932        // Clear and add new leaves
5933        let mut trie =
5934            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5935        trie.update_leaf(new_leaf1_path, new_value1.clone(), DefaultTrieNodeProvider).unwrap();
5936        trie.update_leaf(new_leaf2_path, new_value2.clone(), DefaultTrieNodeProvider).unwrap();
5937        trie.update_leaf(new_leaf3_path, new_value3.clone(), DefaultTrieNodeProvider).unwrap();
5938
5939        // Verify new structure has extension
5940        ctx.assert_upper_subtrie(&trie)
5941            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0x1, 0x2, 0x3]));
5942
5943        // Verify lower subtrie path was correctly updated to 0x123
5944        ctx.assert_subtrie_path(&trie, [0x1, 0x2], [0x1, 0x2, 0x3]);
5945
5946        // Verify lower subtrie - all three leaves should be properly inserted
5947        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
5948            .has_branch(&Nibbles::from_nibbles([0x1, 0x2, 0x3]), &[0x4, 0x5, 0x6]) // All three children
5949            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4]), &Nibbles::default())
5950            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x5]), &Nibbles::default())
5951            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x6]), &Nibbles::default())
5952            .has_value(&new_leaf1_path, &new_value1)
5953            .has_value(&new_leaf2_path, &new_value2)
5954            .has_value(&new_leaf3_path, &new_value3);
5955    }
5956
5957    #[test]
5958    fn test_update_shared_prefix_patterns() {
5959        let ctx = ParallelSparseTrieTestContext;
5960        let mut trie =
5961            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
5962
5963        // Test edge case: different patterns of shared prefixes
5964        //
5965        // Final trie structure:
5966        // Upper trie:
5967        //   0x: Branch { state_mask: 0x6 }
5968        //       ├── 0x1: Leaf { key: 0x234 }
5969        //       └── 0x2: Extension { key: 0x3 }
5970        //           └── Subtrie (0x23): pointer
5971        //
5972        // Lower subtrie (0x23):
5973        //   0x23: Branch { state_mask: 0x10 | 0x20 }
5974        //   ├── 0x234: Leaf { key: 0x5 }
5975        //   └── 0x235: Leaf { key: 0x6 }
5976
5977        // Create leaves with different shared prefix patterns
5978        let (leaf1_path, value1) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4], 1);
5979        let (leaf2_path, value2) = ctx.create_test_leaf([0x2, 0x3, 0x4, 0x5], 2);
5980        let (leaf3_path, value3) = ctx.create_test_leaf([0x2, 0x3, 0x5, 0x6], 3);
5981
5982        trie.update_leaf(leaf1_path, value1, DefaultTrieNodeProvider).unwrap();
5983        trie.update_leaf(leaf2_path, value2.clone(), DefaultTrieNodeProvider).unwrap();
5984        trie.update_leaf(leaf3_path, value3.clone(), DefaultTrieNodeProvider).unwrap();
5985
5986        // Verify upper trie structure
5987        ctx.assert_upper_subtrie(&trie)
5988            .has_branch(&Nibbles::default(), &[0x1, 0x2])
5989            .has_leaf(&Nibbles::from_nibbles([0x1]), &Nibbles::from_nibbles([0x2, 0x3, 0x4]))
5990            .has_extension(&Nibbles::from_nibbles([0x2]), &Nibbles::from_nibbles([0x3]));
5991
5992        // Verify lower subtrie structure
5993        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x2, 0x3]))
5994            .has_branch(&Nibbles::from_nibbles([0x2, 0x3]), &[0x4, 0x5])
5995            .has_leaf(&Nibbles::from_nibbles([0x2, 0x3, 0x4]), &Nibbles::from_nibbles([0x5]))
5996            .has_leaf(&Nibbles::from_nibbles([0x2, 0x3, 0x5]), &Nibbles::from_nibbles([0x6]))
5997            .has_value(&leaf2_path, &value2)
5998            .has_value(&leaf3_path, &value3);
5999    }
6000
6001    #[test]
6002    fn test_progressive_branch_creation() {
6003        let ctx = ParallelSparseTrieTestContext;
6004        let mut trie =
6005            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
6006
6007        // Test starting with a single leaf and progressively adding leaves
6008        // that create branch nodes at shorter and shorter paths
6009        //
6010        // Step 1: Add leaf at 0x12345
6011        // Upper trie:
6012        //   0x: Leaf { key: 0x12345 }
6013        //
6014        // Step 2: Add leaf at 0x12346
6015        // Upper trie:
6016        //   0x: Extension { key: 0x1234 }
6017        //       └── Subtrie (0x12): pointer
6018        // Lower subtrie (0x12):
6019        //   0x1234: Branch { state_mask: 0x60 }  // bits 5 and 6 set
6020        //       ├── 0x12345: Leaf { key: 0x }
6021        //       └── 0x12346: Leaf { key: 0x }
6022        //
6023        // Step 3: Add leaf at 0x1235
6024        // Lower subtrie (0x12) updates to:
6025        //   0x123: Branch { state_mask: 0x30 }  // bits 4 and 5 set
6026        //       ├── 0x1234: Branch { state_mask: 0x60 }
6027        //       │   ├── 0x12345: Leaf { key: 0x }
6028        //       │   └── 0x12346: Leaf { key: 0x }
6029        //       └── 0x1235: Leaf { key: 0x }
6030        //
6031        // Step 4: Add leaf at 0x124
6032        // Lower subtrie (0x12) updates to:
6033        //   0x12: Branch { state_mask: 0x18 }  // bits 3 and 4 set
6034        //       ├── 0x123: Branch { state_mask: 0x30 }
6035        //       │   ├── 0x1234: Branch { state_mask: 0x60 }
6036        //       │   │   ├── 0x12345: Leaf { key: 0x }
6037        //       │   │   └── 0x12346: Leaf { key: 0x }
6038        //       │   └── 0x1235: Leaf { key: 0x }
6039        //       └── 0x124: Leaf { key: 0x }
6040
6041        // Step 1: Add first leaf - initially stored as leaf in upper trie
6042        let (leaf1_path, value1) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4, 0x5], 1);
6043        trie.update_leaf(leaf1_path, value1.clone(), DefaultTrieNodeProvider).unwrap();
6044
6045        // Verify leaf node in upper trie (optimized single-leaf case)
6046        ctx.assert_upper_subtrie(&trie)
6047            .has_leaf(&Nibbles::default(), &Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4, 0x5]))
6048            .has_value(&leaf1_path, &value1);
6049
6050        // Step 2: Add leaf at 0x12346 - creates branch at 0x1234
6051        let (leaf2_path, value2) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x4, 0x6], 2);
6052        trie.update_leaf(leaf2_path, value2.clone(), DefaultTrieNodeProvider).unwrap();
6053
6054        // Verify extension now goes to 0x1234
6055        ctx.assert_upper_subtrie(&trie)
6056            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4]));
6057
6058        // Verify subtrie path updated to 0x1234
6059        ctx.assert_subtrie_path(&trie, [0x1, 0x2], [0x1, 0x2, 0x3, 0x4]);
6060
6061        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
6062            .has_branch(&Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4]), &[0x5, 0x6])
6063            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4, 0x5]), &Nibbles::default())
6064            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4, 0x6]), &Nibbles::default())
6065            .has_value(&leaf1_path, &value1)
6066            .has_value(&leaf2_path, &value2);
6067
6068        // Step 3: Add leaf at 0x1235 - creates branch at 0x123
6069        let (leaf3_path, value3) = ctx.create_test_leaf([0x1, 0x2, 0x3, 0x5], 3);
6070        trie.update_leaf(leaf3_path, value3.clone(), DefaultTrieNodeProvider).unwrap();
6071
6072        // Verify extension now goes to 0x123
6073        ctx.assert_upper_subtrie(&trie)
6074            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0x1, 0x2, 0x3]));
6075
6076        // Verify subtrie path updated to 0x123
6077        ctx.assert_subtrie_path(&trie, [0x1, 0x2], [0x1, 0x2, 0x3]);
6078
6079        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
6080            .has_branch(&Nibbles::from_nibbles([0x1, 0x2, 0x3]), &[0x4, 0x5])
6081            .has_branch(&Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4]), &[0x5, 0x6])
6082            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x5]), &Nibbles::default())
6083            .has_value(&leaf1_path, &value1)
6084            .has_value(&leaf2_path, &value2)
6085            .has_value(&leaf3_path, &value3);
6086
6087        // Step 4: Add leaf at 0x124 - creates branch at 0x12 (subtrie root)
6088        let (leaf4_path, value4) = ctx.create_test_leaf([0x1, 0x2, 0x4], 4);
6089        trie.update_leaf(leaf4_path, value4.clone(), DefaultTrieNodeProvider).unwrap();
6090
6091        // Verify extension now goes to 0x12
6092        ctx.assert_upper_subtrie(&trie)
6093            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles([0x1, 0x2]));
6094
6095        // Verify subtrie path updated to 0x12
6096        ctx.assert_subtrie_path(&trie, [0x1, 0x2], [0x1, 0x2]);
6097
6098        // Verify final structure
6099        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0x1, 0x2]))
6100            .has_branch(&Nibbles::from_nibbles([0x1, 0x2]), &[0x3, 0x4])
6101            .has_branch(&Nibbles::from_nibbles([0x1, 0x2, 0x3]), &[0x4, 0x5])
6102            .has_branch(&Nibbles::from_nibbles([0x1, 0x2, 0x3, 0x4]), &[0x5, 0x6])
6103            .has_leaf(&Nibbles::from_nibbles([0x1, 0x2, 0x4]), &Nibbles::default())
6104            .has_value(&leaf1_path, &value1)
6105            .has_value(&leaf2_path, &value2)
6106            .has_value(&leaf3_path, &value3)
6107            .has_value(&leaf4_path, &value4);
6108    }
6109
6110    #[test]
6111    fn test_update_max_depth_paths() {
6112        let ctx = ParallelSparseTrieTestContext;
6113        let mut trie =
6114            ParallelSparseTrie::from_root(TrieNode::EmptyRoot, TrieMasks::none(), true).unwrap();
6115
6116        // Test edge case: very long paths (64 nibbles - max for addresses/storage)
6117        //
6118        // Final trie structure:
6119        // Upper trie:
6120        //   0x: Extension { key: 0xFF }
6121        //       └── Subtrie (0xFF): pointer
6122        //
6123        // Lower subtrie (0xFF):
6124        //   Has very long paths with slight differences at the end
6125
6126        // Create two 64-nibble paths that differ only in the last nibble
6127        let mut path1_nibbles = vec![0xF; 63];
6128        path1_nibbles.push(0x0);
6129        let mut path2_nibbles = vec![0xF; 63];
6130        path2_nibbles.push(0x1);
6131
6132        let (leaf1_path, value1) = ctx.create_test_leaf(&path1_nibbles, 1);
6133        let (leaf2_path, value2) = ctx.create_test_leaf(&path2_nibbles, 2);
6134
6135        trie.update_leaf(leaf1_path, value1.clone(), DefaultTrieNodeProvider).unwrap();
6136        trie.update_leaf(leaf2_path, value2.clone(), DefaultTrieNodeProvider).unwrap();
6137
6138        // The common prefix of 63 F's will create a very long extension
6139        let extension_key = vec![0xF; 63];
6140        ctx.assert_upper_subtrie(&trie)
6141            .has_extension(&Nibbles::default(), &Nibbles::from_nibbles(&extension_key));
6142
6143        // Verify the subtrie has the branch at the end
6144        ctx.assert_subtrie(&trie, Nibbles::from_nibbles([0xF, 0xF]))
6145            .has_branch(&Nibbles::from_nibbles(&path1_nibbles[..63]), &[0x0, 0x1])
6146            .has_value(&leaf1_path, &value1)
6147            .has_value(&leaf2_path, &value2);
6148    }
6149
6150    #[test]
6151    fn test_hoodie_block_1_data() {
6152        // Reveal node at path Nibbles(0x) - root branch node
6153        let root_branch_stack = vec![
6154            hex!("a0550b6aba4dd4582a2434d2cbdad8d3007d09f622d7a6e6eaa7a49385823c2fa2"),
6155            hex!("a04788a4975a9e1efd29b834fd80fdfe8a57cc1b1c5ace6d30ce5a36a15e0092b3"),
6156            hex!("a093aeccf87da304e6f7d09edc5d7bd3a552808866d2149dd0940507a8f9bfa910"),
6157            hex!("a08b5b423ba68d0dec2eca1f408076f9170678505eb4a5db2abbbd83bb37666949"),
6158            hex!("a08592f62216af4218098a78acad7cf472a727fb55e6c27d3cfdf2774d4518eb83"),
6159            hex!("a0ef02aeee845cb64c11f85edc1a3094227c26445952554b8a9248915d80c746c3"),
6160            hex!("a0df2529ee3a1ce4df5a758cf17e6a86d0fb5ea22ab7071cf60af6412e9b0a428a"),
6161            hex!("a0acaa1092db69cd5a63676685827b3484c4b80dc1d3361f6073bbb9240101e144"),
6162            hex!("a09c3f2bb2a729d71f246a833353ade65667716bb330e0127a3299a42d11200f93"),
6163            hex!("a0ce978470f4c0b1f8069570563a14d2b79d709add2db4bf22dd9b6aed3271c566"),
6164            hex!("a095f783cd1d464a60e3c8adcadc28c6eb9fec7306664df39553be41dccc909606"),
6165            hex!("a0a9083f5fb914b255e1feb5d951a4dfddacf3c8003ef1d1ec6a13bb6ba5b2ac62"),
6166            hex!("a0fec113d537d8577cd361e0cabf5e95ef58f1cc34318292fdecce9fae57c3e094"),
6167            hex!("a08b7465f5fe8b3e3c0d087cb7521310d4065ef2a0ee43bf73f68dee8a5742b3dd"),
6168            hex!("a0c589aa1ae3d5fd87d8640957f7d5184a4ac06f393b453a8e8ed7e8fba0d385c8"),
6169            hex!("a0b516d6f3352f87beab4ed6e7322f191fc7a147686500ef4de7dd290ad784ef51"),
6170        ];
6171
6172        let root_branch_rlp_stack: Vec<RlpNode> = root_branch_stack
6173            .iter()
6174            .map(|hex_str| RlpNode::from_raw_rlp(&hex_str[..]).unwrap())
6175            .collect();
6176
6177        let root_branch_node = BranchNode::new(
6178            root_branch_rlp_stack,
6179            TrieMask::new(0b1111111111111111), // state_mask: all 16 children present
6180        );
6181
6182        let root_branch_masks = TrieMasks {
6183            hash_mask: Some(TrieMask::new(0b1111111111111111)),
6184            tree_mask: Some(TrieMask::new(0b1111111111111111)),
6185        };
6186
6187        let mut trie = ParallelSparseTrie::from_root(
6188            TrieNode::Branch(root_branch_node),
6189            root_branch_masks,
6190            true,
6191        )
6192        .unwrap();
6193
6194        // Reveal node at path Nibbles(0x3) - branch node
6195        let branch_0x3_stack = vec![
6196            hex!("a09da7d9755fe0c558b3c3de9fdcdf9f28ae641f38c9787b05b73ab22ae53af3e2"),
6197            hex!("a0d9990bf0b810d1145ecb2b011fd68c63cc85564e6724166fd4a9520180706e5f"),
6198            hex!("a0f60eb4b12132a40df05d9bbdb88bbde0185a3f097f3c76bf4200c23eda26cf86"),
6199            hex!("a0ca976997ddaf06f18992f6207e4f6a05979d07acead96568058789017cc6d06b"),
6200            hex!("a04d78166b48044fdc28ed22d2fd39c8df6f8aaa04cb71d3a17286856f6893ff83"),
6201            hex!("a021d4f90c34d3f1706e78463b6482bca77a3aa1cd059a3f326c42a1cfd30b9b60"),
6202            hex!("a0fc3b71c33e2e6b77c5e494c1db7fdbb447473f003daf378c7a63ba9bf3f0049d"),
6203            hex!("a0e33ed2be194a3d93d343e85642447c93a9d0cfc47a016c2c23d14c083be32a7c"),
6204            hex!("a07b8e7a21c1178d28074f157b50fca85ee25c12568ff8e9706dcbcdacb77bf854"),
6205            hex!("a0973274526811393ea0bf4811ca9077531db00d06b86237a2ecd683f55ba4bcb0"),
6206            hex!("a03a93d726d7487874e51b52d8d534c63aa2a689df18e3b307c0d6cb0a388b00f3"),
6207            hex!("a06aa67101d011d1c22fe739ef83b04b5214a3e2f8e1a2625d8bfdb116b447e86f"),
6208            hex!("a02dd545b33c62d33a183e127a08a4767fba891d9f3b94fc20a2ca02600d6d1fff"),
6209            hex!("a0fe6db87d00f06d53bff8169fa497571ff5af1addfb715b649b4d79dd3e394b04"),
6210            hex!("a0d9240a9d2d5851d05a97ff3305334dfdb0101e1e321fc279d2bb3cad6afa8fc8"),
6211            hex!("a01b69c6ab5173de8a8ec53a6ebba965713a4cc7feb86cb3e230def37c230ca2b2"),
6212        ];
6213
6214        let branch_0x3_rlp_stack: Vec<RlpNode> = branch_0x3_stack
6215            .iter()
6216            .map(|hex_str| RlpNode::from_raw_rlp(&hex_str[..]).unwrap())
6217            .collect();
6218
6219        let branch_0x3_node = BranchNode::new(
6220            branch_0x3_rlp_stack,
6221            TrieMask::new(0b1111111111111111), // state_mask: all 16 children present
6222        );
6223
6224        let branch_0x3_masks = TrieMasks {
6225            hash_mask: Some(TrieMask::new(0b0100010000010101)),
6226            tree_mask: Some(TrieMask::new(0b0100000000000000)),
6227        };
6228
6229        // Reveal node at path Nibbles(0x37) - leaf node
6230        let leaf_path = Nibbles::from_nibbles([0x3, 0x7]);
6231        let leaf_key = Nibbles::unpack(
6232            &hex!("d65eaa92c6bc4c13a5ec45527f0c18ea8932588728769ec7aecfe6d9f32e42")[..],
6233        );
6234        let leaf_value = hex!("f8440180a056e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421a0f57acd40259872606d76197ef052f3d35588dadf919ee1f0e3cb9b62d3f4b02c").to_vec();
6235
6236        let leaf_node = LeafNode::new(leaf_key, leaf_value);
6237        let leaf_masks = TrieMasks::none();
6238
6239        trie.reveal_nodes(vec![
6240            RevealedSparseNode {
6241                path: Nibbles::from_nibbles([0x3]),
6242                node: TrieNode::Branch(branch_0x3_node),
6243                masks: branch_0x3_masks,
6244            },
6245            RevealedSparseNode {
6246                path: leaf_path,
6247                node: TrieNode::Leaf(leaf_node),
6248                masks: leaf_masks,
6249            },
6250        ])
6251        .unwrap();
6252
6253        // Update leaf with its new value
6254        let mut leaf_full_path = leaf_path;
6255        leaf_full_path.extend(&leaf_key);
6256
6257        let leaf_new_value = vec![
6258            248, 68, 1, 128, 160, 224, 163, 152, 169, 122, 160, 155, 102, 53, 41, 0, 47, 28, 205,
6259            190, 199, 5, 215, 108, 202, 22, 138, 70, 196, 178, 193, 208, 18, 96, 95, 63, 238, 160,
6260            245, 122, 205, 64, 37, 152, 114, 96, 109, 118, 25, 126, 240, 82, 243, 211, 85, 136,
6261            218, 223, 145, 158, 225, 240, 227, 203, 155, 98, 211, 244, 176, 44,
6262        ];
6263
6264        trie.update_leaf(leaf_full_path, leaf_new_value.clone(), DefaultTrieNodeProvider).unwrap();
6265
6266        // Sanity checks before calculating the root
6267        assert_eq!(
6268            Some(&leaf_new_value),
6269            trie.lower_subtrie_for_path(&leaf_path).unwrap().inner.values.get(&leaf_full_path)
6270        );
6271        assert!(trie.upper_subtrie.inner.values.is_empty());
6272
6273        // Assert the root hash matches the expected value
6274        let expected_root =
6275            b256!("0x29b07de8376e9ce7b3a69e9b102199869514d3f42590b5abc6f7d48ec9b8665c");
6276        assert_eq!(trie.root(), expected_root);
6277    }
6278
6279    #[test]
6280    fn find_leaf_existing_leaf() {
6281        // Create a simple trie with one leaf
6282        let provider = DefaultTrieNodeProvider;
6283        let mut sparse = ParallelSparseTrie::default();
6284        let path = Nibbles::from_nibbles([0x1, 0x2, 0x3]);
6285        let value = b"test_value".to_vec();
6286
6287        sparse.update_leaf(path, value.clone(), &provider).unwrap();
6288
6289        // Check that the leaf exists
6290        let result = sparse.find_leaf(&path, None);
6291        assert_matches!(result, Ok(LeafLookup::Exists));
6292
6293        // Check with expected value matching
6294        let result = sparse.find_leaf(&path, Some(&value));
6295        assert_matches!(result, Ok(LeafLookup::Exists));
6296    }
6297
6298    #[test]
6299    fn find_leaf_value_mismatch() {
6300        // Create a simple trie with one leaf
6301        let provider = DefaultTrieNodeProvider;
6302        let mut sparse = ParallelSparseTrie::default();
6303        let path = Nibbles::from_nibbles([0x1, 0x2, 0x3]);
6304        let value = b"test_value".to_vec();
6305        let wrong_value = b"wrong_value".to_vec();
6306
6307        sparse.update_leaf(path, value, &provider).unwrap();
6308
6309        // Check with wrong expected value
6310        let result = sparse.find_leaf(&path, Some(&wrong_value));
6311        assert_matches!(
6312            result,
6313            Err(LeafLookupError::ValueMismatch { path: p, expected: Some(e), actual: _a }) if p == path && e == wrong_value
6314        );
6315    }
6316
6317    #[test]
6318    fn find_leaf_not_found_empty_trie() {
6319        // Empty trie
6320        let sparse = ParallelSparseTrie::default();
6321        let path = Nibbles::from_nibbles([0x1, 0x2, 0x3]);
6322
6323        // Leaf should not exist
6324        let result = sparse.find_leaf(&path, None);
6325        assert_matches!(result, Ok(LeafLookup::NonExistent));
6326    }
6327
6328    #[test]
6329    fn find_leaf_empty_trie() {
6330        let sparse = ParallelSparseTrie::default();
6331        let path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
6332
6333        let result = sparse.find_leaf(&path, None);
6334        assert_matches!(result, Ok(LeafLookup::NonExistent));
6335    }
6336
6337    #[test]
6338    fn find_leaf_exists_no_value_check() {
6339        let provider = DefaultTrieNodeProvider;
6340        let mut sparse = ParallelSparseTrie::default();
6341        let path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
6342        sparse.update_leaf(path, encode_account_value(0), &provider).unwrap();
6343
6344        let result = sparse.find_leaf(&path, None);
6345        assert_matches!(result, Ok(LeafLookup::Exists));
6346    }
6347
6348    #[test]
6349    fn find_leaf_exists_with_value_check_ok() {
6350        let provider = DefaultTrieNodeProvider;
6351        let mut sparse = ParallelSparseTrie::default();
6352        let path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
6353        let value = encode_account_value(0);
6354        sparse.update_leaf(path, value.clone(), &provider).unwrap();
6355
6356        let result = sparse.find_leaf(&path, Some(&value));
6357        assert_matches!(result, Ok(LeafLookup::Exists));
6358    }
6359
6360    #[test]
6361    fn find_leaf_exclusion_branch_divergence() {
6362        let provider = DefaultTrieNodeProvider;
6363        let mut sparse = ParallelSparseTrie::default();
6364        let path1 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); // Creates branch at 0x12
6365        let path2 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x5, 0x6]); // Belongs to same branch
6366        let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x7, 0x8]); // Diverges at nibble 7
6367
6368        sparse.update_leaf(path1, encode_account_value(0), &provider).unwrap();
6369        sparse.update_leaf(path2, encode_account_value(1), &provider).unwrap();
6370
6371        let result = sparse.find_leaf(&search_path, None);
6372        assert_matches!(result, Ok(LeafLookup::NonExistent))
6373    }
6374
6375    #[test]
6376    fn find_leaf_exclusion_extension_divergence() {
6377        let provider = DefaultTrieNodeProvider;
6378        let mut sparse = ParallelSparseTrie::default();
6379        // This will create an extension node at root with key 0x12
6380        let path1 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4, 0x5, 0x6]);
6381        // This path diverges from the extension key
6382        let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x7, 0x8]);
6383
6384        sparse.update_leaf(path1, encode_account_value(0), &provider).unwrap();
6385
6386        let result = sparse.find_leaf(&search_path, None);
6387        assert_matches!(result, Ok(LeafLookup::NonExistent))
6388    }
6389
6390    #[test]
6391    fn find_leaf_exclusion_leaf_divergence() {
6392        let provider = DefaultTrieNodeProvider;
6393        let mut sparse = ParallelSparseTrie::default();
6394        let existing_leaf_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
6395        let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4, 0x5, 0x6]);
6396
6397        sparse.update_leaf(existing_leaf_path, encode_account_value(0), &provider).unwrap();
6398
6399        let result = sparse.find_leaf(&search_path, None);
6400        assert_matches!(result, Ok(LeafLookup::NonExistent))
6401    }
6402
6403    #[test]
6404    fn find_leaf_exclusion_path_ends_at_branch() {
6405        let provider = DefaultTrieNodeProvider;
6406        let mut sparse = ParallelSparseTrie::default();
6407        let path1 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]); // Creates branch at 0x12
6408        let path2 = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x5, 0x6]);
6409        let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2]); // Path of the branch itself
6410
6411        sparse.update_leaf(path1, encode_account_value(0), &provider).unwrap();
6412        sparse.update_leaf(path2, encode_account_value(1), &provider).unwrap();
6413
6414        let result = sparse.find_leaf(&search_path, None);
6415        assert_matches!(result, Ok(LeafLookup::NonExistent));
6416    }
6417
6418    #[test]
6419    fn find_leaf_error_blinded_node_at_leaf_path() {
6420        // Scenario: The node *at* the leaf path is blinded.
6421        let blinded_hash = B256::repeat_byte(0xBB);
6422        let leaf_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
6423
6424        let sparse = new_test_trie(
6425            [
6426                (
6427                    // Ext 0x12
6428                    Nibbles::default(),
6429                    SparseNode::new_ext(Nibbles::from_nibbles_unchecked([0x1, 0x2])),
6430                ),
6431                (
6432                    // Ext 0x123
6433                    Nibbles::from_nibbles_unchecked([0x1, 0x2]),
6434                    SparseNode::new_ext(Nibbles::from_nibbles_unchecked([0x3])),
6435                ),
6436                (
6437                    // Branch at 0x123, child 4
6438                    Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3]),
6439                    SparseNode::new_branch(TrieMask::new(0b10000)),
6440                ),
6441                (
6442                    // Blinded node at 0x1234
6443                    leaf_path,
6444                    SparseNode::Hash(blinded_hash),
6445                ),
6446            ]
6447            .into_iter(),
6448        );
6449
6450        let result = sparse.find_leaf(&leaf_path, None);
6451
6452        // Should error because it hit the blinded node exactly at the leaf path
6453        assert_matches!(result, Err(LeafLookupError::BlindedNode { path, hash })
6454            if path == leaf_path && hash == blinded_hash
6455        );
6456    }
6457
6458    #[test]
6459    fn find_leaf_error_blinded_node() {
6460        let blinded_hash = B256::repeat_byte(0xAA);
6461        let path_to_blind = Nibbles::from_nibbles_unchecked([0x1]);
6462        let search_path = Nibbles::from_nibbles_unchecked([0x1, 0x2, 0x3, 0x4]);
6463
6464        let sparse = new_test_trie(
6465            [
6466                // Root is a branch with child 0x1 (blinded) and 0x5 (revealed leaf)
6467                // So we set Bit 1 and Bit 5 in the state_mask
6468                (Nibbles::default(), SparseNode::new_branch(TrieMask::new(0b100010))),
6469                (path_to_blind, SparseNode::Hash(blinded_hash)),
6470                (
6471                    Nibbles::from_nibbles_unchecked([0x5]),
6472                    SparseNode::new_leaf(Nibbles::from_nibbles_unchecked([0x6, 0x7, 0x8])),
6473                ),
6474            ]
6475            .into_iter(),
6476        );
6477
6478        let result = sparse.find_leaf(&search_path, None);
6479
6480        // Should error because it hit the blinded node at path 0x1
6481        assert_matches!(result, Err(LeafLookupError::BlindedNode { path, hash })
6482            if path == path_to_blind && hash == blinded_hash
6483        );
6484    }
6485}