perf(tree): re-use intermediate nodes (#9836)

bin/reth/src/commands/debug_cmd/build_block.rs

@@ -288,7 +288,6 @@ impl Command {
                 let (state_root, trie_updates) = StateRoot::overlay_root_with_updates(
                     provider_factory.provider()?.tx_ref(),
                     hashed_post_state.clone(),
-                    Default::default(),
                 )?;
 
                 if state_root != block_with_senders.state_root {

bin/reth/src/commands/debug_cmd/in_memory_merkle.rs

@@ -154,7 +154,6 @@ impl Command {
         let (in_memory_state_root, in_memory_updates) = StateRoot::overlay_root_with_updates(
             provider.tx_ref(),
             execution_outcome.hash_state_slow(),
-            Default::default(),
         )?;
 
         if in_memory_state_root == block.state_root {

crates/chain-state/src/in_memory.rs

@@ -802,7 +802,7 @@ mod tests {
     use reth_storage_api::{
         AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
     };
-    use reth_trie::{AccountProof, HashedStorage};
+    use reth_trie::{prefix_set::TriePrefixSetsMut, AccountProof, HashedStorage};
 
     fn create_mock_state(
         test_block_builder: &mut TestBlockBuilder,
@@ -876,6 +876,15 @@ mod tests {
             Ok(B256::random())
         }
 
+        fn hashed_state_root_from_nodes(
+            &self,
+            _nodes: TrieUpdates,
+            _post_state: HashedPostState,
+            _prefix_sets: TriePrefixSetsMut,
+        ) -> ProviderResult<B256> {
+            Ok(B256::random())
+        }
+
         fn hashed_state_root_with_updates(
             &self,
             _hashed_state: HashedPostState,
@@ -883,6 +892,15 @@ mod tests {
             Ok((B256::random(), TrieUpdates::default()))
         }
 
+        fn hashed_state_root_from_nodes_with_updates(
+            &self,
+            _nodes: TrieUpdates,
+            _post_state: HashedPostState,
+            _prefix_sets: TriePrefixSetsMut,
+        ) -> ProviderResult<(B256, TrieUpdates)> {
+            Ok((B256::random(), TrieUpdates::default()))
+        }
+
         fn hashed_storage_root(
             &self,
             _address: Address,

crates/chain-state/src/memory_overlay.rs

@@ -1,5 +1,3 @@
-use std::collections::HashMap;
-
 use super::ExecutedBlock;
 use reth_errors::ProviderResult;
 use reth_primitives::{
@@ -9,7 +7,11 @@ use reth_storage_api::{
     AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateProviderBox,
     StateRootProvider,
 };
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage};
+use reth_trie::{
+    prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
+    HashedStorage,
+};
+use std::collections::HashMap;
 
 /// A state provider that stores references to in-memory blocks along with their state as well as
 /// the historical state provider for fallback lookups.
@@ -19,6 +21,8 @@ pub struct MemoryOverlayStateProvider {
     pub(crate) in_memory: Vec<ExecutedBlock>,
     /// The collection of hashed state from in-memory blocks.
     pub(crate) hashed_post_state: HashedPostState,
+    /// The collection of aggregated in-memory trie updates.
+    pub(crate) trie_updates: TrieUpdates,
     /// Historical state provider for state lookups that are not found in in-memory blocks.
     pub(crate) historical: Box<dyn StateProvider>,
 }
@@ -33,10 +37,12 @@ impl MemoryOverlayStateProvider {
     ///   database.
     pub fn new(in_memory: Vec<ExecutedBlock>, historical: Box<dyn StateProvider>) -> Self {
         let mut hashed_post_state = HashedPostState::default();
+        let mut trie_updates = TrieUpdates::default();
         for block in in_memory.iter().rev() {
             hashed_post_state.extend(block.hashed_state.as_ref().clone());
+            trie_updates.extend(block.trie.as_ref().clone());
         }
-        Self { in_memory, hashed_post_state, historical }
+        Self { in_memory, hashed_post_state, trie_updates, historical }
     }
 
     /// Turn this state provider into a [`StateProviderBox`]
@@ -91,21 +97,47 @@ impl AccountReader for MemoryOverlayStateProvider {
 }
 
 impl StateRootProvider for MemoryOverlayStateProvider {
-    // TODO: Currently this does not reuse available in-memory trie nodes.
     fn hashed_state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
-        let mut state = self.hashed_post_state.clone();
+        let prefix_sets = hashed_state.construct_prefix_sets();
-        state.extend(hashed_state);
+        self.hashed_state_root_from_nodes(TrieUpdates::default(), hashed_state, prefix_sets)
-        self.historical.hashed_state_root(state)
+    }
+
+    fn hashed_state_root_from_nodes(
+        &self,
+        nodes: TrieUpdates,
+        hashed_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<B256> {
+        let mut trie_nodes = self.trie_updates.clone();
+        trie_nodes.extend(nodes);
+        let mut state = self.hashed_post_state.clone();
+        state.extend(hashed_state);
+        self.historical.hashed_state_root_from_nodes(trie_nodes, state, prefix_sets)
     }
 
-    // TODO: Currently this does not reuse available in-memory trie nodes.
     fn hashed_state_root_with_updates(
         &self,
         hashed_state: HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)> {
+        let prefix_sets = hashed_state.construct_prefix_sets();
+        self.hashed_state_root_from_nodes_with_updates(
+            TrieUpdates::default(),
+            hashed_state,
+            prefix_sets,
+        )
+    }
+
+    fn hashed_state_root_from_nodes_with_updates(
+        &self,
+        nodes: TrieUpdates,
+        hashed_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
+        let mut trie_nodes = self.trie_updates.clone();
+        trie_nodes.extend(nodes);
         let mut state = self.hashed_post_state.clone();
         state.extend(hashed_state);
-        self.historical.hashed_state_root_with_updates(state)
+        self.historical.hashed_state_root_from_nodes_with_updates(trie_nodes, state, prefix_sets)
     }
 
     // TODO: Currently this does not reuse available in-memory trie nodes.
@@ -119,7 +151,6 @@ impl StateRootProvider for MemoryOverlayStateProvider {
 }
 
 impl StateProofProvider for MemoryOverlayStateProvider {
-    // TODO: Currently this does not reuse available in-memory trie nodes.
     fn hashed_proof(
         &self,
         hashed_state: HashedPostState,

crates/revm/src/test_utils.rs

@@ -6,7 +6,10 @@ use reth_storage_api::{
     AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
 };
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage};
+use reth_trie::{
+    prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
+    HashedStorage,
+};
 
 #[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
@@ -72,6 +75,15 @@ impl StateRootProvider for StateProviderTest {
         unimplemented!("state root computation is not supported")
     }
 
+    fn hashed_state_root_from_nodes(
+        &self,
+        _nodes: TrieUpdates,
+        _hashed_state: HashedPostState,
+        _prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<B256> {
+        unimplemented!("state root computation is not supported")
+    }
+
     fn hashed_state_root_with_updates(
         &self,
         _hashed_state: HashedPostState,
@@ -79,6 +91,15 @@ impl StateRootProvider for StateProviderTest {
         unimplemented!("state root computation is not supported")
     }
 
+    fn hashed_state_root_from_nodes_with_updates(
+        &self,
+        _nodes: TrieUpdates,
+        _hashed_state: HashedPostState,
+        _prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
+        unimplemented!("state root computation is not supported")
+    }
+
     fn hashed_storage_root(
         &self,
         _address: Address,

crates/rpc/rpc-eth-types/src/cache/db.rs

@@ -25,6 +25,15 @@ impl<'a> reth_storage_api::StateRootProvider for StateProviderTraitObjWrapper<'a
         self.0.hashed_state_root(hashed_state)
     }
 
+    fn hashed_state_root_from_nodes(
+        &self,
+        nodes: reth_trie::updates::TrieUpdates,
+        hashed_state: reth_trie::HashedPostState,
+        prefix_sets: reth_trie::prefix_set::TriePrefixSetsMut,
+    ) -> reth_errors::ProviderResult<B256> {
+        self.0.hashed_state_root_from_nodes(nodes, hashed_state, prefix_sets)
+    }
+
     fn hashed_state_root_with_updates(
         &self,
         hashed_state: reth_trie::HashedPostState,
@@ -32,6 +41,15 @@ impl<'a> reth_storage_api::StateRootProvider for StateProviderTraitObjWrapper<'a
         self.0.hashed_state_root_with_updates(hashed_state)
     }
 
+    fn hashed_state_root_from_nodes_with_updates(
+        &self,
+        nodes: reth_trie::updates::TrieUpdates,
+        hashed_state: reth_trie::HashedPostState,
+        prefix_sets: reth_trie::prefix_set::TriePrefixSetsMut,
+    ) -> reth_errors::ProviderResult<(B256, reth_trie::updates::TrieUpdates)> {
+        self.0.hashed_state_root_from_nodes_with_updates(nodes, hashed_state, prefix_sets)
+    }
+
     fn hashed_storage_root(
         &self,
         address: Address,

crates/storage/provider/src/providers/bundle_state_provider.rs

@@ -1,13 +1,15 @@
-use std::collections::HashMap;
-
 use crate::{
     AccountReader, BlockHashReader, ExecutionDataProvider, StateProvider, StateRootProvider,
 };
 use reth_primitives::{Account, Address, BlockNumber, Bytecode, Bytes, B256};
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage};
+use reth_trie::{
+    prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
+    HashedStorage,
+};
 use revm::db::BundleState;
+use std::collections::HashMap;
 
 /// A state provider that resolves to data from either a wrapped [`crate::ExecutionOutcome`]
 /// or an underlying state provider.
@@ -80,6 +82,15 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateRootProvider
         self.state_provider.hashed_state_root(state)
     }
 
+    fn hashed_state_root_from_nodes(
+        &self,
+        _nodes: TrieUpdates,
+        _hashed_state: HashedPostState,
+        _prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<B256> {
+        unimplemented!()
+    }
+
     fn state_root_with_updates(
         &self,
         bundle_state: &BundleState,
@@ -99,6 +110,24 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateRootProvider
         self.state_provider.hashed_state_root_with_updates(state)
     }
 
+    fn hashed_state_root_from_nodes_with_updates(
+        &self,
+        nodes: TrieUpdates,
+        hashed_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
+        let bundle_state = self.block_execution_data_provider.execution_outcome().state();
+        let mut state = HashedPostState::from_bundle_state(&bundle_state.state);
+        let mut state_prefix_sets = state.construct_prefix_sets();
+        state.extend(hashed_state);
+        state_prefix_sets.extend(prefix_sets);
+        self.state_provider.hashed_state_root_from_nodes_with_updates(
+            nodes,
+            state,
+            state_prefix_sets,
+        )
+    }
+
     fn hashed_storage_root(
         &self,
         address: Address,

crates/storage/provider/src/providers/state/historical.rs

@@ -16,8 +16,8 @@ use reth_primitives::{
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{
-    proof::Proof, updates::TrieUpdates, witness::TrieWitness, AccountProof, HashedPostState,
+    prefix_set::TriePrefixSetsMut, proof::Proof, updates::TrieUpdates, witness::TrieWitness,
-    HashedStorage, StateRoot, StorageRoot,
+    AccountProof, HashedPostState, HashedStorage, StateRoot, StorageRoot,
 };
 use reth_trie_db::{
     DatabaseHashedPostState, DatabaseProof, DatabaseStateRoot, DatabaseStorageRoot,
@@ -266,7 +266,21 @@ impl<'b, TX: DbTx> StateRootProvider for HistoricalStateProviderRef<'b, TX> {
     fn hashed_state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
         let mut revert_state = self.revert_state()?;
         revert_state.extend(hashed_state);
-        StateRoot::overlay_root(self.tx, revert_state, Default::default())
+        StateRoot::overlay_root(self.tx, revert_state)
+            .map_err(|err| ProviderError::Database(err.into()))
+    }
+
+    fn hashed_state_root_from_nodes(
+        &self,
+        nodes: TrieUpdates,
+        hashed_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<B256> {
+        let mut revert_state = self.revert_state()?;
+        let mut revert_prefix_sets = revert_state.construct_prefix_sets();
+        revert_state.extend(hashed_state);
+        revert_prefix_sets.extend(prefix_sets);
+        StateRoot::overlay_root_from_nodes(self.tx, nodes, revert_state, revert_prefix_sets)
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
@@ -276,10 +290,29 @@ impl<'b, TX: DbTx> StateRootProvider for HistoricalStateProviderRef<'b, TX> {
     ) -> ProviderResult<(B256, TrieUpdates)> {
         let mut revert_state = self.revert_state()?;
         revert_state.extend(hashed_state);
-        StateRoot::overlay_root_with_updates(self.tx, revert_state, Default::default())
+        StateRoot::overlay_root_with_updates(self.tx, revert_state)
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
+    fn hashed_state_root_from_nodes_with_updates(
+        &self,
+        nodes: TrieUpdates,
+        hashed_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
+        let mut revert_state = self.revert_state()?;
+        let mut revert_prefix_sets = revert_state.construct_prefix_sets();
+        revert_state.extend(hashed_state);
+        revert_prefix_sets.extend(prefix_sets);
+        StateRoot::overlay_root_from_nodes_with_updates(
+            self.tx,
+            nodes,
+            revert_state,
+            revert_prefix_sets,
+        )
+        .map_err(|err| ProviderError::Database(err.into()))
+    }
+
     fn hashed_storage_root(
         &self,
         address: Address,

crates/storage/provider/src/providers/state/latest.rs

@@ -16,8 +16,8 @@ use reth_primitives::{
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
 use reth_trie::{
-    proof::Proof, updates::TrieUpdates, witness::TrieWitness, AccountProof, HashedPostState,
+    prefix_set::TriePrefixSetsMut, proof::Proof, updates::TrieUpdates, witness::TrieWitness,
-    HashedStorage, StateRoot, StorageRoot,
+    AccountProof, HashedPostState, HashedStorage, StateRoot, StorageRoot,
 };
 use reth_trie_db::{DatabaseProof, DatabaseStateRoot, DatabaseStorageRoot, DatabaseTrieWitness};
 
@@ -82,7 +82,17 @@ impl<'b, TX: DbTx> BlockHashReader for LatestStateProviderRef<'b, TX> {
 
 impl<'b, TX: DbTx> StateRootProvider for LatestStateProviderRef<'b, TX> {
     fn hashed_state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256> {
-        StateRoot::overlay_root(self.tx, hashed_state, Default::default())
+        StateRoot::overlay_root(self.tx, hashed_state)
+            .map_err(|err| ProviderError::Database(err.into()))
+    }
+
+    fn hashed_state_root_from_nodes(
+        &self,
+        nodes: TrieUpdates,
+        hashed_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<B256> {
+        StateRoot::overlay_root_from_nodes(self.tx, nodes, hashed_state, prefix_sets)
             .map_err(|err| ProviderError::Database(err.into()))
     }
 
@@ -90,7 +100,17 @@ impl<'b, TX: DbTx> StateRootProvider for LatestStateProviderRef<'b, TX> {
         &self,
         hashed_state: HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        StateRoot::overlay_root_with_updates(self.tx, hashed_state, Default::default())
+        StateRoot::overlay_root_with_updates(self.tx, hashed_state)
+            .map_err(|err| ProviderError::Database(err.into()))
+    }
+
+    fn hashed_state_root_from_nodes_with_updates(
+        &self,
+        nodes: TrieUpdates,
+        hashed_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
+        StateRoot::overlay_root_from_nodes_with_updates(self.tx, nodes, hashed_state, prefix_sets)
             .map_err(|err| ProviderError::Database(err.into()))
     }
 

crates/storage/provider/src/providers/state/macros.rs

@@ -44,8 +44,10 @@ macro_rules! delegate_provider_impls {
             StateRootProvider $(where [$($generics)*])? {
                 fn state_root(&self, state: &revm::db::BundleState) -> reth_storage_errors::provider::ProviderResult<reth_primitives::B256>;
                 fn hashed_state_root(&self, state: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult<reth_primitives::B256>;
+                fn hashed_state_root_from_nodes(&self,  nodes: reth_trie::updates::TrieUpdates, state: reth_trie::HashedPostState, prefix_sets: reth_trie::prefix_set::TriePrefixSetsMut) -> reth_storage_errors::provider::ProviderResult<reth_primitives::B256>;
                 fn state_root_with_updates(&self, state: &revm::db::BundleState) -> reth_storage_errors::provider::ProviderResult<(reth_primitives::B256, reth_trie::updates::TrieUpdates)>;
                 fn hashed_state_root_with_updates(&self, state: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult<(reth_primitives::B256, reth_trie::updates::TrieUpdates)>;
+                fn hashed_state_root_from_nodes_with_updates(&self,  nodes: reth_trie::updates::TrieUpdates, state: reth_trie::HashedPostState, prefix_sets: reth_trie::prefix_set::TriePrefixSetsMut) -> reth_storage_errors::provider::ProviderResult<(reth_primitives::B256, reth_trie::updates::TrieUpdates)>;
                 fn hashed_storage_root(&self, address: reth_primitives::Address, storage: reth_trie::HashedStorage) ->  reth_storage_errors::provider::ProviderResult<reth_primitives::B256>;
             }
             StateProofProvider $(where [$($generics)*])? {

crates/storage/provider/src/test_utils/mock.rs

@@ -19,7 +19,10 @@ use reth_primitives::{
 use reth_stages_types::{StageCheckpoint, StageId};
 use reth_storage_api::{StageCheckpointReader, StateProofProvider};
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage};
+use reth_trie::{
+    prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
+    HashedStorage,
+};
 use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg};
 use std::{
     collections::{BTreeMap, HashMap},
@@ -564,8 +567,16 @@ impl StageCheckpointReader for MockEthProvider {
 
 impl StateRootProvider for MockEthProvider {
     fn hashed_state_root(&self, _state: HashedPostState) -> ProviderResult<B256> {
-        let state_root = self.state_roots.lock().pop().unwrap_or_default();
+        Ok(self.state_roots.lock().pop().unwrap_or_default())
-        Ok(state_root)
+    }
+
+    fn hashed_state_root_from_nodes(
+        &self,
+        _nodes: TrieUpdates,
+        _hashed_state: HashedPostState,
+        _prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<B256> {
+        Ok(self.state_roots.lock().pop().unwrap_or_default())
     }
 
     fn hashed_state_root_with_updates(
@@ -583,6 +594,16 @@ impl StateRootProvider for MockEthProvider {
     ) -> ProviderResult<B256> {
         Ok(B256::default())
     }
+
+    fn hashed_state_root_from_nodes_with_updates(
+        &self,
+        _nodes: TrieUpdates,
+        _hashed_state: HashedPostState,
+        _prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
+        let state_root = self.state_roots.lock().pop().unwrap_or_default();
+        Ok((state_root, Default::default()))
+    }
 }
 
 impl StateProofProvider for MockEthProvider {

crates/storage/provider/src/test_utils/noop.rs

@@ -22,7 +22,9 @@ use reth_prune_types::{PruneCheckpoint, PruneSegment};
 use reth_stages_types::{StageCheckpoint, StageId};
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
+use reth_trie::{
+    prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
+};
 use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg};
 use tokio::sync::{broadcast, watch};
 
@@ -321,6 +323,15 @@ impl StateRootProvider for NoopProvider {
         Ok(B256::default())
     }
 
+    fn hashed_state_root_from_nodes(
+        &self,
+        _nodes: TrieUpdates,
+        _hashed_state: HashedPostState,
+        _prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<B256> {
+        Ok(B256::default())
+    }
+
     fn hashed_state_root_with_updates(
         &self,
         _state: HashedPostState,
@@ -328,6 +339,15 @@ impl StateRootProvider for NoopProvider {
         Ok((B256::default(), TrieUpdates::default()))
     }
 
+    fn hashed_state_root_from_nodes_with_updates(
+        &self,
+        _nodes: TrieUpdates,
+        _hashed_state: HashedPostState,
+        _prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
+        Ok((B256::default(), TrieUpdates::default()))
+    }
+
     fn hashed_storage_root(
         &self,
         _address: Address,

crates/storage/provider/src/writer/mod.rs

@@ -1413,7 +1413,6 @@ mod tests {
                         Vec::new()
                     )
                     .hash_state_slow(),
-                    Default::default()
                 )
                 .unwrap(),
                 state_root(expected.clone().into_iter().map(|(address, (account, storage))| (

crates/storage/storage-api/src/trie.rs

@@ -1,6 +1,9 @@
 use reth_primitives::{Address, Bytes, B256};
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState, HashedStorage};
+use reth_trie::{
+    prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
+    HashedStorage,
+};
 use revm::db::BundleState;
 use std::collections::HashMap;
 
@@ -21,6 +24,16 @@ pub trait StateRootProvider: Send + Sync {
     /// Returns the state root of the `HashedPostState` on top of the current state.
     fn hashed_state_root(&self, hashed_state: HashedPostState) -> ProviderResult<B256>;
 
+    /// Returns the state root of the `HashedPostState` on top of the current state but re-uses the
+    /// intermediate nodes to speed up the computation. It's up to the caller to construct the
+    /// prefix sets and inform the provider of the trie paths that have changes.
+    fn hashed_state_root_from_nodes(
+        &self,
+        nodes: TrieUpdates,
+        hashed_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<B256>;
+
     /// Returns the state root of the BundleState on top of the current state with trie
     /// updates to be committed to the database.
     fn state_root_with_updates(
@@ -37,6 +50,15 @@ pub trait StateRootProvider: Send + Sync {
         hashed_state: HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)>;
 
+    /// Returns state root and trie updates.
+    /// See [`StateRootProvider::hashed_state_root_from_nodes`] for more info.
+    fn hashed_state_root_from_nodes_with_updates(
+        &self,
+        nodes: TrieUpdates,
+        hashed_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> ProviderResult<(B256, TrieUpdates)>;
+
     /// Returns the storage root of the `HashedStorage` for target address on top of the current
     /// state.
     fn hashed_storage_root(

crates/trie/db/src/state.rs

@@ -9,8 +9,9 @@ use reth_execution_errors::StateRootError;
 use reth_primitives::{keccak256, Account, Address, BlockNumber, B256, U256};
 use reth_storage_errors::db::DatabaseError;
 use reth_trie::{
-    hashed_cursor::HashedPostStateCursorFactory, trie_cursor::InMemoryTrieCursorFactory,
+    hashed_cursor::HashedPostStateCursorFactory, prefix_set::TriePrefixSetsMut,
-    updates::TrieUpdates, HashedPostState, HashedStorage, StateRoot, StateRootProgress,
+    trie_cursor::InMemoryTrieCursorFactory, updates::TrieUpdates, HashedPostState, HashedStorage,
+    StateRoot, StateRootProgress,
 };
 use std::{
     collections::{hash_map, HashMap},
@@ -92,29 +93,38 @@ pub trait DatabaseStateRoot<'a, TX>: Sized {
     ///     Some(Account { nonce: 1, balance: U256::from(10), bytecode_hash: None }),
     /// );
     ///
-    /// // Initialize intermediate nodes if any.
-    /// let intermediate_nodes = TrieUpdates::default();
-    ///
     /// // Calculate the state root
     /// let tx = db.tx().expect("failed to create transaction");
-    /// let state_root = StateRoot::overlay_root(&tx, hashed_state, intermediate_nodes);
+    /// let state_root = StateRoot::overlay_root(&tx, hashed_state);
     /// ```
     ///
     /// # Returns
     ///
     /// The state root for this [`HashedPostState`].
-    fn overlay_root(
+    fn overlay_root(tx: &'a TX, post_state: HashedPostState) -> Result<B256, StateRootError>;
-        tx: &'a TX,
-        post_state: HashedPostState,
-        intermediate_nodes: TrieUpdates,
-    ) -> Result<B256, StateRootError>;
 
     /// Calculates the state root for this [`HashedPostState`] and returns it alongside trie
     /// updates. See [`Self::overlay_root`] for more info.
     fn overlay_root_with_updates(
         tx: &'a TX,
         post_state: HashedPostState,
+    ) -> Result<(B256, TrieUpdates), StateRootError>;
+
+    /// Calculates the state root for provided [`HashedPostState`] using cached intermediate nodes.
+    fn overlay_root_from_nodes(
+        tx: &'a TX,
         intermediate_nodes: TrieUpdates,
+        post_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> Result<B256, StateRootError>;
+
+    /// Calculates the state root and trie updates for provided [`HashedPostState`] using
+    /// cached intermediate nodes.
+    fn overlay_root_from_nodes_with_updates(
+        tx: &'a TX,
+        intermediate_nodes: TrieUpdates,
+        post_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
     ) -> Result<(B256, TrieUpdates), StateRootError>;
 }
 
@@ -164,16 +174,11 @@ impl<'a, TX: DbTx> DatabaseStateRoot<'a, TX>
         Self::incremental_root_calculator(tx, range)?.root_with_progress()
     }
 
-    fn overlay_root(
+    fn overlay_root(tx: &'a TX, post_state: HashedPostState) -> Result<B256, StateRootError> {
-        tx: &'a TX,
-        post_state: HashedPostState,
-        intermediate_nodes: TrieUpdates,
-    ) -> Result<B256, StateRootError> {
         let prefix_sets = post_state.construct_prefix_sets().freeze();
         let state_sorted = post_state.into_sorted();
-        let nodes_sorted = intermediate_nodes.into_sorted();
         StateRoot::new(
-            InMemoryTrieCursorFactory::new(DatabaseTrieCursorFactory::new(tx), &nodes_sorted),
+            DatabaseTrieCursorFactory::new(tx),
             HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
         )
         .with_prefix_sets(prefix_sets)
@@ -183,16 +188,46 @@ impl<'a, TX: DbTx> DatabaseStateRoot<'a, TX>
     fn overlay_root_with_updates(
         tx: &'a TX,
         post_state: HashedPostState,
-        intermediate_nodes: TrieUpdates,
     ) -> Result<(B256, TrieUpdates), StateRootError> {
         let prefix_sets = post_state.construct_prefix_sets().freeze();
+        let state_sorted = post_state.into_sorted();
+        StateRoot::new(
+            DatabaseTrieCursorFactory::new(tx),
+            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
+        )
+        .with_prefix_sets(prefix_sets)
+        .root_with_updates()
+    }
+
+    fn overlay_root_from_nodes(
+        tx: &'a TX,
+        intermediate_nodes: TrieUpdates,
+        post_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> Result<B256, StateRootError> {
         let state_sorted = post_state.into_sorted();
         let nodes_sorted = intermediate_nodes.into_sorted();
         StateRoot::new(
             InMemoryTrieCursorFactory::new(DatabaseTrieCursorFactory::new(tx), &nodes_sorted),
             HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
         )
-        .with_prefix_sets(prefix_sets)
+        .with_prefix_sets(prefix_sets.freeze())
+        .root()
+    }
+
+    fn overlay_root_from_nodes_with_updates(
+        tx: &'a TX,
+        intermediate_nodes: TrieUpdates,
+        post_state: HashedPostState,
+        prefix_sets: TriePrefixSetsMut,
+    ) -> Result<(B256, TrieUpdates), StateRootError> {
+        let state_sorted = post_state.into_sorted();
+        let nodes_sorted = intermediate_nodes.into_sorted();
+        StateRoot::new(
+            InMemoryTrieCursorFactory::new(DatabaseTrieCursorFactory::new(tx), &nodes_sorted),
+            HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &state_sorted),
+        )
+        .with_prefix_sets(prefix_sets.freeze())
         .root_with_updates()
     }
 }
@@ -277,7 +312,7 @@ mod tests {
         let db = create_test_rw_db();
         let tx = db.tx().expect("failed to create transaction");
         assert_eq!(
-            StateRoot::overlay_root(&tx, post_state, Default::default()).unwrap(),
+            StateRoot::overlay_root(&tx, post_state).unwrap(),
             hex!("b464525710cafcf5d4044ac85b72c08b1e76231b8d91f288fe438cc41d8eaafd")
         );
     }

crates/trie/trie/src/prefix_set.rs

@@ -18,6 +18,15 @@ pub struct TriePrefixSetsMut {
 }
 
 impl TriePrefixSetsMut {
+    /// Extends prefix sets with contents of another prefix set.
+    pub fn extend(&mut self, other: Self) {
+        self.account_prefix_set.extend(other.account_prefix_set.keys);
+        for (hashed_address, prefix_set) in other.storage_prefix_sets {
+            self.storage_prefix_sets.entry(hashed_address).or_default().extend(prefix_set.keys);
+        }
+        self.destroyed_accounts.extend(other.destroyed_accounts);
+    }
+
     /// Returns a `TriePrefixSets` with the same elements as these sets.
     ///
     /// If not yet sorted, the elements will be sorted and deduplicated.