feat(trie): Add root_node method to v2 ProofCalculator (#21906)

.changelog/zesty-birds-smile.md

@@ -0,0 +1,6 @@
+---
+reth-trie: minor
+reth-trie-parallel: minor
+---
+
+Added `root_node` and `storage_root_node` methods to proof calculators for efficient root-only calculations. These methods directly return the root node without requiring dummy targets, replacing the previous workaround of passing fake targets to proof generation.

crates/trie/parallel/src/proof_task.rs

@@ -528,13 +528,6 @@ where
             panic!("compute_v2_storage_proof only accepts StorageProofInput::V2")
         };
 
-        // If targets is empty it means the caller only wants the root hash. The V2 proof calculator
-        // will do nothing given no targets, so instead we give it a fake target so it always
-        // returns at least the root.
-        if targets.is_empty() {
-            targets.push(proof_v2::Target::new(B256::ZERO));
-        }
-
         let span = debug_span!(
             target: "trie::proof_task",
             "V2 Storage proof calculation",
@@ -545,7 +538,15 @@ where
         let _span_guard = span.enter();
 
         let proof_start = Instant::now();
-        let proof = calculator.storage_proof(hashed_address, &mut targets)?;
+
+        // If targets is empty it means the caller only wants the root node.
+        let proof = if targets.is_empty() {
+            let root_node = calculator.storage_root_node(hashed_address)?;
+            vec![root_node]
+        } else {
+            calculator.storage_proof(hashed_address, &mut targets)?
+        };
+
         let root = calculator.compute_root_hash(&proof)?;
 
         trace!(

crates/trie/parallel/src/value_encoder.rs

@@ -177,11 +177,10 @@ where
                         stats.borrow_mut().dispatched_missing_root_count += 1;
 
                         let mut calculator = storage_calculator.borrow_mut();
-                        let proof =
-                            calculator.storage_proof(hashed_address, &mut [B256::ZERO.into()])?;
+                        let root_node = calculator.storage_root_node(hashed_address)?;
                         let storage_root = calculator
-                            .compute_root_hash(&proof)?
-                            .expect("storage_proof with dummy target always returns root");
+                            .compute_root_hash(&[root_node])?
+                            .expect("storage_root_node returns a node at empty path");
 
                         cached_storage_roots.insert(hashed_address, storage_root);
                         storage_root
@@ -195,10 +194,10 @@ where
                 let hashed_address = *hashed_address;
                 let account = *account;
                 let mut calculator = storage_calculator.borrow_mut();
-                let proof = calculator.storage_proof(hashed_address, &mut [B256::ZERO.into()])?;
+                let root_node = calculator.storage_root_node(hashed_address)?;
                 let storage_root = calculator
-                    .compute_root_hash(&proof)?
-                    .expect("storage_proof with dummy target always returns root");
+                    .compute_root_hash(&[root_node])?
+                    .expect("storage_root_node returns a node at empty path");
 
                 cached_storage_roots.insert(hashed_address, storage_root);
                 (account, storage_root)

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

@@ -182,10 +182,13 @@ where
     /// ```
     fn should_retain<'a>(
         &self,
-        targets: &mut TargetsCursor<'a>,
+        targets: &mut Option<TargetsCursor<'a>>,
         path: &Nibbles,
         check_min_len: bool,
     ) -> bool {
+        // If no targets are given then we never retain anything
+        let Some(targets) = targets.as_mut() else { return false };
+
         let (mut lower, mut upper) = targets.current();
 
         trace!(target: TRACE_TARGET, ?path, target = ?lower, "should_retain: called");
@@ -252,7 +255,7 @@ where
     /// therefore can be retained as a proof node if applicable.
     fn commit_child<'a>(
         &mut self,
-        targets: &mut TargetsCursor<'a>,
+        targets: &mut Option<TargetsCursor<'a>>,
         child_path: Nibbles,
         child: ProofTrieBranchChild<VE::DeferredEncoder>,
     ) -> Result<RlpNode, StateProofError> {
@@ -339,7 +342,7 @@ where
     /// to this method.
     fn commit_last_child<'a>(
         &mut self,
-        targets: &mut TargetsCursor<'a>,
+        targets: &mut Option<TargetsCursor<'a>>,
     ) -> Result<(), StateProofError> {
         let Some(child_path) = self.last_child_path() else { return Ok(()) };
         let child =
@@ -373,7 +376,7 @@ where
     /// - If the leaf's nibble is already set in the branch's `state_mask`.
     fn push_new_leaf<'a>(
         &mut self,
-        targets: &mut TargetsCursor<'a>,
+        targets: &mut Option<TargetsCursor<'a>>,
         leaf_nibble: u8,
         leaf_short_key: Nibbles,
         leaf_val: VE::DeferredEncoder,
@@ -478,7 +481,10 @@ where
     /// # Panics
     ///
     /// This method panics if `branch_stack` is empty.
-    fn pop_branch<'a>(&mut self, targets: &mut TargetsCursor<'a>) -> Result<(), StateProofError> {
+    fn pop_branch<'a>(
+        &mut self,
+        targets: &mut Option<TargetsCursor<'a>>,
+    ) -> Result<(), StateProofError> {
         trace!(
             target: TRACE_TARGET,
             branch = ?self.branch_stack.last(),
@@ -564,7 +570,7 @@ where
     /// creating a new one depending on the path of the key.
     fn push_leaf<'a>(
         &mut self,
-        targets: &mut TargetsCursor<'a>,
+        targets: &mut Option<TargetsCursor<'a>>,
         key: Nibbles,
         val: VE::DeferredEncoder,
     ) -> Result<(), StateProofError> {
@@ -652,7 +658,7 @@ where
     fn calculate_key_range<'a>(
         &mut self,
         value_encoder: &mut VE,
-        targets: &mut TargetsCursor<'a>,
+        targets: &mut Option<TargetsCursor<'a>>,
         hashed_cursor_current: &mut Option<(Nibbles, VE::DeferredEncoder)>,
         lower_bound: Nibbles,
         upper_bound: Option<Nibbles>,
@@ -720,7 +726,7 @@ where
     /// cached branch will be a child of that splitting branch.
     fn push_cached_branch<'a>(
         &mut self,
-        targets: &mut TargetsCursor<'a>,
+        targets: &mut Option<TargetsCursor<'a>>,
         cached_path: Nibbles,
         cached_branch: &BranchNodeCompact,
     ) -> Result<(), StateProofError> {
@@ -902,7 +908,7 @@ where
     #[instrument(target = TRACE_TARGET, level = "trace", skip_all)]
     fn next_uncached_key_range<'a>(
         &mut self,
-        targets: &mut TargetsCursor<'a>,
+        targets: &mut Option<TargetsCursor<'a>>,
         trie_cursor_state: &mut TrieCursorState,
         sub_trie_prefix: &Nibbles,
         sub_trie_upper_bound: Option<&Nibbles>,
@@ -1132,8 +1138,13 @@ where
     ) -> Result<(), StateProofError> {
         let sub_trie_upper_bound = sub_trie_targets.upper_bound();
 
-        // Wrap targets into a `TargetsCursor`.
-        let mut targets = TargetsCursor::new(sub_trie_targets.targets);
+        // Wrap targets into a `TargetsCursor`.  targets can be empty if we only want to calculate
+        // the root, in which case we don't need a cursor.
+        let mut targets = if sub_trie_targets.targets.is_empty() {
+            None
+        } else {
+            Some(TargetsCursor::new(sub_trie_targets.targets))
+        };
 
         // Ensure initial state is cleared. By the end of the method call these should be empty once
         // again.
@@ -1284,7 +1295,7 @@ where
             return Ok(Vec::new())
         }
 
-        // Initialize the variables which track the state of the two cursors. Both indicated the
+        // Initialize the variables which track the state of the two cursors. Both indicate the
         // cursors are unseeked.
         let mut trie_cursor_state = TrieCursorState::unseeked();
         let mut hashed_cursor_current: Option<(Nibbles, VE::DeferredEncoder)> = None;
@@ -1307,14 +1318,7 @@ where
         );
         Ok(core::mem::take(&mut self.retained_proofs))
     }
