perf(trie): optimize TrieNodeIter by skipping redundant seek (#15841)

Signed-off-by: 7suyash7 <[email protected]>
Co-authored-by: Alexey Shekhirin <[email protected]>

Author: 7suyash7

crates/trie/trie/src/metrics.rs

@@ -83,8 +83,6 @@ pub struct TrieNodeIterMetrics {
     /// iterator. It does not mean the database seek was actually done, as the trie node
     /// iterator caches the last hashed cursor seek.
     leaf_nodes_same_seeked_total: Counter,
-    /// The number of times the same leaf node as we just advanced to was seeked by the iterator.
-    leaf_nodes_same_seeked_as_advanced_total: Counter,
     /// The number of leaf nodes seeked by the iterator.
     leaf_nodes_seeked_total: Counter,
     /// The number of leaf nodes advanced by the iterator.
@@ -109,11 +107,6 @@ impl TrieNodeIterMetrics {
         self.leaf_nodes_same_seeked_total.increment(1);
     }
 
-    /// Increment `leaf_nodes_same_seeked_as_advanced_total`.
-    pub fn inc_leaf_nodes_same_seeked_as_advanced(&self) {
-        self.leaf_nodes_same_seeked_as_advanced_total.increment(1);
-    }
-
     /// Increment `leaf_nodes_seeked_total`.
     pub fn inc_leaf_nodes_seeked(&self) {
         self.leaf_nodes_seeked_total.increment(1);

crates/trie/trie/src/node_iter.rs

@@ -64,9 +64,10 @@ pub struct TrieNodeIter<C, H: HashedCursor> {
 
     #[cfg(feature = "metrics")]
     metrics: crate::metrics::TrieNodeIterMetrics,
-    /// The key that the [`HashedCursor`] previously advanced to using [`HashedCursor::next`].
-    #[cfg(feature = "metrics")]
-    previously_advanced_to_key: Option<B256>,
+    /// Stores the result of the last successful [`Self::next_hashed_entry`], used to avoid a
+    /// redundant [`Self::seek_hashed_entry`] call if the walker points to the same key that
+    /// was just returned by `next()`.
+    last_next_result: Option<(B256, H::Value)>,
 }
 
 impl<C, H: HashedCursor> TrieNodeIter<C, H>
@@ -108,8 +109,7 @@ where
             last_seeked_hashed_entry: None,
             #[cfg(feature = "metrics")]
             metrics: crate::metrics::TrieNodeIterMetrics::new(trie_type),
-            #[cfg(feature = "metrics")]
-            previously_advanced_to_key: None,
+            last_next_result: None,
         }
     }
 
@@ -126,6 +126,18 @@ where
     ///
     /// If `metrics` feature is enabled, also updates the metrics.
     fn seek_hashed_entry(&mut self, key: B256) -> Result<Option<(B256, H::Value)>, DatabaseError> {
+        if let Some((last_key, last_value)) = self.last_next_result {
+            if last_key == key {
+                trace!(target: "trie::node_iter", seek_key = ?key, "reusing result from last next() call instead of seeking");
+                self.last_next_result = None; // Consume the cached value
+
+                let result = Some((last_key, last_value));
+                self.last_seeked_hashed_entry = Some(SeekedHashedEntry { seeked_key: key, result });
+
+                return Ok(result);
+            }
+        }
+
         if let Some(entry) = self
             .last_seeked_hashed_entry
             .as_ref()
@@ -137,16 +149,13 @@ where
             return Ok(entry);
         }
 
+        trace!(target: "trie::node_iter", ?key, "performing hashed cursor seek");
         let result = self.hashed_cursor.seek(key)?;
         self.last_seeked_hashed_entry = Some(SeekedHashedEntry { seeked_key: key, result });
 
         #[cfg(feature = "metrics")]
         {
             self.metrics.inc_leaf_nodes_seeked();
-
-            if Some(key) == self.previously_advanced_to_key {
-                self.metrics.inc_leaf_nodes_same_seeked_as_advanced();
-            }
         }
         Ok(result)
     }
@@ -156,12 +165,12 @@ where
     /// If `metrics` feature is enabled, also updates the metrics.
     fn next_hashed_entry(&mut self) -> Result<Option<(B256, H::Value)>, DatabaseError> {
         let result = self.hashed_cursor.next();
+
+        self.last_next_result = result.clone()?;
+
         #[cfg(feature = "metrics")]
         {
             self.metrics.inc_leaf_nodes_advanced();
-
-            self.previously_advanced_to_key =
-                result.as_ref().ok().and_then(|result| result.as_ref().map(|(k, _)| *k));
         }
         result
     }
@@ -518,12 +527,6 @@ mod tests {
                 // Collect the siblings of the modified account
                 KeyVisit { visit_type: KeyVisitType::Next, visited_key: Some(account_4) },
                 KeyVisit { visit_type: KeyVisitType::Next, visited_key: Some(account_5) },
-                // We seek the account 5 because its hash is not in the branch node, but we already
-                // walked it before, so there should be no need for it.
-                KeyVisit {
-                    visit_type: KeyVisitType::SeekNonExact(account_5),
-                    visited_key: Some(account_5)
-                },
                 KeyVisit { visit_type: KeyVisitType::Next, visited_key: None },
             ],
         );