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/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`.