-}
 
-impl<TC, HC, VE> ProofCalculator<TC, HC, VE>
-where
-    TC: TrieCursor,
-    HC: HashedCursor,
-    VE: LeafValueEncoder<Value = HC::Value>,
-{
     /// Generate a proof for the given targets.
     ///
     /// Given a set of [`Target`]s, returns nodes whose paths are a prefix of any target. The
@@ -1333,6 +1337,75 @@ where
         self.hashed_cursor.reset();
         self.proof_inner(value_encoder, targets)
     }
+
+    /// Computes the root hash from a set of proof nodes.
+    ///
+    /// Returns `None` if there is no root node (partial proof), otherwise returns the hash of the
+    /// root node.
+    ///
+    /// This method reuses the internal RLP encode buffer for efficiency.
+    pub fn compute_root_hash(
+        &mut self,
+        proof_nodes: &[ProofTrieNode],
+    ) -> Result<Option<B256>, StateProofError> {
+        // Find the root node (node at empty path)
+        let root_node = proof_nodes.iter().find(|node| node.path.is_empty());
+
+        let Some(root) = root_node else {
+            return Ok(None);
+        };
+
+        // Compute the hash of the root node
+        self.rlp_encode_buf.clear();
+        root.node.encode(&mut self.rlp_encode_buf);
+        let root_hash = keccak256(&self.rlp_encode_buf);
+
+        Ok(Some(root_hash))
+    }
+
+    /// Calculates the root node of the trie.
+    ///
+    /// This method does not accept targets nor retain proofs. Returns the root node which can
+    /// be used to compute the root hash via [`Self::compute_root_hash`].
+    #[instrument(target = TRACE_TARGET, level = "trace", skip(self, value_encoder))]
+    pub fn root_node(&mut self, value_encoder: &mut VE) -> Result<ProofTrieNode, StateProofError> {
+        // Initialize the variables which track the state of the two cursors. Both indicate the
+        // cursors are unseeked.
+        let mut trie_cursor_state = TrieCursorState::unseeked();
+        let mut hashed_cursor_current: Option<(Nibbles, VE::DeferredEncoder)> = None;
+
+        static EMPTY_TARGETS: [Target; 0] = [];
+        let sub_trie_targets =
+            SubTrieTargets { prefix: Nibbles::new(), targets: &EMPTY_TARGETS, retain_root: true };
+
+        self.proof_subtrie(
+            value_encoder,
+            &mut trie_cursor_state,
+            &mut hashed_cursor_current,
+            sub_trie_targets,
+        )?;
+
+        // proof_subtrie will retain the root node if retain_proof is true, regardless of if there
+        // are any targets.
+        let mut proofs = core::mem::take(&mut self.retained_proofs);
+        trace!(
+            target: TRACE_TARGET,
+            proofs_len = ?proofs.len(),
+            "root_node: extracting root",
+        );
+
+        // The root node is at the empty path - it must exist since retain_root is true. Otherwise
+        // targets was empty, so there should be no other retained proofs.
+        debug_assert_eq!(
+            proofs.len(), 1,
+            "prefix is empty, retain_root is true, and targets is empty, so there must be only the root node"
+        );
+
+        // Find and remove the root node (node at empty path)
+        let root_node = proofs.pop().expect("prefix is empty, retain_root is true, and targets is empty, so there must be only the root node");
+
+        Ok(root_node)
+    }
 }
 
 /// A proof calculator for storage tries.
