diff --git a/crates/trie/trie/src/witness.rs b/crates/trie/trie/src/witness.rs index 57b91f8264..ce40a01e1c 100644 --- a/crates/trie/trie/src/witness.rs +++ b/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 { 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, } @@ -43,6 +50,7 @@ impl TrieWitness { trie_cursor_factory, hashed_cursor_factory, prefix_sets: TriePrefixSetsMut::default(), + always_include_root_node: false, witness: HashMap::default(), } } @@ -53,6 +61,7 @@ impl TrieWitness { 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 TrieWitness { 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 TrieWitness { 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 TrieWitness @@ -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, 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())) {