test(trie): improve in-memory nodes fuzz tests (paradigmxyz#10479)

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<B256, U256>, state_updates: [BTreeMap<B256, U256>; 10]) {
+    fn fuzz_in_memory_account_nodes(mut init_state: BTreeMap<B256, U256>, state_updates: [BTreeMap<B256, Option<U256>>; 10]) {
         let factory = create_test_provider_factory();
         let provider = factory.provider_rw().unwrap();
         let mut hashed_account_cursor = provider.tx_ref().cursor_write::<tables::HashedAccounts>().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<B256, U256>, storage_updates: [BTreeMap<B256, U256>; 10]) {
+    fn fuzz_in_memory_storage_nodes(mut 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();
@@ -83,22 +91,25 @@ 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
             let mut trie_nodes = TrieUpdates::default();
             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);