@@ -1383,29 +1456,32 @@ where
         self.proof_inner(&mut storage_value_encoder, targets)
     }
 
-    /// Computes the root hash from a set of proof nodes.
+    /// Calculates the root node of a storage trie.
     ///
-    /// Returns `None` if there is no root node (partial proof), otherwise returns the hash of the
-    /// root node.
-    ///
-    /// This method reuses the internal RLP encode buffer for efficiency.
-    pub fn compute_root_hash(
+    /// This method does not accept targets nor retain proofs. Returns the root node which can
+    /// be used to compute the root hash via [`Self::compute_root_hash`].
+    #[instrument(target = TRACE_TARGET, level = "trace", skip(self))]
+    pub fn storage_root_node(
         &mut self,
-        proof_nodes: &[ProofTrieNode],
-    ) -> Result<Option<B256>, StateProofError> {
-        // Find the root node (node at empty path)
-        let root_node = proof_nodes.iter().find(|node| node.path.is_empty());
+        hashed_address: B256,
+    ) -> Result<ProofTrieNode, StateProofError> {
+        self.hashed_cursor.set_hashed_address(hashed_address);
 
-        let Some(root) = root_node else {
-            return Ok(None);
-        };
+        if self.hashed_cursor.is_storage_empty()? {
+            return Ok(ProofTrieNode {
+                path: Nibbles::default(),
+                node: TrieNode::EmptyRoot,
+                masks: None,
+            })
+        }
 
-        // Compute the hash of the root node
-        self.rlp_encode_buf.clear();
-        root.node.encode(&mut self.rlp_encode_buf);
-        let root_hash = keccak256(&self.rlp_encode_buf);
+        // Don't call `set_hashed_address` on the trie cursor until after the previous shortcut has
+        // been checked.
+        self.trie_cursor.set_hashed_address(hashed_address);
 
-        Ok(Some(root_hash))
+        // Create a mutable storage value encoder
+        let mut storage_value_encoder = StorageValueEncoder;
+        self.root_node(&mut storage_value_encoder)
     }
 }
 
