From d5906f3f89653059e37afff71a9913a9b43a5039 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Mon, 2 Sep 2024 03:34:51 -0700 Subject: [PATCH] test(trie): improve in-memory nodes fuzz tests (#10479) --- crates/trie/db/tests/fuzz_in_memory_nodes.rs | 48 +++++++++++++------- 1 file changed, 31 insertions(+), 17 deletions(-) diff --git a/crates/trie/db/tests/fuzz_in_memory_nodes.rs b/crates/trie/db/tests/fuzz_in_memory_nodes.rs index 32737c198..3637dc8ba 100644 --- a/crates/trie/db/tests/fuzz_in_memory_nodes.rs +++ b/crates/trie/db/tests/fuzz_in_memory_nodes.rs @@ -1,15 +1,17 @@ use proptest::prelude::*; -use reth_db::{cursor::DbCursorRW, tables, transaction::DbTxMut}; +use reth_db::{ + cursor::{DbCursorRO, DbCursorRW, DbDupCursorRW}, + tables, + transaction::DbTxMut, +}; use reth_primitives::{Account, StorageEntry, B256, U256}; use reth_provider::test_utils::create_test_provider_factory; use reth_trie::{ - prefix_set::{PrefixSetMut, TriePrefixSets}, test_utils::{state_root_prehashed, storage_root_prehashed}, trie_cursor::InMemoryTrieCursorFactory, updates::TrieUpdates, - StateRoot, StorageRoot, + HashedPostState, HashedStorage, StateRoot, StorageRoot, }; -use reth_trie_common::Nibbles; use reth_trie_db::{DatabaseStateRoot, DatabaseStorageRoot, DatabaseTrieCursorFactory}; use std::collections::BTreeMap; @@ -19,7 +21,7 @@ proptest! { })] #[test] - fn fuzz_in_memory_nodes(mut init_state: BTreeMap, state_updates: [BTreeMap; 10]) { + fn fuzz_in_memory_account_nodes(mut init_state: BTreeMap, state_updates: [BTreeMap>; 10]) { let factory = create_test_provider_factory(); let provider = factory.provider_rw().unwrap(); let mut hashed_account_cursor = provider.tx_ref().cursor_write::().unwrap(); @@ -35,17 +37,24 @@ proptest! { .unwrap(); let mut state = init_state; - for mut state_update in state_updates { + for state_update in state_updates { // Insert state updates into database - let mut changes = PrefixSetMut::default(); - for (hashed_address, balance) in state_update.clone() { - hashed_account_cursor.upsert(hashed_address, Account { balance, ..Default::default() }).unwrap(); - changes.insert(Nibbles::unpack(hashed_address)); + let mut hashed_state = HashedPostState::default(); + for (hashed_address, balance) in state_update { + if let Some(balance) = balance { + let account = Account { balance, ..Default::default() }; + hashed_account_cursor.upsert(hashed_address, account).unwrap(); + hashed_state.accounts.insert(hashed_address, Some(account)); + state.insert(hashed_address, balance); + } else { + hashed_state.accounts.insert(hashed_address, None); + state.remove(&hashed_address); + } } // Compute root with in-memory trie nodes overlay let (state_root, trie_updates) = StateRoot::from_tx(provider.tx_ref()) - .with_prefix_sets(TriePrefixSets { account_prefix_set: changes.freeze(), ..Default::default() }) + .with_prefix_sets(hashed_state.construct_prefix_sets().freeze()) .with_trie_cursor_factory(InMemoryTrieCursorFactory::new( DatabaseTrieCursorFactory::new(provider.tx_ref()), &trie_nodes.clone().into_sorted()) ) @@ -55,7 +64,6 @@ proptest! { trie_nodes.extend(trie_updates); // Verify the result - state.append(&mut state_update); let expected_root = state_root_prehashed( state.iter().map(|(&key, &balance)| (key, (Account { balance, ..Default::default() }, std::iter::empty()))) ); @@ -64,7 +72,7 @@ proptest! { } #[test] - fn fuzz_in_memory_storage_nodes(mut init_storage: BTreeMap, storage_updates: [BTreeMap; 10]) { + fn fuzz_in_memory_storage_nodes(mut 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(); @@ -83,14 +91,17 @@ proptest! { StorageRoot::from_tx_hashed(provider.tx_ref(), hashed_address).root_with_updates().unwrap(); let mut storage = init_storage; - for mut storage_update in storage_updates { + for (is_deleted, mut storage_update) in storage_updates { // Insert state updates into database - let mut changes = PrefixSetMut::default(); + if is_deleted && hashed_storage_cursor.seek_exact(hashed_address).unwrap().is_some() { + hashed_storage_cursor.delete_current_duplicates().unwrap(); + } + let mut hashed_storage = HashedStorage::new(is_deleted); for (hashed_slot, value) in storage_update.clone() { hashed_storage_cursor .upsert(hashed_address, StorageEntry { key: hashed_slot, value }) .unwrap(); - changes.insert(Nibbles::unpack(hashed_slot)); + hashed_storage.storage.insert(hashed_slot, value); } // Compute root with in-memory trie nodes overlay @@ -98,7 +109,7 @@ proptest! { trie_nodes.insert_storage_updates(hashed_address, storage_trie_nodes.clone()); let (storage_root, _, trie_updates) = StorageRoot::from_tx_hashed(provider.tx_ref(), hashed_address) - .with_prefix_set(changes.freeze()) + .with_prefix_set(hashed_storage.construct_prefix_set().freeze()) .with_trie_cursor_factory(InMemoryTrieCursorFactory::new( DatabaseTrieCursorFactory::new(provider.tx_ref()), &trie_nodes.into_sorted(), @@ -109,6 +120,9 @@ proptest! { storage_trie_nodes.extend(trie_updates); // 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);