feat(witness): always_include_root_node flag (#15679)

crates/trie/trie/src/witness.rs

@@ -4,6 +4,8 @@ use crate::{
     proof::{Proof, ProofBlindedProviderFactory},
     trie_cursor::TrieCursorFactory,
 };
+use alloy_rlp::EMPTY_STRING_CODE;
+use alloy_trie::EMPTY_ROOT_HASH;
 use reth_trie_common::HashedPostState;
 
 use alloy_primitives::{
@@ -32,6 +34,11 @@ pub struct TrieWitness<T, H> {
     hashed_cursor_factory: H,
     /// A set of prefix sets that have changes.
     prefix_sets: TriePrefixSetsMut,
+    /// Flag indicating whether the root node should always be included (even if the target state
+    /// is empty). This setting is useful if the caller wants to verify the witness against the
+    /// parent state root.
+    /// Set to `false` by default.
+    always_include_root_node: bool,
     /// Recorded witness.
     witness: B256Map<Bytes>,
 }
@@ -43,6 +50,7 @@ impl<T, H> TrieWitness<T, H> {
             trie_cursor_factory,
             hashed_cursor_factory,
             prefix_sets: TriePrefixSetsMut::default(),
+            always_include_root_node: false,
             witness: HashMap::default(),
         }
     }
@@ -53,6 +61,7 @@ impl<T, H> TrieWitness<T, H> {
             trie_cursor_factory,
             hashed_cursor_factory: self.hashed_cursor_factory,
             prefix_sets: self.prefix_sets,
+            always_include_root_node: self.always_include_root_node,
             witness: self.witness,
         }
     }
@@ -63,6 +72,7 @@ impl<T, H> TrieWitness<T, H> {
             trie_cursor_factory: self.trie_cursor_factory,
             hashed_cursor_factory,
             prefix_sets: self.prefix_sets,
+            always_include_root_node: self.always_include_root_node,
             witness: self.witness,
         }
     }
@@ -72,6 +82,14 @@ impl<T, H> TrieWitness<T, H> {
         self.prefix_sets = prefix_sets;
         self
     }
+
+    /// Set `always_include_root_node` to true. Root node will be included even on empty state.
+    /// This setting is useful if the caller wants to verify the witness against the
+    /// parent state root.
+    pub const fn always_include_root_node(mut self) -> Self {
+        self.always_include_root_node = true;
+        self
+    }
 }
 
 impl<T, H> TrieWitness<T, H>
@@ -86,16 +104,34 @@ where
     ///
     /// `state` - state transition containing both modified and touched accounts and storage slots.
     pub fn compute(mut self, state: HashedPostState) -> Result<B256Map<Bytes>, TrieWitnessError> {
-        if state.is_empty() {
-            return Ok(self.witness)
+        let is_state_empty = state.is_empty();
+        if is_state_empty && !self.always_include_root_node {
+            return Ok(Default::default())
         }
 
-        let proof_targets = self.get_proof_targets(&state)?;
+        let proof_targets = if is_state_empty {
+            MultiProofTargets::account(B256::ZERO)
+        } else {
+            self.get_proof_targets(&state)?
+        };
         let multiproof =
             Proof::new(self.trie_cursor_factory.clone(), self.hashed_cursor_factory.clone())
                 .with_prefix_sets_mut(self.prefix_sets.clone())
                 .multiproof(proof_targets.clone())?;
 
+        // No need to reconstruct the rest of the trie, we just need to include
+        // the root node and return.
+        if is_state_empty {
+            let (root_hash, root_node) = if let Some(root_node) =
+                multiproof.account_subtree.into_inner().remove(&Nibbles::default())
+            {
+                (keccak256(&root_node), root_node)
+            } else {
+                (EMPTY_ROOT_HASH, Bytes::from([EMPTY_STRING_CODE]))
+            };
+            return Ok(B256Map::from_iter([(root_hash, root_node)]))
+        }
+
         // Record all nodes from multiproof in the witness
         for account_node in multiproof.account_subtree.values() {
             if let Entry::Vacant(entry) = self.witness.entry(keccak256(account_node.as_ref())) {