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

megaeth-labs · Aug 13, 2024 · ac3d62b · ac3d62b
1 parent 8a802da
commit ac3d62b
Show file tree

Hide file tree

Showing 16 changed files with 325 additions and 49 deletions.
diff --git a/bin/reth/src/commands/debug_cmd/build_block.rs b/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 {

diff --git a/bin/reth/src/commands/debug_cmd/in_memory_merkle.rs b/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 {

diff --git a/crates/chain-state/src/in_memory.rs b/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,13 +876,31 @@ 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,
         ) -> ProviderResult<(B256, TrieUpdates)> {
             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,

diff --git a/crates/chain-state/src/memory_overlay.rs b/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 prefix_sets = hashed_state.construct_prefix_sets();
+        self.hashed_state_root_from_nodes(TrieUpdates::default(), hashed_state, prefix_sets)
+    }
+
+    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(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,

diff --git a/crates/revm/src/test_utils.rs b/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,13 +75,31 @@ 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,
     ) -> ProviderResult<(B256, TrieUpdates)> {
         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,

diff --git a/crates/rpc/rpc-eth-types/src/cache/db.rs b/crates/rpc/rpc-eth-types/src/cache/db.rs
@@ -25,13 +25,31 @@ 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,
     ) -> reth_errors::ProviderResult<(B256, reth_trie::updates::TrieUpdates)> {
         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,

diff --git a/crates/storage/provider/src/providers/bundle_state_provider.rs b/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,

diff --git a/crates/storage/provider/src/providers/state/historical.rs b/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,
-    HashedStorage, StateRoot, StorageRoot,
+    prefix_set::TriePrefixSetsMut, proof::Proof, updates::TrieUpdates, witness::TrieWitness,
+    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,