From 23cfaebd9b2a6daf42104ba60ddcc7343990ca13 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Fri, 27 Sep 2024 08:39:42 +0200 Subject: [PATCH] test: witness --- Cargo.lock | 4 +- crates/trie/db/tests/witness.rs | 200 ++++++++++++++++++++++++++++++++ crates/trie/trie/src/witness.rs | 23 ++-- 3 files changed, 218 insertions(+), 9 deletions(-) create mode 100644 crates/trie/db/tests/witness.rs diff --git a/Cargo.lock b/Cargo.lock index 103da9225540..ea8eacd43af6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4424,7 +4424,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4979f22fdb869068da03c9f7528f8297c6fd2606bc3a4affe42e6a823fdb8da4" dependencies = [ "cfg-if", - "windows-targets 0.52.6", + "windows-targets 0.48.5", ] [[package]] @@ -11293,7 +11293,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] diff --git a/crates/trie/db/tests/witness.rs b/crates/trie/db/tests/witness.rs new file mode 100644 index 000000000000..62afbf061718 --- /dev/null +++ b/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, storage_updates: [(bool, BTreeMap); 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::().unwrap().upsert(hashed_address, Default::default()).unwrap(); + + // Insert init state into database + let mut hashed_storage_cursor = provider.tx_ref().cursor_write::().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::() + .unwrap() + .upsert(hashed_address, Default::default()) + .unwrap(); + + // Insert init state into database + let mut hashed_storage_cursor = + provider.tx_ref().cursor_write::().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; + } +} diff --git a/crates/trie/trie/src/witness.rs b/crates/trie/trie/src/witness.rs index 61576aabe36e..df86bed5536d 100644 --- a/crates/trie/trie/src/witness.rs +++ b/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; } }