test: witness

crates/trie/db/tests/witness.rs

@@ -0,0 +1,200 @@
+use alloy_primitives::{B256, U256};
+use proptest::prelude::*;
+use reth_db::{
+    cursor::{DbCursorRO, DbCursorRW, DbDupCursorRW},
+    tables,
+    transaction::DbTxMut,
+};
+use reth_primitives::StorageEntry;
+use reth_provider::{test_utils::create_test_provider_factory, StorageTrieWriter};
+use reth_trie::{
+    prefix_set::TriePrefixSetsMut, test_utils::storage_root_prehashed, witness::TrieWitness,
+    HashedPostState, HashedStorage, StorageRoot, EMPTY_ROOT_HASH,
+};
+use reth_trie_db::{DatabaseStorageRoot, DatabaseTrieWitness};
+use std::collections::{BTreeMap, HashMap};
+
+// #[test]
+// fn f() {
+//     let hashed_address = B256::random();
+//     let factory = create_test_provider_factory();
+//     let provider = factory.provider_rw().unwrap();
+
+//     let (_, _, mut storage_trie_nodes) =
+//         StorageRoot::from_tx_hashed(provider.tx_ref(),
+// hashed_address).root_with_updates().unwrap();
+
+//     provider.write_storage_trie_updates(storage_tries).unwrap();
+//     TriePrefixSetsMut { storage_prefix_sets: HashMap::from([(hashed_address,)]) }
+// }
+
+proptest! {
+    #![proptest_config(ProptestConfig {
+        cases: 128, ..ProptestConfig::default()
+    })]
+
+    // Test that the witness can always be computed.
+    #[test]
+    fn fuzz_in_execution_witness_storage(init_storage: BTreeMap<B256, U256>, storage_updates: [(bool, BTreeMap<B256, U256>); 10]) {
+        let hashed_address = B256::random();
+        let factory = create_test_provider_factory();
+        let provider = factory.provider_rw().unwrap();
+
+        // Insert account
+        provider.tx_ref().cursor_write::<tables::HashedAccounts>().unwrap().upsert(hashed_address, Default::default()).unwrap();
+
+        // Insert init state into database
+        let mut hashed_storage_cursor = provider.tx_ref().cursor_write::<tables::HashedStorages>().unwrap();
+        for (hashed_slot, value) in init_storage.clone() {
+            hashed_storage_cursor
+                .upsert(hashed_address, StorageEntry { key: hashed_slot, value })
+                .unwrap();
+        }
+
+        // Compute initial storage root and updates
+        let (mut previous_storage_root, _, trie_updates) =
+            StorageRoot::from_tx_hashed(provider.tx_ref(), hashed_address).root_with_updates().unwrap();
+
+        provider.write_storage_trie_updates(&HashMap::from([(hashed_address, trie_updates)])).unwrap();
+
+        let mut storage = init_storage;
+        for (is_deleted, mut storage_update) in storage_updates {
+            let hashed_storage = HashedStorage::from_iter(is_deleted, storage_update.clone());
+
+            let prefix_set = hashed_storage.construct_prefix_set();
+            let witness = TrieWitness::from_tx(provider.tx_ref())
+                .with_prefix_sets_mut(TriePrefixSetsMut {
+                    storage_prefix_sets: HashMap::from([(hashed_address, prefix_set.clone())]),
+                    ..Default::default()
+                })
+                .compute(HashedPostState {
+                    accounts: HashMap::from([(hashed_address, Some(Default::default()))]),
+                    storages: HashMap::from([(hashed_address, hashed_storage)]),
+                })
+                .unwrap();
+            assert!(!witness.is_empty());
+            if previous_storage_root != EMPTY_ROOT_HASH {
+                assert!(witness.contains_key(&previous_storage_root));
+            }
+
+            // Insert state updates into database
+            if is_deleted && hashed_storage_cursor.seek_exact(hashed_address).unwrap().is_some() {
+                hashed_storage_cursor.delete_current_duplicates().unwrap();
+            }
+            for (hashed_slot, value) in storage_update.clone() {
+                hashed_storage_cursor
+                    .upsert(hashed_address, StorageEntry { key: hashed_slot, value })
+                    .unwrap();
+            }
+
+            // Compute root with in-memory trie nodes overlay
+            let (storage_root, _, trie_updates) =
+                StorageRoot::from_tx_hashed(provider.tx_ref(), hashed_address)
+                    .with_prefix_set(prefix_set.freeze())
+                    .root_with_updates()
+                    .unwrap();
+            provider.write_storage_trie_updates(&HashMap::from([(hashed_address, trie_updates)])).unwrap();
+
+            // Verify the result
+            if is_deleted {
+                storage.clear();
+            }
+            storage.append(&mut storage_update);
+            let expected_root = storage_root_prehashed(storage.clone());
+            assert_eq!(expected_root, storage_root);
+
+            previous_storage_root = storage_root;
+        }
+    }
+}
+
+#[test]
+fn execution_witness_storage() {
+    let init_storage = BTreeMap::default();
+    let storage_updates = [
+        (false, BTreeMap::default()),
+        // (false, BTreeMap::from([(B256::with_last_byte(69), U256::ZERO)])),
+        // (false, BTreeMap::from([(B256::with_last_byte(42), U256::ZERO)])),
+        (false, BTreeMap::from([((B256::with_last_byte(69), U256::from(1)))])),
+        (false, BTreeMap::from([(B256::with_last_byte(42), U256::from(2))])),
+    ];
+
+    let hashed_address = B256::random();
+    let factory = create_test_provider_factory();
+    let provider = factory.provider_rw().unwrap();
+
+    // Insert account
+    provider
+        .tx_ref()
+        .cursor_write::<tables::HashedAccounts>()
+        .unwrap()
+        .upsert(hashed_address, Default::default())
+        .unwrap();
+
+    // Insert init state into database
+    let mut hashed_storage_cursor =
+        provider.tx_ref().cursor_write::<tables::HashedStorages>().unwrap();
+    for (hashed_slot, value) in init_storage.clone() {
+        hashed_storage_cursor
+            .upsert(hashed_address, StorageEntry { key: hashed_slot, value })
+            .unwrap();
+    }
+
+    // Compute initial storage root and updates
+    let (mut previous_storage_root, _, trie_updates) =
+        StorageRoot::from_tx_hashed(provider.tx_ref(), hashed_address).root_with_updates().unwrap();
+
+    provider.write_storage_trie_updates(&HashMap::from([(hashed_address, trie_updates)])).unwrap();
+
+    let mut storage = init_storage;
+    for (is_deleted, mut storage_update) in storage_updates {
+        // Insert state updates into database
+        let hashed_storage = HashedStorage::from_iter(is_deleted, storage_update.clone());
+
+        let prefix_set = hashed_storage.construct_prefix_set();
+        let witness = TrieWitness::from_tx(provider.tx_ref())
+            .with_prefix_sets_mut(TriePrefixSetsMut {
+                storage_prefix_sets: HashMap::from([(hashed_address, prefix_set.clone())]),
+                ..Default::default()
+            })
+            .compute(HashedPostState {
+                accounts: HashMap::from([(hashed_address, Some(Default::default()))]),
+                storages: HashMap::from([(hashed_address, hashed_storage)]),
+            })
+            .unwrap();
+        println!("prev storage root {previous_storage_root} witness {:?}", witness);
+        assert!(!witness.is_empty());
+        if previous_storage_root != EMPTY_ROOT_HASH {
+            assert!(witness.contains_key(&previous_storage_root));
+        }
+
+        if is_deleted && hashed_storage_cursor.seek_exact(hashed_address).unwrap().is_some() {
+            hashed_storage_cursor.delete_current_duplicates().unwrap();
+        }
+        for (hashed_slot, value) in storage_update.clone() {
+            hashed_storage_cursor
+                .upsert(hashed_address, StorageEntry { key: hashed_slot, value })
+                .unwrap();
+        }
+
+        // Compute root with in-memory trie nodes overlay
+        let (storage_root, _, trie_updates) =
+            StorageRoot::from_tx_hashed(provider.tx_ref(), hashed_address)
+                .with_prefix_set(prefix_set.freeze())
+                .root_with_updates()
+                .unwrap();
+        provider
+            .write_storage_trie_updates(&HashMap::from([(hashed_address, trie_updates)]))
+            .unwrap();
+
+        // Verify the result
+        if is_deleted {
+            storage.clear();
+        }
+        storage.append(&mut storage_update);
+        let expected_root = storage_root_prehashed(storage.clone());
+        assert_eq!(expected_root, storage_root);
+
+        previous_storage_root = storage_root;
+    }
+}

