From ac3d62ba02038ebeb602e8b7dbfcfadf1b8b5e33 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Tue, 13 Aug 2024 12:57:22 -0700 Subject: [PATCH] perf(tree): re-use intermediate nodes (#9836) --- .../src/commands/debug_cmd/build_block.rs | 1 - .../commands/debug_cmd/in_memory_merkle.rs | 1 - crates/chain-state/src/in_memory.rs | 20 ++++- crates/chain-state/src/memory_overlay.rs | 53 ++++++++++--- crates/revm/src/test_utils.rs | 23 +++++- crates/rpc/rpc-eth-types/src/cache/db.rs | 18 +++++ .../src/providers/bundle_state_provider.rs | 35 ++++++++- .../src/providers/state/historical.rs | 41 +++++++++- .../provider/src/providers/state/latest.rs | 28 ++++++- .../provider/src/providers/state/macros.rs | 2 + .../storage/provider/src/test_utils/mock.rs | 27 ++++++- .../storage/provider/src/test_utils/noop.rs | 22 +++++- crates/storage/provider/src/writer/mod.rs | 1 - crates/storage/storage-api/src/trie.rs | 24 +++++- crates/trie/db/src/state.rs | 77 ++++++++++++++----- crates/trie/trie/src/prefix_set.rs | 9 +++ 16 files changed, 329 insertions(+), 53 deletions(-) diff --git a/bin/reth/src/commands/debug_cmd/build_block.rs b/bin/reth/src/commands/debug_cmd/build_block.rs index d93b62564..9f0380605 100644 --- 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 index b0cae6f90..0b8c7ade2 100644 --- 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 index 889c70365..942398f5c 100644 --- 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,6 +876,15 @@ mod tests { Ok(B256::random()) } + fn hashed_state_root_from_nodes( + &self, + _nodes: TrieUpdates, + _post_state: HashedPostState, + _prefix_sets: TriePrefixSetsMut, + ) -> ProviderResult { + 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, diff --git a/crates/chain-state/src/memory_overlay.rs b/crates/chain-state/src/memory_overlay.rs index c1050331d..72d2c8a12 100644 --- 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, /// 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, } @@ -33,10 +37,12 @@ impl MemoryOverlayStateProvider { /// database. pub fn new(in_memory: Vec, historical: Box) -> 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 { - let mut state = self.hashed_post_state.clone(); - state.extend(hashed_state); - self.historical.hashed_state_root(state) + 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 { + 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, diff --git a/crates/revm/src/test_utils.rs b/crates/revm/src/test_utils.rs index 94234e68a..d4a699ccb 100644 --- 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,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 { + 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, diff --git a/crates/rpc/rpc-eth-types/src/cache/db.rs b/crates/rpc/rpc-eth-types/src/cache/db.rs index 1b7c956f3..c72db5c19 100644 --- a/crates/rpc/rpc-eth-types/src/cache/db.rs +++ b/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 { + 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, diff --git a/crates/storage/provider/src/providers/bundle_state_provider.rs b/crates/storage/provider/src/providers/bundle_state_provider.rs index 3c41966bb..0ea481ca8 100644 --- 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 StateRootProvider self.state_provider.hashed_state_root(state) } + fn hashed_state_root_from_nodes( + &self, + _nodes: TrieUpdates, + _hashed_state: HashedPostState, + _prefix_sets: TriePrefixSetsMut, + ) -> ProviderResult { + unimplemented!() + } + fn state_root_with_updates( &self, bundle_state: &BundleState, @@ -99,6 +110,24 @@ impl 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 index 8a95bea60..b2245f074 100644 --- 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 { 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 { + 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, diff --git a/crates/storage/provider/src/providers/state/latest.rs b/crates/storage/provider/src/providers/state/latest.rs index afdd23aae..fe152dccb 100644 --- a/crates/storage/provider/src/providers/state/latest.rs +++ b/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, - HashedStorage, StateRoot, StorageRoot, + prefix_set::TriePrefixSetsMut, proof::Proof, updates::TrieUpdates, witness::TrieWitness, + 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 { - 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 { + 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())) } diff --git a/crates/storage/provider/src/providers/state/macros.rs b/crates/storage/provider/src/providers/state/macros.rs index 9e67e3db5..040093c76 100644 --- a/crates/storage/provider/src/providers/state/macros.rs +++ b/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; fn hashed_state_root(&self, state: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult; + 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; 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; } StateProofProvider $(where [$($generics)*])? { diff --git a/crates/storage/provider/src/test_utils/mock.rs b/crates/storage/provider/src/test_utils/mock.rs index 43f3d30e1..3bd4b772d 100644 --- a/crates/storage/provider/src/test_utils/mock.rs +++ b/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 { - let state_root = self.state_roots.lock().pop().unwrap_or_default(); - Ok(state_root) + Ok(self.state_roots.lock().pop().unwrap_or_default()) + } + + fn hashed_state_root_from_nodes( + &self, + _nodes: TrieUpdates, + _hashed_state: HashedPostState, + _prefix_sets: TriePrefixSetsMut, + ) -> ProviderResult { + Ok(self.state_roots.lock().pop().unwrap_or_default()) } fn hashed_state_root_with_updates( @@ -583,6 +594,16 @@ impl StateRootProvider for MockEthProvider { ) -> ProviderResult { 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 { diff --git a/crates/storage/provider/src/test_utils/noop.rs b/crates/storage/provider/src/test_utils/noop.rs index 4796a0e02..8ff774447 100644 --- a/crates/storage/provider/src/test_utils/noop.rs +++ b/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 { + 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, diff --git a/crates/storage/provider/src/writer/mod.rs b/crates/storage/provider/src/writer/mod.rs index b3250bef0..24a4a4fc2 100644 --- a/crates/storage/provider/src/writer/mod.rs +++ b/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))| ( diff --git a/crates/storage/storage-api/src/trie.rs b/crates/storage/storage-api/src/trie.rs index 961aaa5e8..27800b9f9 100644 --- a/crates/storage/storage-api/src/trie.rs +++ b/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; + /// 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; + /// 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( diff --git a/crates/trie/db/src/state.rs b/crates/trie/db/src/state.rs index 942fbfb18..a138ee579 100644 --- a/crates/trie/db/src/state.rs +++ b/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, - updates::TrieUpdates, HashedPostState, HashedStorage, StateRoot, StateRootProgress, + hashed_cursor::HashedPostStateCursorFactory, prefix_set::TriePrefixSetsMut, + 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( - tx: &'a TX, - post_state: HashedPostState, - intermediate_nodes: TrieUpdates, - ) -> Result; + fn overlay_root(tx: &'a TX, post_state: HashedPostState) -> Result; /// 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; + + /// 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( - tx: &'a TX, - post_state: HashedPostState, - intermediate_nodes: TrieUpdates, - ) -> Result { + fn overlay_root(tx: &'a TX, post_state: HashedPostState) -> Result { 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 { 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") ); } diff --git a/crates/trie/trie/src/prefix_set.rs b/crates/trie/trie/src/prefix_set.rs index 95542e1e1..c94d15eb8 100644 --- a/crates/trie/trie/src/prefix_set.rs +++ b/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.