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

@@ -637,6 +637,73 @@ where
         }
     }
 
+    /// Like [`Self::push_leaf`], but won't pop branches at or above `min_keep_prefix`.
+    ///
+    /// This optimization reduces unnecessary branch pop/push operations when processing
+    /// keys with shared prefixes within sub-tries.
+    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_with_min_prefix: loop",
+            );
+
+            let curr_branch_state_mask = match self.branch_stack.last() {
+                Some(curr_branch) => curr_branch.state_mask,
+                None if self.child_stack.is_empty() => {
+                    self.child_stack
+                        .push(ProofTrieBranchChild::Leaf { short_key: key, value: val });
+                    return Ok(())
+                }
+                None => {
+                    debug_assert_eq!(self.child_stack.len(), 1);
+                    debug_assert!(!self
+                        .child_stack
+                        .last()
+                        .expect("already checked for emptiness")
+                        .short_key()
+                        .is_empty());
+                    let (nibble, short_key) = self.push_new_branch(key);
+                    self.push_new_leaf(targets, nibble, short_key, val)?;
+                    return Ok(())
+                }
+            };
+
+            let common_prefix_len = self.branch_path.common_prefix_length(&key);
+
+            if common_prefix_len < self.branch_path.len() {
+                if common_prefix_len >= min_keep_prefix {
+                    self.pop_branch(targets)?;
+                    continue
+                }
+                self.pop_branch(targets)?;
+                continue
+            }
+
+            let nibble = key.get_unchecked(common_prefix_len);
+            if curr_branch_state_mask.is_bit_set(nibble) {
+                let (nibble, short_key) = self.push_new_branch(key);
+                self.push_new_leaf(targets, nibble, short_key, val)?;
+            } else {
+                let short_key = key.slice_unchecked(common_prefix_len + 1, key.len());
+                self.push_new_leaf(targets, nibble, short_key, val)?;
+            }
+
+            return Ok(())
+        }
+    }
+
     /// Given the lower and upper bounds (exclusive) of a range of keys, iterates over the
     /// `hashed_cursor` and calculates all trie nodes possible based on those keys. If the upper
     /// bound is None then it is considered unbounded.

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

@@ -96,6 +96,47 @@ impl<'a> SubTrieTargets<'a> {
     }
 }
 
+/// Helper to group targets by common prefix for efficient traversal.
+#[derive(Debug)]
+pub struct PrefixBatch<'a> {
+    /// The common prefix shared by all targets in this batch
+    pub common_prefix: Nibbles,
+    /// The targets in this batch
+    pub targets: &'a [Nibbles],
+}
+
+impl<'a> PrefixBatch<'a> {
+    /// Group sorted targets 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
+    }
+}
+
 /// Given a set of [`Target`]s, returns an iterator over those same [`Target`]s chunked by the
 /// sub-tries they apply to within the overall trie.
 pub(crate) fn iter_sub_trie_targets<'a>(