fix(trie): sparse trie walk should be done in a sorted manner (#12087)

crates/trie/sparse/src/trie.rs

@@ -11,7 +11,7 @@ use reth_trie_common::{
     EMPTY_ROOT_HASH,
 };
 use smallvec::SmallVec;
-use std::{collections::HashSet, fmt};
+use std::fmt;
 
 /// Inner representation of the sparse trie.
 /// Sparse trie is blind by default until nodes are revealed.
@@ -579,29 +579,29 @@ impl RevealedSparseTrie {
     /// Returns a list of paths to the nodes that are located at the provided depth when counting
     /// from the root node. If there's a leaf at a depth less than the provided depth, it will be
     /// included in the result.
-    fn get_nodes_at_depth(&self, depth: usize) -> HashSet<Nibbles> {
+    fn get_nodes_at_depth(&self, depth: usize) -> Vec<Nibbles> {
         let mut paths = Vec::from([(Nibbles::default(), 0)]);
-        let mut targets = HashSet::<Nibbles>::default();
+        let mut targets = Vec::new();
 
         while let Some((mut path, level)) = paths.pop() {
             match self.nodes.get(&path).unwrap() {
                 SparseNode::Empty | SparseNode::Hash(_) => {}
                 SparseNode::Leaf { .. } => {
-                    targets.insert(path);
+                    targets.push(path);
                 }
                 SparseNode::Extension { key, .. } => {
                     if level >= depth {
-                        targets.insert(path);
+                        targets.push(path);
                     } else {
                         path.extend_from_slice_unchecked(key);
                         paths.push((path, level + 1));
                     }
                 }
                 SparseNode::Branch { state_mask, .. } => {
                     if level >= depth {
-                        targets.insert(path);
+                        targets.push(path);
                     } else {
-                        for bit in CHILD_INDEX_RANGE {
+                        for bit in CHILD_INDEX_RANGE.rev() {
                             if state_mask.is_bit_set(bit) {
                                 let mut child_path = path.clone();
                                 child_path.push_unchecked(bit);
@@ -666,21 +666,27 @@ impl RevealedSparseTrie {
                     }
 
                     branch_child_buf.clear();
-                    for bit in CHILD_INDEX_RANGE {
+                    // Walk children in a reverse order from `f` to `0`, so we pop the `0` first
+                    // from the stack.
+                    for bit in CHILD_INDEX_RANGE.rev() {
                         if state_mask.is_bit_set(bit) {
                             let mut child = path.clone();
                             child.push_unchecked(bit);
                             branch_child_buf.push(child);
                         }
                     }
 
-                    branch_value_stack_buf.clear();
-                    for child_path in &branch_child_buf {
+                    branch_value_stack_buf.resize(branch_child_buf.len(), Default::default());
+                    let mut added_children = false;
+                    for (i, child_path) in branch_child_buf.iter().enumerate() {
                         if rlp_node_stack.last().map_or(false, |e| &e.0 == child_path) {
                             let (_, child) = rlp_node_stack.pop().unwrap();
-                            branch_value_stack_buf.push(child);
+                            // Insert children in the resulting buffer in a normal order, because
+                            // initially we iterated in reverse.
+                            branch_value_stack_buf[branch_child_buf.len() - i - 1] = child;
+                            added_children = true;
                         } else {
-                            debug_assert!(branch_value_stack_buf.is_empty());
+                            debug_assert!(!added_children);
                             path_stack.push(path);
                             path_stack.extend(branch_child_buf.drain(..));
                             continue 'main
@@ -1568,38 +1574,35 @@ mod tests {
             .unwrap();
         sparse.update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]), value).unwrap();
 
-        assert_eq!(sparse.get_nodes_at_depth(0), HashSet::from([Nibbles::default()]));
-        assert_eq!(
-            sparse.get_nodes_at_depth(1),
-            HashSet::from([Nibbles::from_nibbles_unchecked([0x5])])
-        );
+        assert_eq!(sparse.get_nodes_at_depth(0), vec![Nibbles::default()]);
+        assert_eq!(sparse.get_nodes_at_depth(1), vec![Nibbles::from_nibbles_unchecked([0x5])]);
         assert_eq!(
             sparse.get_nodes_at_depth(2),
-            HashSet::from([
+            vec![
                 Nibbles::from_nibbles_unchecked([0x5, 0x0]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x2]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x3])
-            ])
+            ]
         );
         assert_eq!(
             sparse.get_nodes_at_depth(3),
-            HashSet::from([
+            vec![
                 Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x2]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x1]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3])
-            ])
+            ]
         );
         assert_eq!(
             sparse.get_nodes_at_depth(4),
-            HashSet::from([
+            vec![
                 Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x1]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x3]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x2]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x1]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3, 0x0]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x3, 0x3, 0x2])
-            ])
+            ]
         );
     }
 }