crates/trie/trie/src/witness.rs

@@ -211,9 +211,14 @@ where
                 TrieNode::Branch(branch) => {
                     next_path.push(key[path.len()]);
                     let children = branch_node_children(path.clone(), &branch);
-                    for (child_path, node_hash) in children {
+                    for (child_path, value) in children {
                         if !key.starts_with(&child_path) {
-                            trie_nodes.insert(child_path, Either::Left(node_hash));
+                            let value = if value.len() < B256::len_bytes() {
+                                Either::Right(value.to_vec())
+                            } else {
+                                Either::Left(B256::from_slice(&value[1..]))
+                            };
+                            trie_nodes.insert(child_path, value);
                         }
                     }
                 }
@@ -269,9 +274,13 @@ where
                             let node = trie_node_provider(path.clone())?;
                             match TrieNode::decode(&mut &node[..])? {
                                 TrieNode::Branch(branch) => {
-                                    let children = branch_node_children(path, &branch);
-                                    for (child_path, branch_hash) in children {
-                                        hash_builder.add_branch(child_path, branch_hash, false);
+                                    for (child_path, value) in branch_node_children(path, &branch) {
+                                        if value.len() < B256::len_bytes() {
+                                            hash_builder.add_leaf(child_path, value);
+                                        } else {
+                                            let hash = B256::from_slice(&value[1..]);
+                                            hash_builder.add_branch(child_path, hash, false);
+                                        }
                                     }
                                     break
                                 }
@@ -301,14 +310,14 @@ where
 }
 
 /// Returned branch node children with keys in order.
-fn branch_node_children(prefix: Nibbles, node: &BranchNode) -> Vec<(Nibbles, B256)> {
+fn branch_node_children(prefix: Nibbles, node: &BranchNode) -> Vec<(Nibbles, &[u8])> {
     let mut children = Vec::with_capacity(node.state_mask.count_ones() as usize);
     let mut stack_ptr = node.as_ref().first_child_index();
     for index in CHILD_INDEX_RANGE {
         if node.state_mask.is_bit_set(index) {
             let mut child_path = prefix.clone();
             child_path.push(index);
-            children.push((child_path, B256::from_slice(&node.stack[stack_ptr][1..])));
+            children.push((child_path, &node.stack[stack_ptr][..]));
             stack_ptr += 1;
         }
     }