fix(trie): move to sibling on invalid tree mask (paradigmxyz#12193)

Co-authored-by: Federico Gimenez <[email protected]>

crates/trie/trie/src/metrics.rs

@@ -1,5 +1,5 @@
 use crate::stats::TrieStats;
-use metrics::Histogram;
+use metrics::{Counter, Histogram};
 use reth_metrics::Metrics;
 
 /// Wrapper for state root metrics.
@@ -63,3 +63,23 @@ impl TrieType {
         }
     }
 }
+
+/// Metrics for trie walker
+#[derive(Clone, Metrics)]
+#[metrics(scope = "trie.walker")]
+pub struct WalkerMetrics {
+    /// The number of subnodes out of order due to wrong tree mask.
+    out_of_order_subnode: Counter,
+}
+
+impl WalkerMetrics {
+    /// Create new metrics for the given trie type.
+    pub fn new(ty: TrieType) -> Self {
+        Self::new_with_labels(&[("type", ty.as_str())])
+    }
+
+    /// Increment `out_of_order_subnode`.
+    pub fn inc_out_of_order_subnode(&self, amount: u64) {
+        self.out_of_order_subnode.increment(amount);
+    }
+}

crates/trie/trie/src/walker.rs

@@ -7,6 +7,9 @@ use alloy_primitives::B256;
 use reth_storage_errors::db::DatabaseError;
 use std::collections::HashSet;
 
+#[cfg(feature = "metrics")]
+use crate::metrics::WalkerMetrics;
+
 /// `TrieWalker` is a structure that enables traversal of a Merkle trie.
 /// It allows moving through the trie in a depth-first manner, skipping certain branches
 /// if they have not changed.
@@ -24,13 +27,23 @@ pub struct TrieWalker<C> {
     pub changes: PrefixSet,
     /// The retained trie node keys that need to be removed.
     removed_keys: Option<HashSet<Nibbles>>,
+    #[cfg(feature = "metrics")]
+    /// Walker metrics.
+    metrics: WalkerMetrics,
 }
 
 impl<C> TrieWalker<C> {
     /// Constructs a new `TrieWalker` from existing stack and a cursor.
     pub fn from_stack(cursor: C, stack: Vec<CursorSubNode>, changes: PrefixSet) -> Self {
-        let mut this =
-            Self { cursor, changes, stack, can_skip_current_node: false, removed_keys: None };
+        let mut this = Self {
+            cursor,
+            changes,
+            stack,
+            can_skip_current_node: false,
+            removed_keys: None,
+            #[cfg(feature = "metrics")]
+            metrics: WalkerMetrics::default(),
+        };
         this.update_skip_node();
         this
     }
@@ -113,6 +126,8 @@ impl<C: TrieCursor> TrieWalker<C> {
             stack: vec![CursorSubNode::default()],
             can_skip_current_node: false,
             removed_keys: None,
+            #[cfg(feature = "metrics")]
+            metrics: WalkerMetrics::default(),
         };
 
         // Set up the root node of the trie in the stack, if it exists.
@@ -179,6 +194,19 @@ impl<C: TrieCursor> TrieWalker<C> {
             self.stack[0].set_nibble(key[0] as i8);
         }
 
+        // The current tree mask might have been set incorrectly.
+        // Sanity check that the newly retrieved trie node key is the child of the last item
+        // on the stack. If not, advance to the next sibling instead of adding the node to the
+        // stack.
+        if let Some(subnode) = self.stack.last() {
+            if !key.starts_with(subnode.full_key()) {
+                #[cfg(feature = "metrics")]
+                self.metrics.inc_out_of_order_subnode(1);
+                self.move_to_next_sibling(false)?;
+                return Ok(())
+            }
+        }
+
         // Create a new CursorSubNode and push it to the stack.
         let subnode = CursorSubNode::new(key, Some(node));
         let nibble = subnode.nibble();