feat(trie): in-memory trie node overlay (#8199)

Co-authored-by: Roman Krasiuk <rokrassyuk@gmail.com>

crates/trie/trie/src/trie.rs

@@ -546,9 +546,11 @@ where
 mod tests {
     use super::*;
     use crate::{
+        hashed_cursor::HashedPostStateCursorFactory,
         prefix_set::PrefixSetMut,
         test_utils::{state_root, state_root_prehashed, storage_root, storage_root_prehashed},
-        BranchNodeCompact, TrieMask,
+        trie_cursor::TrieUpdatesCursorFactory,
+        BranchNodeCompact, HashedPostState, HashedStorage, TrieMask,
     };
     use proptest::{prelude::ProptestConfig, proptest};
     use proptest_arbitrary_interop::arb;
@@ -562,6 +564,7 @@ mod tests {
     use reth_trie_common::triehash::KeccakHasher;
     use std::{
         collections::{BTreeMap, HashMap},
+        iter,
         ops::Mul,
         str::FromStr,
         sync::Arc,
@@ -1369,4 +1372,126 @@ mod tests {
         assert_eq!(node.root_hash, None);
         assert_eq!(node.hashes.len(), 1);
     }
+
+    #[test]
+    fn trie_updates_across_multiple_iterations() {
+        let address = Address::ZERO;
+        let hashed_address = keccak256(address);
+
+        let factory = create_test_provider_factory();
+
+        let mut hashed_storage = BTreeMap::default();
+        let mut post_state = HashedPostState::default();
+
+        // Block #1
+        // Update specific storage slots
+        let mut modified_storage = BTreeMap::default();
+
+        // 0x0f..
+        let modified_key_prefix = Nibbles::from_nibbles(
+            [0x0, 0xf].into_iter().chain(iter::repeat(0).take(62)).collect::<Vec<_>>(),
+        );
+
+        // 0x0faa0..
+        let mut modified_entry1 = modified_key_prefix.clone();
+        modified_entry1.set_at(2, 0xa);
+        modified_entry1.set_at(3, 0xa);
+
+        // 0x0faaa..
+        let mut modified_entry2 = modified_key_prefix.clone();
+        modified_entry2.set_at(2, 0xa);
+        modified_entry2.set_at(3, 0xa);
+        modified_entry2.set_at(4, 0xa);
+
+        // 0x0fab0..
+        let mut modified_entry3 = modified_key_prefix.clone();
+        modified_entry3.set_at(2, 0xa);
+        modified_entry3.set_at(3, 0xb);
+
+        // 0x0fba0..
+        let mut modified_entry4 = modified_key_prefix;
+        modified_entry4.set_at(2, 0xb);
+        modified_entry4.set_at(3, 0xa);
+
+        [modified_entry1, modified_entry2, modified_entry3.clone(), modified_entry4]
+            .into_iter()
+            .for_each(|key| {
+                modified_storage.insert(B256::from_slice(&key.pack()), U256::from(1));
+            });
+
+        // Update main hashed storage.
+        hashed_storage.extend(modified_storage.clone());
+        post_state.extend(HashedPostState::default().with_storages([(
+            hashed_address,
+            HashedStorage::from_iter(false, modified_storage.clone()),
+        )]));
+
+        let (storage_root, block1_updates) = compute_storage_root(
+            address,
+            factory.provider().unwrap().tx_ref(),
+            &post_state,
+            &TrieUpdates::default(),
+        );
+        assert_eq!(storage_root, storage_root_prehashed(hashed_storage.clone()));
+
+        // Block #2
+        // Set 0x0fab0.. hashed slot to 0
+        modified_storage.insert(B256::from_slice(&modified_entry3.pack()), U256::ZERO);
+
+        // Update main hashed storage.
+        hashed_storage.remove(&B256::from_slice(&modified_entry3.pack()));
+        post_state.extend(HashedPostState::default().with_storages([(
+            hashed_address,
+            HashedStorage::from_iter(false, modified_storage.clone()),
+        )]));
+
+        let (storage_root, block2_updates) = compute_storage_root(
+            address,
+            factory.provider().unwrap().tx_ref(),
+            &post_state,
+            &block1_updates,
+        );
+        assert_eq!(storage_root, storage_root_prehashed(hashed_storage.clone()));
+
+        // Commit trie updates
+        {
+            let mut updates = block1_updates;
+            updates.extend(block2_updates);
+
+            let provider_rw = factory.provider_rw().unwrap();
+            let mut hashed_storage_cursor =
+                provider_rw.tx_ref().cursor_dup_write::<tables::HashedStorages>().unwrap();
+            for (hashed_slot, value) in &hashed_storage {
+                hashed_storage_cursor
+                    .upsert(hashed_address, StorageEntry { key: *hashed_slot, value: *value })
+                    .unwrap();
+            }
+            updates.flush(provider_rw.tx_ref()).unwrap();
+            provider_rw.commit().unwrap();
+        }
+
+        // Recompute storage root for block #3
+        let storage_root =
+            StorageRoot::from_tx(factory.provider().unwrap().tx_ref(), address).root().unwrap();
+        assert_eq!(storage_root, storage_root_prehashed(hashed_storage.clone()));
+    }
+
+    fn compute_storage_root<TX: DbTx>(
+        address: Address,
+        tx: &TX,
+        post_state: &HashedPostState,
+        update: &TrieUpdates,
+    ) -> (B256, TrieUpdates) {
+        let mut prefix_sets = post_state.construct_prefix_sets();
+        let (root, _, updates) = StorageRoot::from_tx(tx, address)
+            .with_hashed_cursor_factory(HashedPostStateCursorFactory::new(
+                tx,
+                &post_state.clone().into_sorted(),
+            ))
+            .with_trie_cursor_factory(TrieUpdatesCursorFactory::new(tx, &update.sorted()))
+            .with_prefix_set(prefix_sets.storage_prefix_sets.remove(&keccak256(address)).unwrap())
+            .root_with_updates()
+            .unwrap();
+        (root, updates)
+    }
 }

crates/trie/trie/src/trie_cursor/database_cursors.rs

@@ -9,18 +9,22 @@ use reth_primitives::B256;
 
 /// Implementation of the trie cursor factory for a database transaction.
 impl<'a, TX: DbTx> TrieCursorFactory for &'a TX {
-    fn account_trie_cursor(&self) -> Result<Box<dyn TrieCursor + '_>, DatabaseError> {
-        Ok(Box::new(DatabaseAccountTrieCursor::new(self.cursor_read::<tables::AccountsTrie>()?)))
+    type AccountTrieCursor = DatabaseAccountTrieCursor<<TX as DbTx>::Cursor<tables::AccountsTrie>>;
+    type StorageTrieCursor =
+        DatabaseStorageTrieCursor<<TX as DbTx>::DupCursor<tables::StoragesTrie>>;
+
+    fn account_trie_cursor(&self) -> Result<Self::AccountTrieCursor, DatabaseError> {
+        Ok(DatabaseAccountTrieCursor::new(self.cursor_read::<tables::AccountsTrie>()?))
     }
 
     fn storage_trie_cursor(
         &self,
         hashed_address: B256,
-    ) -> Result<Box<dyn TrieCursor + '_>, DatabaseError> {
-        Ok(Box::new(DatabaseStorageTrieCursor::new(
+    ) -> Result<Self::StorageTrieCursor, DatabaseError> {
+        Ok(DatabaseStorageTrieCursor::new(
             self.cursor_dup_read::<tables::StoragesTrie>()?,
             hashed_address,
-        )))
+        ))
     }
 }
 

crates/trie/trie/src/trie_cursor/mod.rs

@@ -3,25 +3,32 @@ use reth_db::DatabaseError;
 use reth_primitives::B256;
 mod database_cursors;
 mod subnode;
+mod update;
 
 /// Noop trie cursor implementations.
 pub mod noop;
 
 pub use self::{
     database_cursors::{DatabaseAccountTrieCursor, DatabaseStorageTrieCursor},
     subnode::CursorSubNode,
+    update::*,
 };
 
 /// Factory for creating trie cursors.
 pub trait TrieCursorFactory {
+    /// The account trie cursor type.
+    type AccountTrieCursor: TrieCursor;
+    /// The storage trie cursor type.
+    type StorageTrieCursor: TrieCursor;
+
     /// Create an account trie cursor.
-    fn account_trie_cursor(&self) -> Result<Box<dyn TrieCursor + '_>, DatabaseError>;
+    fn account_trie_cursor(&self) -> Result<Self::AccountTrieCursor, DatabaseError>;
 
     /// Create a storage tries cursor.
     fn storage_trie_cursor(
         &self,
         hashed_address: B256,
-    ) -> Result<Box<dyn TrieCursor + '_>, DatabaseError>;
+    ) -> Result<Self::StorageTrieCursor, DatabaseError>;
 }
 
 /// A cursor for navigating a trie that works with both Tables and DupSort tables.

crates/trie/trie/src/trie_cursor/noop.rs

@@ -9,17 +9,20 @@ use reth_primitives::B256;
 pub struct NoopTrieCursorFactory;
 
 impl TrieCursorFactory for NoopTrieCursorFactory {
+    type AccountTrieCursor = NoopAccountTrieCursor;
+    type StorageTrieCursor = NoopStorageTrieCursor;
+
     /// Generates a Noop account trie cursor.
-    fn account_trie_cursor(&self) -> Result<Box<dyn TrieCursor + '_>, DatabaseError> {
-        Ok(Box::<NoopAccountTrieCursor>::default())
+    fn account_trie_cursor(&self) -> Result<Self::AccountTrieCursor, DatabaseError> {
+        Ok(NoopAccountTrieCursor::default())
     }
 
     /// Generates a Noop storage trie cursor.
     fn storage_trie_cursor(
         &self,
         _hashed_address: B256,
-    ) -> Result<Box<dyn TrieCursor + '_>, DatabaseError> {
-        Ok(Box::<NoopStorageTrieCursor>::default())
+    ) -> Result<Self::StorageTrieCursor, DatabaseError> {
+        Ok(NoopStorageTrieCursor::default())
     }
 }