@@ -1585,6 +1661,8 @@ mod tests {
         trie_cursor_factory: MockTrieCursorFactory,
         /// Mock factory for hashed cursors, populated from `HashedPostState`
         hashed_cursor_factory: MockHashedCursorFactory,
+        /// The expected state root, calculated by `StateRoot`
+        expected_root: B256,
     }
 
     impl ProofTestHarness {
@@ -1612,7 +1690,7 @@ mod tests {
             let hashed_cursor_factory = MockHashedCursorFactory::from_hashed_post_state(post_state);
 
             // Generate TrieUpdates using StateRoot
-            let (_root, mut trie_updates) =
+            let (expected_root, mut trie_updates) =
                 crate::StateRoot::new(empty_trie_cursor_factory, hashed_cursor_factory.clone())
                     .root_with_updates()
                     .expect("StateRoot should succeed");
@@ -1624,7 +1702,7 @@ mod tests {
             // Initialize trie cursor factory from the generated TrieUpdates
             let trie_cursor_factory = MockTrieCursorFactory::from_trie_updates(trie_updates);
 
-            Self { trie_cursor_factory, hashed_cursor_factory }
+            Self { trie_cursor_factory, hashed_cursor_factory, expected_root }
         }
 
         /// Asserts that `ProofCalculator` and legacy `Proof` produce equivalent results for account
@@ -1713,6 +1791,26 @@ mod tests {
             // Basic comparison: both should succeed and produce identical results
             pretty_assertions::assert_eq!(proof_legacy_nodes, proof_v2_result);
 
+            // Also test root_node - get a fresh calculator and verify it returns the root node
+            // that hashes to the expected root
+            let trie_cursor = self.trie_cursor_factory.account_trie_cursor()?;
+            let hashed_cursor = self.hashed_cursor_factory.hashed_account_cursor()?;
+            let mut value_encoder = SyncAccountValueEncoder::new(
+                self.trie_cursor_factory.clone(),
+                self.hashed_cursor_factory.clone(),
+            );
+            let mut proof_calculator = ProofCalculator::new(trie_cursor, hashed_cursor);
+            let root_node = proof_calculator.root_node(&mut value_encoder)?;
+
+            // The root node should be at the empty path
+            assert!(root_node.path.is_empty(), "root_node should return node at empty path");
+
+            // The hash of the root node should match the expected root from legacy StateRoot
+            let root_hash = proof_calculator
+                .compute_root_hash(&[root_node])?
+                .expect("root_node returns a node at empty path");
+            pretty_assertions::assert_eq!(self.expected_root, root_hash);
+
             Ok(())
         }
     }
@@ -1802,7 +1900,6 @@ mod tests {
                 reth_tracing::init_test_tracing();
                 let harness = ProofTestHarness::new(post_state);
 
-                // Pass generated targets to both implementations
                 harness.assert_proof(targets).expect("Proof generation failed");
             }
         }

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

@@ -115,12 +115,10 @@ where
             self.hashed_cursor_factory.hashed_storage_cursor(self.hashed_address)?;
 
         let mut storage_proof_calculator = ProofCalculator::new_storage(trie_cursor, hashed_cursor);
-
-        let proof = storage_proof_calculator
-            .storage_proof(self.hashed_address, &mut [B256::ZERO.into()])?;
+        let root_node = storage_proof_calculator.storage_root_node(self.hashed_address)?;
         let storage_root = storage_proof_calculator
-            .compute_root_hash(&proof)?
-            .expect("storage_proof with dummy target always returns root");
+            .compute_root_hash(&[root_node])?
+            .expect("storage_root_node returns a node at empty path");
 
         let trie_account = self.account.into_trie_account(storage_root);
         trie_account.encode(buf);