feat(trie): Proof Rewrite: Use cached branch nodes (#20075)

Co-authored-by: YK <chiayongkang@hotmail.com>
Co-authored-by: Alexey Shekhirin <5773434+shekhirin@users.noreply.github.com>

Cargo.lock

@@ -10895,6 +10895,7 @@ dependencies = [
  "pretty_assertions",
  "proptest",
  "proptest-arbitrary-interop",
+ "rand 0.9.2",
  "reth-ethereum-primitives",
  "reth-execution-errors",
  "reth-metrics",

crates/trie/trie/Cargo.toml

@@ -64,6 +64,7 @@ parking_lot.workspace = true
 pretty_assertions.workspace = true
 proptest-arbitrary-interop.workspace = true
 proptest.workspace = true
+rand.workspace = true
 
 [features]
 metrics = ["reth-metrics", "dep:metrics"]
@@ -84,6 +85,7 @@ serde = [
     "revm-state/serde",
     "parking_lot/serde",
     "reth-ethereum-primitives/serde",
+    "rand/serde",
 ]
 test-utils = [
     "triehash",

crates/trie/trie/benches/proof_v2.rs

@@ -11,20 +11,19 @@ use reth_trie::{
     proof_v2::StorageProofCalculator,
     trie_cursor::{mock::MockTrieCursorFactory, TrieCursorFactory},
 };
-use reth_trie_common::{HashedPostState, HashedStorage, Nibbles};
-use std::collections::BTreeMap;
+use reth_trie_common::{HashedPostState, HashedStorage};
 
 /// Generate test data for benchmarking.
 ///
 /// Returns a tuple of:
 /// - Hashed address for the storage trie
 /// - `HashedPostState` with random storage slots
-/// - Proof targets (Nibbles) that are 80% from existing slots, 20% random
+/// - Proof targets as B256 (sorted) for V2 implementation
 /// - Equivalent [`B256Set`] for legacy implementation
 fn generate_test_data(
     dataset_size: usize,
     num_targets: usize,
-) -> (B256, HashedPostState, Vec<Nibbles>, B256Set) {
+) -> (B256, HashedPostState, Vec<B256>, B256Set) {
     let mut runner = TestRunner::deterministic();
 
     // Use a fixed hashed address for the storage trie
@@ -68,14 +67,8 @@ fn generate_test_data(
 
     let target_b256s = targets_strategy.new_tree(&mut runner).unwrap().current();
 
-    // Convert B256 targets to sorted Nibbles for V2
-    let mut targets: Vec<Nibbles> = target_b256s
-        .iter()
-        .map(|b256| {
-            // SAFETY: B256 is exactly 32 bytes
-            unsafe { Nibbles::unpack_unchecked(b256.as_slice()) }
-        })
-        .collect();
+    // Sort B256 targets for V2 (storage_proof expects sorted targets)
+    let mut targets: Vec<B256> = target_b256s.clone();
     targets.sort();
 
     // Create B256Set for legacy
@@ -86,19 +79,42 @@ fn generate_test_data(
 
 /// Create cursor factories from a `HashedPostState` for storage trie testing.
 ///
-/// This mimics the test harness pattern from the `proof_v2` tests.
+/// This mimics the test harness pattern from the `proof_v2` tests by using `StateRoot`
+/// to generate `TrieUpdates` from the `HashedPostState`.
 fn create_cursor_factories(
     post_state: &HashedPostState,
 ) -> (MockTrieCursorFactory, MockHashedCursorFactory) {
-    // Ensure that there's a storage trie dataset for every storage trie, even if empty
-    let storage_trie_nodes: B256Map<BTreeMap<_, _>> =
-        post_state.storages.keys().copied().map(|addr| (addr, Default::default())).collect();
+    use reth_trie::{updates::StorageTrieUpdates, StateRoot};
+
+    // Create empty trie cursor factory to serve as the initial state for StateRoot
+    // Ensure that there's a storage trie dataset for every storage account
+    let storage_tries: B256Map<_> = post_state
+        .storages
+        .keys()
+        .copied()
+        .map(|addr| (addr, StorageTrieUpdates::default()))
+        .collect();
+
+    let empty_trie_cursor_factory =
+        MockTrieCursorFactory::from_trie_updates(reth_trie_common::updates::TrieUpdates {
+            storage_tries: storage_tries.clone(),
+            ..Default::default()
+        });
 
     // Create mock hashed cursor factory from the post state
     let hashed_cursor_factory = MockHashedCursorFactory::from_hashed_post_state(post_state.clone());
 
-    // Create empty trie cursor factory (leaf-only calculator doesn't need trie nodes)
-    let trie_cursor_factory = MockTrieCursorFactory::new(BTreeMap::new(), storage_trie_nodes);
+    // Generate TrieUpdates using StateRoot
+    let (_root, mut trie_updates) =
+        StateRoot::new(empty_trie_cursor_factory, hashed_cursor_factory.clone())
+            .root_with_updates()
+            .expect("StateRoot should succeed");
+
+    // Continue using empty storage tries for each account
+    trie_updates.storage_tries = storage_tries;
+
+    // Initialize trie cursor factory from the generated TrieUpdates
+    let trie_cursor_factory = MockTrieCursorFactory::from_trie_updates(trie_updates);
 
     (trie_cursor_factory, hashed_cursor_factory)
 }
@@ -148,7 +164,7 @@ fn bench_proof_algos(c: &mut Criterion) {
                     || targets.clone(),
                     |targets| {
                         proof_calculator
-                            .storage_proof(hashed_address, targets.into_iter())
+                            .storage_proof(hashed_address, targets)
                             .expect("Proof generation failed");
                     },
                     BatchSize::SmallInput,

crates/trie/trie/src/hashed_cursor/mock.rs

@@ -48,7 +48,7 @@ impl MockHashedCursorFactory {
             .collect();
 
         // Extract storages from post state
-        let hashed_storages: B256Map<BTreeMap<B256, U256>> = post_state
+        let mut hashed_storages: B256Map<BTreeMap<B256, U256>> = post_state
             .storages
             .into_iter()
             .map(|(addr, hashed_storage)| {
@@ -62,6 +62,11 @@ impl MockHashedCursorFactory {
             })
             .collect();
 
+        // Ensure all accounts have at least an empty storage
+        for account in hashed_accounts.keys() {
+            hashed_storages.entry(*account).or_default();
+        }
+
         Self::new(hashed_accounts, hashed_storages)
     }
 

crates/trie/trie/src/proof_v2/node.rs

@@ -25,7 +25,12 @@ pub(crate) enum ProofTrieBranchChild<RF> {
         child: RlpNode,
     },
     /// A branch node whose children have already been flattened into [`RlpNode`]s.
-    Branch(BranchNode),
+    Branch {
+        /// The node itself, for use during RLP encoding.
+        node: BranchNode,
+        /// Bitmasks carried over from cached `BranchNodeCompact` values, if any.
+        masks: TrieMasks,
+    },
     /// A node whose type is not known, as it has already been converted to an [`RlpNode`].
     RlpNode(RlpNode),
 }
@@ -64,7 +69,7 @@ impl<RF: DeferredValueEncoder> ProofTrieBranchChild<RF> {
                 ExtensionNodeRef::new(&short_key, child.as_slice()).encode(buf);
                 Ok((RlpNode::from_rlp(buf), None))
             }
-            Self::Branch(branch_node) => {
+            Self::Branch { node: branch_node, .. } => {
                 branch_node.encode(buf);
                 Ok((RlpNode::from_rlp(buf), Some(branch_node.stack)))
             }
@@ -98,8 +103,7 @@ impl<RF: DeferredValueEncoder> ProofTrieBranchChild<RF> {
             Self::Extension { short_key, child } => {
                 (TrieNode::Extension(ExtensionNode { key: short_key, child }), TrieMasks::none())
             }
-            // TODO store trie masks on branch
-            Self::Branch(branch_node) => (TrieNode::Branch(branch_node), TrieMasks::none()),
+            Self::Branch { node, masks } => (TrieNode::Branch(node), masks),
             Self::RlpNode(_) => panic!("Cannot call `into_proof_trie_node` on RlpNode"),
         };
 
@@ -111,7 +115,7 @@ impl<RF: DeferredValueEncoder> ProofTrieBranchChild<RF> {
     pub(crate) fn short_key(&self) -> &Nibbles {
         match self {
             Self::Leaf { short_key, .. } | Self::Extension { short_key, .. } => short_key,
-            Self::Branch(_) | Self::RlpNode(_) => {
+            Self::Branch { .. } | Self::RlpNode(_) => {
                 static EMPTY_NIBBLES: Nibbles = Nibbles::new();
                 &EMPTY_NIBBLES
             }
@@ -136,7 +140,7 @@ impl<RF: DeferredValueEncoder> ProofTrieBranchChild<RF> {
             Self::Leaf { short_key, .. } | Self::Extension { short_key, .. } => {
                 *short_key = trim_nibbles_prefix(short_key, len);
             }
-            Self::Branch(_) | Self::RlpNode(_) => {
+            Self::Branch { .. } | Self::RlpNode(_) => {
                 panic!("Cannot call `trim_short_key_prefix` on Branch or RlpNode")
             }
         }
@@ -153,14 +157,8 @@ pub(crate) struct ProofTrieBranch {
     /// A mask tracking which child nibbles are set on the branch so far. There will be a single
     /// child on the stack for each set bit.
     pub(crate) state_mask: TrieMask,
-    /// A subset of `state_mask`. Each bit is set if the `state_mask` bit is set and:
-    /// - The child is a branch which is stored in the DB.
-    /// - The child is an extension whose child branch is stored in the DB.
-    #[expect(unused)]
-    pub(crate) tree_mask: TrieMask,
-    /// A subset of `state_mask`. Each bit is set if the hash for the child is cached in the DB.
-    #[expect(unused)]
-    pub(crate) hash_mask: TrieMask,
+    /// Bitmasks which are subsets of `state_mask`.
+    pub(crate) masks: TrieMasks,
 }
 
 /// Trims the first `len` nibbles from the head of the given `Nibbles`.

crates/trie/trie/src/proof_v2/value.rs

@@ -7,7 +7,6 @@ use alloy_primitives::{B256, U256};
 use alloy_rlp::Encodable;
 use reth_execution_errors::trie::StateProofError;
 use reth_primitives_traits::Account;
-use reth_trie_common::Nibbles;
 use std::rc::Rc;
 
 /// A trait for deferred RLP-encoding of leaf values.
@@ -124,7 +123,7 @@ where
         // Compute storage root by calling storage_proof with the root path as a target.
         // This returns just the root node of the storage trie.
         let storage_root = storage_proof_calculator
-            .storage_proof(self.hashed_address, [Nibbles::new()])
+            .storage_proof(self.hashed_address, [B256::ZERO])
             .map(|nodes| {
                 // Encode the root node to RLP and hash it
                 let root_node =