test(trie): fuzz in-memory storage nodes

crates/trie/db/tests/fuzz_in_memory_nodes.rs

@@ -1,15 +1,16 @@
 use proptest::prelude::*;
 use reth_db::{cursor::DbCursorRW, tables, transaction::DbTxMut};
-use reth_primitives::{Account, B256, U256};
+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,
+    test_utils::{state_root_prehashed, storage_root_prehashed},
     trie_cursor::InMemoryTrieCursorFactory,
-    StateRoot,
+    updates::TrieUpdates,
+    StateRoot, StorageRoot,
 };
 use reth_trie_common::Nibbles;
-use reth_trie_db::{DatabaseStateRoot, DatabaseTrieCursorFactory};
+use reth_trie_db::{DatabaseStateRoot, DatabaseStorageRoot, DatabaseTrieCursorFactory};
 use std::collections::BTreeMap;
 
 proptest! {
@@ -60,7 +61,57 @@ proptest! {
             );
             assert_eq!(expected_root, state_root);
         }
+    }
+
+    #[test]
+    fn fuzz_in_memory_storage_nodes(mut init_storage: BTreeMap<B256, U256>, storage_updates: [BTreeMap<B256, U256>; 10]) {
+        let hashed_address = B256::random();
+        let factory = create_test_provider_factory();
+        let provider = factory.provider_rw().unwrap();
+        let mut hashed_storage_cursor =
+            provider.tx_ref().cursor_write::<tables::HashedStorages>().unwrap();
+
+        // Insert init state into database
+        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 storage_trie_nodes) =
+            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 {
+            // Insert state updates into database
+            let mut changes = PrefixSetMut::default();
+            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));
+            }
+
+            // 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_trie_cursor_factory(InMemoryTrieCursorFactory::new(
+                        DatabaseTrieCursorFactory::new(provider.tx_ref()),
+                        &trie_nodes.into_sorted(),
+                    ))
+                    .root_with_updates()
+                    .unwrap();
+
+            storage_trie_nodes.extend(trie_updates);
+
+            // Verify the result
+            storage.append(&mut storage_update);
+            let expected_root = storage_root_prehashed(storage.clone());
+            assert_eq!(expected_root, storage_root);
+        }
     }
 }