feat(trie): in-memory trie node overlay

crates/blockchain-tree/src/blockchain_tree.rs

@@ -1775,7 +1775,7 @@ mod tests {
         );
         let block2_chain_id = tree.state.block_indices.get_blocks_chain_id(&block2.hash()).unwrap();
         let block2_chain = tree.state.chains.get(&block2_chain_id).unwrap();
-        assert!(block2_chain.trie_updates().is_none());
+        assert!(block2_chain.trie_updates().is_some());
 
         assert_eq!(
             tree.make_canonical(block2.hash()).unwrap(),
@@ -1810,7 +1810,7 @@ mod tests {
 
         let block5_chain_id = tree.state.block_indices.get_blocks_chain_id(&block5.hash()).unwrap();
         let block5_chain = tree.state.chains.get(&block5_chain_id).unwrap();
-        assert!(block5_chain.trie_updates().is_none());
+        assert!(block5_chain.trie_updates().is_some());
 
         assert_eq!(
             tree.make_canonical(block5.hash()).unwrap(),

crates/blockchain-tree/src/chain.rs

@@ -24,7 +24,7 @@ use reth_provider::{
     BundleStateDataProvider, BundleStateWithReceipts, Chain, ProviderError, StateRootProvider,
 };
 use reth_revm::database::StateProviderDatabase;
-use reth_trie::updates::TrieUpdates;
+use reth_trie::updates::{TrieUpdates, TrieUpdatesSorted};
 use reth_trie_parallel::parallel_root::ParallelStateRoot;
 use std::{
     collections::BTreeMap,
@@ -98,6 +98,7 @@ impl AppendableChain {
             externals,
             block_attachment,
             block_validation_kind,
+            None,
         )?;
 
         Ok(Self { chain: Chain::new(vec![block], bundle_state, trie_updates) })
@@ -144,6 +145,7 @@ impl AppendableChain {
             externals,
             BlockAttachment::HistoricalFork,
             block_validation_kind,
+            None,
         )?;
         // extending will also optimize few things, mostly related to selfdestruct and wiping of
         // storage.
@@ -176,6 +178,7 @@ impl AppendableChain {
         externals: &TreeExternals<DB, E>,
         block_attachment: BlockAttachment,
         block_validation_kind: BlockValidationKind,
+        trie_updates: Option<TrieUpdatesSorted>,
     ) -> RethResult<(BundleStateWithReceipts, Option<TrieUpdates>)>
     where
         BSDP: BundleStateDataProvider,
@@ -227,7 +230,7 @@ impl AppendableChain {
                 let mut state = provider.bundle_state_data_provider.state().clone();
                 state.extend(bundle_state.clone());
                 let hashed_state = state.hash_state_slow();
-                ParallelStateRoot::new(consistent_view, hashed_state)
+                ParallelStateRoot::new(consistent_view, hashed_state, trie_updates)
                     .incremental_root_with_updates()
                     .map(|(root, updates)| (root, Some(updates)))
                     .map_err(ProviderError::from)?
@@ -291,16 +294,17 @@ impl AppendableChain {
             canonical_fork,
         };
 
-        let (block_state, _) = Self::validate_and_execute(
+        let (block_state, trie_updates) = Self::validate_and_execute(
             block.clone(),
             parent_block,
             bundle_state_data,
             externals,
             block_attachment,
             block_validation_kind,
+            self.chain.trie_updates().map(|u| u.sorted()),
         )?;
         // extend the state.
-        self.chain.append_block(block, block_state);
+        self.chain.append_block(block, block_state, trie_updates);
 
         Ok(())
     }

crates/storage/provider/src/chain.rs

@@ -226,10 +226,16 @@ impl Chain {
 
     /// Append a single block with state to the chain.
     /// This method assumes that blocks attachment to the chain has already been validated.
-    pub fn append_block(&mut self, block: SealedBlockWithSenders, state: BundleStateWithReceipts) {
+    pub fn append_block(
+        &mut self,
+        block: SealedBlockWithSenders,
+        state: BundleStateWithReceipts,
+        trie_updates: Option<TrieUpdates>,
+    ) {
         self.blocks.insert(block.number, block);
         self.state.extend(state);
-        self.trie_updates.take(); // reset
+        self.append_trie_updates(trie_updates);
+        // self.trie_updates.take(); // reset
     }
 
     /// Merge two chains by appending the given chain into the current one.
@@ -249,11 +255,24 @@ impl Chain {
         // Insert blocks from other chain
         self.blocks.extend(other.blocks);
         self.state.extend(other.state);
-        self.trie_updates.take(); // reset
+        self.append_trie_updates(other.trie_updates);
+        // self.trie_updates.take(); // reset
 
         Ok(())
     }
 
+    /// Append trie updates.
+    /// If existing or incoming trie updates are not set, reset as neither is valid anymore.
+    fn append_trie_updates(&mut self, other_trie_updates: Option<TrieUpdates>) {
+        if let Some((trie_updates, other)) = self.trie_updates.as_mut().zip(other_trie_updates) {
+            // Extend trie updates.
+            trie_updates.extend(other);
+        } else {
+            // Reset trie updates as they are no longer valid.
+            self.trie_updates.take();
+        }
+    }
+
     /// Split this chain at the given block.
     ///
     /// The given block will be the last block in the first returned chain.

crates/trie-parallel/benches/root.rs

@@ -59,7 +59,7 @@ pub fn calculate_state_root(c: &mut Criterion) {
         // parallel root
         group.bench_function(BenchmarkId::new("parallel root", size), |b| {
             b.to_async(&runtime).iter_with_setup(
-                || ParallelStateRoot::new(view.clone(), updated_state.clone()),
+                || ParallelStateRoot::new(view.clone(), updated_state.clone(), None),
                 |calculator| async { calculator.incremental_root() },
             );
         });

crates/trie-parallel/src/parallel_root.rs

@@ -11,8 +11,8 @@ use reth_provider::{providers::ConsistentDbView, DatabaseProviderFactory, Provid
 use reth_trie::{
     hashed_cursor::{HashedCursorFactory, HashedPostStateCursorFactory},
     node_iter::{AccountNode, AccountNodeIter},
-    trie_cursor::TrieCursorFactory,
-    updates::TrieUpdates,
+    trie_cursor::{TrieCursorFactory, TrieUpdatesCursorFactory},
+    updates::{TrieUpdates, TrieUpdatesSorted},
     walker::TrieWalker,
     HashedPostState, StorageRoot,
 };
@@ -41,17 +41,24 @@ pub struct ParallelStateRoot<DB, Provider> {
     view: ConsistentDbView<DB, Provider>,
     /// Changed hashed state.
     hashed_state: HashedPostState,
+    /// Trie updates.
+    trie_updates: Option<TrieUpdatesSorted>,
     /// Parallel state root metrics.
     #[cfg(feature = "metrics")]
     metrics: ParallelStateRootMetrics,
 }
 
 impl<DB, Provider> ParallelStateRoot<DB, Provider> {
     /// Create new parallel state root calculator.
-    pub fn new(view: ConsistentDbView<DB, Provider>, hashed_state: HashedPostState) -> Self {
+    pub fn new(
+        view: ConsistentDbView<DB, Provider>,
+        hashed_state: HashedPostState,
+        trie_updates: Option<TrieUpdatesSorted>,
+    ) -> Self {
         Self {
             view,
             hashed_state,
+            trie_updates,
             #[cfg(feature = "metrics")]
             metrics: ParallelStateRootMetrics::default(),
         }
@@ -86,6 +93,7 @@ where
             prefix_sets.storage_prefix_sets,
         );
         let hashed_state_sorted = self.hashed_state.into_sorted();
+        let trie_updates_sorted = self.trie_updates.unwrap_or_default();
 
         // Pre-calculate storage roots in parallel for accounts which were changed.
         tracker.set_precomputed_storage_roots(storage_root_targets.len() as u64);
@@ -95,7 +103,7 @@ where
             .map(|(hashed_address, prefix_set)| {
                 let provider_ro = self.view.provider_ro()?;
                 let storage_root_result = StorageRoot::new_hashed(
-                    provider_ro.tx_ref(),
+                    TrieUpdatesCursorFactory::new(provider_ro.tx_ref(), &trie_updates_sorted),
                     HashedPostStateCursorFactory::new(provider_ro.tx_ref(), &hashed_state_sorted),
                     hashed_address,
                     #[cfg(feature = "metrics")]
@@ -265,7 +273,7 @@ mod tests {
         }
 
         assert_eq!(
-            ParallelStateRoot::new(consistent_view.clone(), HashedPostState::default())
+            ParallelStateRoot::new(consistent_view.clone(), HashedPostState::default(), None)
                 .incremental_root()
                 .unwrap(),
             test_utils::state_root(state.clone())
@@ -297,7 +305,7 @@ mod tests {
         }
 
         assert_eq!(
-            ParallelStateRoot::new(consistent_view, hashed_state).incremental_root().unwrap(),
+            ParallelStateRoot::new(consistent_view, hashed_state, None).incremental_root().unwrap(),
             test_utils::state_root(state)
         );
     }

crates/trie/Cargo.toml

@@ -41,13 +41,18 @@ triehash = { version = "0.8", optional = true }
 reth-primitives = { workspace = true, features = ["test-utils", "arbitrary"] }
 reth-db = { workspace = true, features = ["test-utils"] }
 reth-provider = { workspace = true, features = ["test-utils"] }
+reth-tracing = { workspace = true }
 
 # trie
 triehash = "0.8"
 
 # misc
 proptest.workspace = true
-tokio = { workspace = true, default-features = false, features = ["sync", "rt", "macros"] }
+tokio = { workspace = true, default-features = false, features = [
+    "sync",
+    "rt",
+    "macros",
+] }
 tokio-stream.workspace = true
 once_cell.workspace = true
 serde_json.workspace = true