perf(trie): shared-prefix walk to reduce branch pop/push operations

Amp-Thread-ID: https://ampcode.com/threads/T-019bfe25-43f3-75ac-98f7-32bf937b69e1
Co-authored-by: Amp <amp@ampcode.com>

crates/trie/sparse/src/provider.rs

@@ -1,6 +1,6 @@
 //! Traits and default implementations related to retrieval of blinded trie nodes.
 
-use alloy_primitives::{map::HashMap, Bytes, B256};
+use alloy_primitives::{Bytes, B256};
 use reth_execution_errors::SparseTrieError;
 use reth_trie_common::{Nibbles, TrieMask};
 
@@ -35,21 +35,6 @@ pub struct RevealedNode {
 pub trait TrieNodeProvider {
     /// Retrieve trie node by path.
     fn trie_node(&self, path: &Nibbles) -> Result<Option<RevealedNode>, SparseTrieError>;
-
-    /// Batch retrieve trie nodes by paths. Default: sequential fallback.
-    fn trie_nodes_batch(
-        &self,
-        paths: &[Nibbles],
-    ) -> Result<HashMap<Nibbles, RevealedNode>, SparseTrieError> {
-        let mut result = HashMap::default();
-        result.reserve(paths.len());
-        for path in paths {
-            if let Some(node) = self.trie_node(path)? {
-                result.insert(path.clone(), node);
-            }
-        }
-        Ok(result)
-    }
 }
 
 /// Default trie node provider factory that creates [`DefaultTrieNodeProviderFactory`].

crates/trie/trie/src/proof/trie_node.rs

@@ -1,6 +1,6 @@
 use super::{Proof, StorageProof};
 use crate::{hashed_cursor::HashedCursorFactory, trie_cursor::TrieCursorFactory};
-use alloy_primitives::{map::{B256Set, HashMap, HashSet}, B256};
+use alloy_primitives::{map::HashSet, B256};
 use reth_execution_errors::{SparseTrieError, SparseTrieErrorKind};
 use reth_trie_common::{MultiProofTargets, Nibbles};
 use reth_trie_sparse::provider::{
@@ -94,48 +94,6 @@ where
         );
         Ok(node.map(|node| RevealedNode { node, tree_mask, hash_mask }))
     }
-
-    fn trie_nodes_batch(
-        &self,
-        paths: &[Nibbles],
-    ) -> Result<HashMap<Nibbles, RevealedNode>, SparseTrieError> {
-        if paths.is_empty() {
-            return Ok(HashMap::default());
-        }
-
-        let start = enabled!(target: "trie::proof::blinded", Level::TRACE).then(Instant::now);
-
-        let targets: MultiProofTargets =
-            paths.iter().map(|p| (pad_path_to_key(p), HashSet::default())).collect();
-
-        let mut proof = Proof::new(&self.trie_cursor_factory, &self.hashed_cursor_factory)
-            .with_branch_node_masks(true)
-            .multiproof(targets)
-            .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?;
-
-        let mut account_subtree = proof.account_subtree.into_inner();
-        let mut result = HashMap::default();
-        result.reserve(paths.len());
-
-        for path in paths {
-            if let Some(node) = account_subtree.remove(path) {
-                let masks = proof.branch_node_masks.remove(path);
-                let hash_mask = masks.map(|m| m.hash_mask);
-                let tree_mask = masks.map(|m| m.tree_mask);
-                result.insert(path.clone(), RevealedNode { node, tree_mask, hash_mask });
-            }
-        }
-
-        trace!(
-            target: "trie::proof::blinded",
-            elapsed = ?start.unwrap().elapsed(),
-            paths_count = paths.len(),
-            results_count = result.len(),
-            "Batch blinded nodes for account trie"
-        );
-
-        Ok(result)
-    }
 }
 
 /// Blinded provider for retrieving storage trie nodes by path.
@@ -190,50 +148,4 @@ where
         );
         Ok(node.map(|node| RevealedNode { node, tree_mask, hash_mask }))
     }
-
-    fn trie_nodes_batch(
-        &self,
-        paths: &[Nibbles],
-    ) -> Result<HashMap<Nibbles, RevealedNode>, SparseTrieError> {
-        if paths.is_empty() {
-            return Ok(HashMap::default());
-        }
-
-        let start = enabled!(target: "trie::proof::blinded", Level::TRACE).then(Instant::now);
-
-        let targets: B256Set = paths.iter().map(pad_path_to_key).collect();
-
-        let mut proof = StorageProof::new_hashed(
-            &self.trie_cursor_factory,
-            &self.hashed_cursor_factory,
-            self.account,
-        )
-        .with_branch_node_masks(true)
-        .storage_multiproof(targets)
-        .map_err(|error| SparseTrieErrorKind::Other(Box::new(error)))?;
-
-        let mut subtree = proof.subtree.into_inner();
-        let mut result = HashMap::default();
-        result.reserve(paths.len());
-
-        for path in paths {
-            if let Some(node) = subtree.remove(path) {
-                let masks = proof.branch_node_masks.remove(path);
-                let hash_mask = masks.map(|m| m.hash_mask);
-                let tree_mask = masks.map(|m| m.tree_mask);
-                result.insert(path.clone(), RevealedNode { node, tree_mask, hash_mask });
-            }
-        }
-
-        trace!(
-            target: "trie::proof::blinded",
-            account = ?self.account,
-            elapsed = ?start.unwrap().elapsed(),
-            paths_count = paths.len(),
-            results_count = result.len(),
-            "Batch blinded nodes for storage trie"
-        );
-
-        Ok(result)
-    }
 }

