feat(trie): walker branch node seeks metric (#16106)

crates/trie/db/tests/walker.rs

@@ -60,7 +60,7 @@ fn test_cursor<T>(mut trie: T, expected: &[Vec<u8>])
 where
     T: TrieCursor,
 {
-    let mut walker = TrieWalker::new(&mut trie, Default::default());
+    let mut walker = TrieWalker::state_trie(&mut trie, Default::default());
     assert!(walker.key().unwrap().is_empty());
 
     // We're traversing the path in lexicographical order.
@@ -114,7 +114,7 @@ fn cursor_rootnode_with_changesets() {
     let mut trie = DatabaseStorageTrieCursor::new(cursor, hashed_address);
 
     // No changes
-    let mut cursor = TrieWalker::new(&mut trie, Default::default());
+    let mut cursor = TrieWalker::state_trie(&mut trie, Default::default());
     assert_eq!(cursor.key().cloned(), Some(Nibbles::new())); // root
     assert!(cursor.can_skip_current_node); // due to root_hash
     cursor.advance().unwrap(); // skips to the end of trie
@@ -123,7 +123,7 @@ fn cursor_rootnode_with_changesets() {
     // We insert something that's not part of the existing trie/prefix.
     let mut changed = PrefixSetMut::default();
     changed.insert(Nibbles::from_nibbles([0xF, 0x1]));
-    let mut cursor = TrieWalker::new(&mut trie, changed.freeze());
+    let mut cursor = TrieWalker::state_trie(&mut trie, changed.freeze());
 
     // Root node
     assert_eq!(cursor.key().cloned(), Some(Nibbles::new()));

crates/trie/parallel/src/proof.rs

@@ -209,7 +209,7 @@ where
         );
 
         // Create the walker.
-        let walker = TrieWalker::new(
+        let walker = TrieWalker::state_trie(
             trie_cursor_factory.account_trie_cursor().map_err(ProviderError::Database)?,
             prefix_sets.account_prefix_set,
         )

crates/trie/parallel/src/root.rs

@@ -145,7 +145,7 @@ where
             &hashed_state_sorted,
         );
 
-        let walker = TrieWalker::new(
+        let walker = TrieWalker::state_trie(
             trie_cursor_factory.account_trie_cursor().map_err(ProviderError::Database)?,
             prefix_sets.account_prefix_set,
         )

crates/trie/sparse/benches/root.rs

@@ -131,7 +131,7 @@ fn calculate_root_from_leaves_repeated(c: &mut Criterion) {
                                         )
                                     };
 
-                                let walker = TrieWalker::new(
+                                let walker = TrieWalker::storage_trie(
                                     InMemoryStorageTrieCursor::new(
                                         B256::ZERO,
                                         NoopStorageTrieCursor::default(),

crates/trie/sparse/src/trie.rs

@@ -69,7 +69,7 @@ pub enum SparseTrie<P = DefaultBlindedProvider> {
     /// The trie is blind -- no nodes have been revealed
     ///
     /// This is the default state. In this state,
-    /// the trie cannot be directly queried or modified until nodes are revealed.    
+    /// the trie cannot be directly queried or modified until nodes are revealed.
     #[default]
     Blind,
     /// Some nodes in the Trie have been revealed.
@@ -2405,7 +2405,7 @@ mod tests {
         prefix_set.extend_keys(state.clone().into_iter().map(|(nibbles, _)| nibbles));
         prefix_set.extend_keys(destroyed_accounts.iter().map(Nibbles::unpack));
         let walker =
-            TrieWalker::new(trie_cursor, prefix_set.freeze()).with_deletions_retained(true);
+            TrieWalker::state_trie(trie_cursor, prefix_set.freeze()).with_deletions_retained(true);
         let hashed_post_state = HashedPostState::default()
             .with_accounts(state.into_iter().map(|(nibbles, account)| {
                 (nibbles.pack().into_inner().unwrap().into(), Some(account))

crates/trie/trie/src/metrics.rs

@@ -50,6 +50,8 @@ impl TrieRootMetrics {
 #[derive(Clone, Metrics)]
 #[metrics(scope = "trie.walker")]
 pub struct WalkerMetrics {
+    /// The number of branch nodes seeked by the walker.
+    branch_nodes_seeked_total: Counter,
     /// The number of subnodes out of order due to wrong tree mask.
     out_of_order_subnode: Counter,
 }
@@ -60,6 +62,11 @@ impl WalkerMetrics {
         Self::new_with_labels(&[("type", ty.as_str())])
     }
 
+    /// Increment `branch_nodes_seeked_total`.
+    pub fn inc_branch_nodes_seeked(&self) {
+        self.branch_nodes_seeked_total.increment(1);
+    }
+
     /// 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/node_iter.rs

@@ -330,7 +330,7 @@ mod tests {
 
         let mut prefix_set = PrefixSetMut::default();
         prefix_set.extend_keys(state.clone().into_iter().map(|(nibbles, _)| nibbles));
-        let walker = TrieWalker::new(NoopAccountTrieCursor, prefix_set.freeze());
+        let walker = TrieWalker::state_trie(NoopAccountTrieCursor, prefix_set.freeze());
 
         let hashed_post_state = HashedPostState::default()
             .with_accounts(state.into_iter().map(|(nibbles, account)| {
@@ -460,7 +460,7 @@ mod tests {
         let prefix_set = prefix_set.freeze();
 
         let walker =
-            TrieWalker::new(trie_cursor_factory.account_trie_cursor().unwrap(), prefix_set);
+            TrieWalker::state_trie(trie_cursor_factory.account_trie_cursor().unwrap(), prefix_set);
 
         let hashed_cursor_factory = MockHashedCursorFactory::new(
             BTreeMap::from([

crates/trie/trie/src/proof/mod.rs

@@ -111,7 +111,7 @@ where
         // Create the walker.
         let mut prefix_set = self.prefix_sets.account_prefix_set.clone();
         prefix_set.extend_keys(targets.keys().map(Nibbles::unpack));
-        let walker = TrieWalker::new(trie_cursor, prefix_set.freeze());
+        let walker = TrieWalker::state_trie(trie_cursor, prefix_set.freeze());
 
         // Create a hash builder to rebuild the root node since it is not available in the database.
         let retainer = targets.keys().map(Nibbles::unpack).collect();
@@ -282,7 +282,7 @@ where
         self.prefix_set.extend_keys(target_nibbles.clone());
 
         let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;
-        let walker = TrieWalker::new(trie_cursor, self.prefix_set.freeze());
+        let walker = TrieWalker::storage_trie(trie_cursor, self.prefix_set.freeze());
 
         let retainer = ProofRetainer::from_iter(target_nibbles);
         let mut hash_builder = HashBuilder::default()

crates/trie/trie/src/trie.rs

@@ -159,7 +159,7 @@ where
         let (mut hash_builder, mut account_node_iter) = match self.previous_state {
             Some(state) => {
                 let hash_builder = state.hash_builder.with_updates(retain_updates);
-                let walker = TrieWalker::from_stack(
+                let walker = TrieWalker::state_trie_from_stack(
                     trie_cursor,
                     state.walker_stack,
                     self.prefix_sets.account_prefix_set,
@@ -171,8 +171,9 @@ where
             }
             None => {
                 let hash_builder = HashBuilder::default().with_updates(retain_updates);
-                let walker = TrieWalker::new(trie_cursor, self.prefix_sets.account_prefix_set)
-                    .with_deletions_retained(retain_updates);
+                let walker =
+                    TrieWalker::state_trie(trie_cursor, self.prefix_sets.account_prefix_set)
+                        .with_deletions_retained(retain_updates);
                 let node_iter = TrieNodeIter::state_trie(walker, hashed_account_cursor);
                 (hash_builder, node_iter)
             }
@@ -407,8 +408,8 @@ where
 
         let mut tracker = TrieTracker::default();
         let trie_cursor = self.trie_cursor_factory.storage_trie_cursor(self.hashed_address)?;
-        let walker =
-            TrieWalker::new(trie_cursor, self.prefix_set).with_deletions_retained(retain_updates);
+        let walker = TrieWalker::storage_trie(trie_cursor, self.prefix_set)
+            .with_deletions_retained(retain_updates);
 
         let mut hash_builder = HashBuilder::default().with_updates(retain_updates);
 

crates/trie/trie/src/walker.rs

@@ -33,8 +33,39 @@ pub struct TrieWalker<C> {
 }
 
 impl<C> TrieWalker<C> {
+    /// Constructs a new `TrieWalker` for the state trie from existing stack and a cursor.
+    pub fn state_trie_from_stack(cursor: C, stack: Vec<CursorSubNode>, changes: PrefixSet) -> Self {
+        Self::from_stack(
+            cursor,
+            stack,
+            changes,
+            #[cfg(feature = "metrics")]
+            crate::TrieType::State,
+        )
+    }
+
+    /// Constructs a new `TrieWalker` for the storage trie from existing stack and a cursor.
+    pub fn storage_trie_from_stack(
+        cursor: C,
+        stack: Vec<CursorSubNode>,
+        changes: PrefixSet,
+    ) -> Self {
+        Self::from_stack(
+            cursor,
+            stack,
+            changes,
+            #[cfg(feature = "metrics")]
+            crate::TrieType::Storage,
+        )
+    }
+
     /// Constructs a new `TrieWalker` from existing stack and a cursor.
-    pub fn from_stack(cursor: C, stack: Vec<CursorSubNode>, changes: PrefixSet) -> Self {
+    fn from_stack(
+        cursor: C,
+        stack: Vec<CursorSubNode>,
+        changes: PrefixSet,
+        #[cfg(feature = "metrics")] trie_type: crate::TrieType,
+    ) -> Self {
         let mut this = Self {
             cursor,
             changes,
@@ -42,7 +73,7 @@ impl<C> TrieWalker<C> {
             can_skip_current_node: false,
             removed_keys: None,
             #[cfg(feature = "metrics")]
-            metrics: WalkerMetrics::default(),
+            metrics: WalkerMetrics::new(trie_type),
         };
         this.update_skip_node();
         this
@@ -128,8 +159,32 @@ impl<C> TrieWalker<C> {
 }
 
 impl<C: TrieCursor> TrieWalker<C> {
+    /// Constructs a new [`TrieWalker`] for the state trie.
+    pub fn state_trie(cursor: C, changes: PrefixSet) -> Self {
+        Self::new(
+            cursor,
+            changes,
+            #[cfg(feature = "metrics")]
+            crate::TrieType::State,
+        )
+    }
+
+    /// Constructs a new [`TrieWalker`] for the storage trie.
+    pub fn storage_trie(cursor: C, changes: PrefixSet) -> Self {
+        Self::new(
+            cursor,
+            changes,
+            #[cfg(feature = "metrics")]
+            crate::TrieType::Storage,
+        )
+    }
+
     /// Constructs a new `TrieWalker`, setting up the initial state of the stack and cursor.
-    pub fn new(cursor: C, changes: PrefixSet) -> Self {
+    fn new(
+        cursor: C,
+        changes: PrefixSet,
+        #[cfg(feature = "metrics")] trie_type: crate::TrieType,
+    ) -> Self {
         // Initialize the walker with a single empty stack element.
         let mut this = Self {
             cursor,
@@ -138,7 +193,7 @@ impl<C: TrieCursor> TrieWalker<C> {
             can_skip_current_node: false,
             removed_keys: None,
             #[cfg(feature = "metrics")]
-            metrics: WalkerMetrics::default(),
+            metrics: WalkerMetrics::new(trie_type),
         };
 
         // Set up the root node of the trie in the stack, if it exists.
@@ -188,6 +243,8 @@ impl<C: TrieCursor> TrieWalker<C> {
     fn node(&mut self, exact: bool) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
         let key = self.key().expect("key must exist").clone();
         let entry = if exact { self.cursor.seek_exact(key)? } else { self.cursor.seek(key)? };
+        #[cfg(feature = "metrics")]
+        self.metrics.inc_branch_nodes_seeked();
 
         if let Some((_, node)) = &entry {
             assert!(!node.state_mask.is_empty());