test(trie): witness empty root node

crates/trie/common/src/proofs.rs

@@ -2,7 +2,7 @@
 
 use crate::{Nibbles, TrieAccount};
 use alloy_primitives::{keccak256, Address, Bytes, B256, U256};
-use alloy_rlp::{encode_fixed_size, Decodable};
+use alloy_rlp::{encode_fixed_size, Decodable, EMPTY_STRING_CODE};
 use alloy_trie::{
     nodes::TrieNode,
     proof::{verify_proof, ProofNodes, ProofVerificationError},
@@ -86,13 +86,18 @@ pub struct StorageMultiProof {
     pub subtree: ProofNodes,
 }
 
-impl Default for StorageMultiProof {
-    fn default() -> Self {
-        Self { root: EMPTY_ROOT_HASH, subtree: Default::default() }
+impl StorageMultiProof {
+    /// Create new storage multiproof for empty trie.
+    pub fn empty() -> Self {
+        Self {
+            root: EMPTY_ROOT_HASH,
+            subtree: ProofNodes::from_iter([(
+                Nibbles::default(),
+                Bytes::from([EMPTY_STRING_CODE]),
+            )]),
+        }
     }
-}
 
-impl StorageMultiProof {
     /// Return storage proofs for the target storage slot (unhashed).
     pub fn storage_proof(&self, slot: B256) -> Result<StorageProof, alloy_rlp::Error> {
         let nibbles = Nibbles::unpack(keccak256(slot));

crates/trie/db/tests/witness.rs

@@ -0,0 +1,56 @@
+use alloy_primitives::{
+    keccak256,
+    map::{HashMap, HashSet},
+    Address,
+};
+use alloy_rlp::EMPTY_STRING_CODE;
+use reth_primitives::{constants::EMPTY_ROOT_HASH, Account, Bytes, B256, U256};
+use reth_provider::{test_utils::create_test_provider_factory, HashingWriter};
+use reth_trie::{proof::Proof, witness::TrieWitness, HashedPostState, HashedStorage, StateRoot};
+use reth_trie_db::{DatabaseProof, DatabaseStateRoot, DatabaseTrieWitness};
+
+#[test]
+fn includes_empty_node_preimage() {
+    let factory = create_test_provider_factory();
+    let provider = factory.provider_rw().unwrap();
+
+    let address = Address::random();
+    let hashed_address = keccak256(address);
+    let hashed_slot = B256::random();
+
+    // witness includes empty state trie root node
+    assert_eq!(
+        TrieWitness::from_tx(provider.tx_ref())
+            .compute(HashedPostState {
+                accounts: HashMap::from([(hashed_address, Some(Account::default()))]),
+                storages: HashMap::default(),
+            })
+            .unwrap(),
+        HashMap::from([(EMPTY_ROOT_HASH, Bytes::from([EMPTY_STRING_CODE]))])
+    );
+
+    // Insert account into database
+    provider.insert_account_for_hashing([(address, Some(Account::default()))]).unwrap();
+
+    let state_root = StateRoot::from_tx(provider.tx_ref()).root().unwrap();
+    let multiproof = Proof::from_tx(provider.tx_ref())
+        .with_target((hashed_address, HashSet::from([hashed_slot])))
+        .multiproof()
+        .unwrap();
+
+    let witness = TrieWitness::from_tx(provider.tx_ref())
+        .compute(HashedPostState {
+            accounts: HashMap::from([(hashed_address, Some(Account::default()))]),
+            storages: HashMap::from([(
+                hashed_address,
+                HashedStorage::from_iter(false, [(hashed_slot, U256::from(1))]),
+            )]),
+        })
+        .unwrap();
+    assert!(witness.contains_key(&state_root));
+    for node in multiproof.account_subtree.values() {
+        assert_eq!(witness.get(&keccak256(&node)), Some(node));
+    }
+    // witness includes empty state trie root node
+    assert_eq!(witness.get(&EMPTY_ROOT_HASH), Some(&Bytes::from([EMPTY_STRING_CODE])));
+}

crates/trie/trie/src/proof.rs

@@ -149,7 +149,7 @@ where
 
         // short circuit on empty storage
         if hashed_storage_cursor.is_storage_empty()? {
-            return Ok(StorageMultiProof::default())
+            return Ok(StorageMultiProof::empty())
         }
 
         let target_nibbles = self

crates/trie/trie/src/witness.rs

@@ -14,7 +14,7 @@ use itertools::{Either, Itertools};
 use reth_execution_errors::TrieWitnessError;
 use reth_primitives::constants::EMPTY_ROOT_HASH;
 use reth_trie_common::{
-    BranchNode, HashBuilder, Nibbles, TrieAccount, TrieNode, CHILD_INDEX_RANGE,
+    BranchNode, HashBuilder, Nibbles, StorageMultiProof, TrieAccount, TrieNode, CHILD_INDEX_RANGE,
 };
 
 /// State transition witness for the trie.
@@ -103,8 +103,10 @@ where
         let mut account_rlp = Vec::with_capacity(128);
         let mut account_trie_nodes = BTreeMap::default();
         for (hashed_address, hashed_slots) in proof_targets {
-            let storage_multiproof =
-                account_multiproof.storages.remove(&hashed_address).unwrap_or_default();
+            let storage_multiproof = account_multiproof
+                .storages
+                .remove(&hashed_address)
+                .unwrap_or_else(StorageMultiProof::empty);
 
             // Gather and record account trie nodes.
             let account = state
@@ -202,7 +204,8 @@ where
         proof: impl IntoIterator<Item = (&'b Nibbles, &'b Bytes)>,
     ) -> Result<BTreeMap<Nibbles, Either<B256, Vec<u8>>>, TrieWitnessError> {
         let mut trie_nodes = BTreeMap::default();
-        for (path, encoded) in proof {
+        let mut proof_iter = proof.into_iter().enumerate().peekable();
+        while let Some((idx, (path, encoded))) = proof_iter.next() {
             // Record the node in witness.
             self.witness.insert(keccak256(encoded.as_ref()), encoded.clone());
 
@@ -226,7 +229,11 @@ where
                         trie_nodes.insert(next_path.clone(), Either::Right(leaf.value.clone()));
                     }
                 }
-                TrieNode::EmptyRoot => return Err(TrieWitnessError::UnexpectedEmptyRoot(next_path)),
+                TrieNode::EmptyRoot => {
+                    if idx != 0 || proof_iter.peek().is_some() {
+                        return Err(TrieWitnessError::UnexpectedEmptyRoot(next_path))
+                    }
+                }
             };
         }