crates/trie/trie/src/proof_v2/mod.rs

@@ -567,15 +567,31 @@ where
         targets: &mut TargetsCursor<'a>,
         key: Nibbles,
         val: VE::DeferredEncoder,
+    ) -> Result<(), StateProofError> {
+        self.push_leaf_with_min_prefix(targets, key, val, 0)
+    }
+
+    /// Adds a single leaf for a key to the stack with a minimum prefix hint.
+    ///
+    /// The `min_keep_prefix` parameter hints that branches at or above this depth will be needed
+    /// by subsequent leaves and should not be popped. This avoids redundant pop/push operations
+    /// when processing leaves that share a common prefix.
+    fn push_leaf_with_min_prefix<'a>(
+        &mut self,
+        targets: &mut TargetsCursor<'a>,
+        key: Nibbles,
+        val: VE::DeferredEncoder,
+        min_keep_prefix: usize,
     ) -> Result<(), StateProofError> {
         loop {
             trace!(
                 target: TRACE_TARGET,
                 ?key,
+                ?min_keep_prefix,
                 branch_stack_len = ?self.branch_stack.len(),
                 branch_path = ?self.branch_path,
                 child_stack_len = ?self.child_stack.len(),
-                "push_leaf: loop",
+                "push_leaf_with_min_prefix: loop",
             );
 
             // Get the `state_mask` of the branch currently being built. If there are no branches
@@ -612,7 +628,21 @@ where
             // If the current branch does not share all of its nibbles with the new key then it is
             // not the parent of the new key. In this case the current branch will have no more
             // children. We can pop it and loop back to the top to try again with its parent branch.
+            //
+            // However, if we're at or below the min_keep_prefix depth, we should not pop branches
+            // that will be needed by subsequent leaves sharing this prefix.
             if common_prefix_len < self.branch_path.len() {
+                if self.branch_path.len() <= min_keep_prefix {
+                    // Branch is at or above the shared prefix depth - don't pop it
+                    // This leaf doesn't belong under this branch, which shouldn't happen
+                    // if min_keep_prefix is set correctly by the caller
+                    debug_assert!(
+                        false,
+                        "push_leaf_with_min_prefix called with key {key:?} that doesn't share \
+                         min_keep_prefix {min_keep_prefix} with branch_path {:?}",
+                        self.branch_path
+                    );
+                }
                 self.pop_branch(targets)?;
                 continue
             }

crates/trie/trie/src/proof_v2/target.rs

@@ -2,6 +2,49 @@ use crate::proof_v2::increment_and_strip_trailing_zeros;
 use alloy_primitives::B256;
 use reth_trie_common::Nibbles;
 
+/// Groups sorted targets by common prefix for efficient batch traversal.
+#[derive(Debug)]
+pub struct PrefixBatch<'a> {
+    /// The common prefix shared by all targets in this batch.
+    pub common_prefix: Nibbles,
+    /// The targets belonging to this batch.
+    pub targets: &'a [Nibbles],
+}
+
+impl<'a> PrefixBatch<'a> {
+    /// Group sorted nibble paths by their longest common prefix.
+    pub fn from_sorted(targets: &'a [Nibbles]) -> Vec<PrefixBatch<'a>> {
+        if targets.is_empty() {
+            return Vec::new();
+        }
+
+        let mut result = Vec::new();
+        let mut start = 0;
+
+        while start < targets.len() {
+            let first = &targets[start];
+            let mut end = start + 1;
+            let mut prefix_len = first.len();
+
+            while end < targets.len() {
+                let common = first.common_prefix_length(&targets[end]);
+                if common == 0 {
+                    break;
+                }
+                prefix_len = prefix_len.min(common);
+                end += 1;
+            }
+
+            result.push(PrefixBatch {
+                common_prefix: first.slice(..prefix_len),
+                targets: &targets[start..end],
+            });
+            start = end;
+        }
+        result
+    }
+}
+
 /// Target describes a proof target. For every proof target given, the
 /// [`crate::proof_v2::ProofCalculator`] will calculate and return all nodes whose path is a prefix
 /// of the target's `key`.