From fdd64972b9ff3d434f24bd7dbd8b04fd107b112d Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Mon, 16 Sep 2024 11:18:23 +0200 Subject: [PATCH] feat(trie): integrate `TrieInput` into `StateProofProvider` methods (#10929) --- crates/chain-state/src/in_memory.rs | 4 +-- crates/chain-state/src/memory_overlay.rs | 18 ++++++------ .../engine/invalid-block-hooks/src/witness.rs | 2 +- crates/revm/src/test_utils.rs | 4 +-- crates/rpc/rpc-eth-types/src/cache/db.rs | 8 +++--- crates/rpc/rpc/src/debug.rs | 5 ++-- .../src/providers/bundle_state_provider.rs | 14 ++++------ .../src/providers/state/historical.rs | 15 ++++------ .../provider/src/providers/state/latest.rs | 8 +++--- .../provider/src/providers/state/macros.rs | 4 +-- .../storage/provider/src/test_utils/mock.rs | 4 +-- .../storage/provider/src/test_utils/noop.rs | 4 +-- crates/storage/storage-api/src/trie.rs | 4 +-- crates/trie/db/src/proof.rs | 28 ++++++++++++------- crates/trie/db/src/witness.rs | 26 ++++++++++------- crates/trie/trie/src/input.rs | 2 +- crates/trie/trie/src/proof.rs | 10 +++++++ crates/trie/trie/src/witness.rs | 10 +++++++ 18 files changed, 98 insertions(+), 72 deletions(-) diff --git a/crates/chain-state/src/in_memory.rs b/crates/chain-state/src/in_memory.rs index 4125f80a8..b1d9531b1 100644 --- a/crates/chain-state/src/in_memory.rs +++ b/crates/chain-state/src/in_memory.rs @@ -945,7 +945,7 @@ mod tests { impl StateProofProvider for MockStateProvider { fn proof( &self, - _hashed_state: HashedPostState, + _input: TrieInput, _address: Address, _slots: &[B256], ) -> ProviderResult { @@ -954,7 +954,7 @@ mod tests { fn witness( &self, - _overlay: HashedPostState, + _input: TrieInput, _target: HashedPostState, ) -> ProviderResult> { Ok(HashMap::default()) diff --git a/crates/chain-state/src/memory_overlay.rs b/crates/chain-state/src/memory_overlay.rs index fb88e4ab3..fbd30a4d9 100644 --- a/crates/chain-state/src/memory_overlay.rs +++ b/crates/chain-state/src/memory_overlay.rs @@ -136,27 +136,25 @@ impl StorageRootProvider for MemoryOverlayStateProvider { } impl StateProofProvider for MemoryOverlayStateProvider { - // TODO: Currently this does not reuse available in-memory trie nodes. fn proof( &self, - state: HashedPostState, + mut input: TrieInput, address: Address, slots: &[B256], ) -> ProviderResult { - let mut hashed_state = self.trie_state().state.clone(); - hashed_state.extend(state); - self.historical.proof(hashed_state, address, slots) + let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone(); + input.prepend_cached(nodes, state); + self.historical.proof(input, address, slots) } - // TODO: Currently this does not reuse available in-memory trie nodes. fn witness( &self, - overlay: HashedPostState, + mut input: TrieInput, target: HashedPostState, ) -> ProviderResult> { - let mut hashed_state = self.trie_state().state.clone(); - hashed_state.extend(overlay); - self.historical.witness(hashed_state, target) + let MemoryOverlayTrieState { nodes, state } = self.trie_state().clone(); + input.prepend_cached(nodes, state); + self.historical.witness(input, target) } } diff --git a/crates/engine/invalid-block-hooks/src/witness.rs b/crates/engine/invalid-block-hooks/src/witness.rs index ead7d8e79..c3829e94e 100644 --- a/crates/engine/invalid-block-hooks/src/witness.rs +++ b/crates/engine/invalid-block-hooks/src/witness.rs @@ -145,7 +145,7 @@ where // Generate an execution witness for the aggregated state of accessed accounts. // Destruct the cache database to retrieve the state provider. let state_provider = db.database.into_inner(); - let witness = state_provider.witness(HashedPostState::default(), hashed_state.clone())?; + let witness = state_provider.witness(Default::default(), hashed_state.clone())?; // Write the witness to the output directory. let mut file = File::options() diff --git a/crates/revm/src/test_utils.rs b/crates/revm/src/test_utils.rs index 1a1c1d8c4..e32e98e3b 100644 --- a/crates/revm/src/test_utils.rs +++ b/crates/revm/src/test_utils.rs @@ -103,7 +103,7 @@ impl StorageRootProvider for StateProviderTest { impl StateProofProvider for StateProviderTest { fn proof( &self, - _hashed_state: HashedPostState, + _input: TrieInput, _address: Address, _slots: &[B256], ) -> ProviderResult { @@ -112,7 +112,7 @@ impl StateProofProvider for StateProviderTest { fn witness( &self, - _overlay: HashedPostState, + _input: TrieInput, _target: HashedPostState, ) -> ProviderResult> { unimplemented!("witness generation is not supported") diff --git a/crates/rpc/rpc-eth-types/src/cache/db.rs b/crates/rpc/rpc-eth-types/src/cache/db.rs index 21dec6bdf..1dcfae25f 100644 --- a/crates/rpc/rpc-eth-types/src/cache/db.rs +++ b/crates/rpc/rpc-eth-types/src/cache/db.rs @@ -60,19 +60,19 @@ impl<'a> reth_storage_api::StorageRootProvider for StateProviderTraitObjWrapper< impl<'a> reth_storage_api::StateProofProvider for StateProviderTraitObjWrapper<'a> { fn proof( &self, - hashed_state: reth_trie::HashedPostState, + input: reth_trie::TrieInput, address: revm_primitives::Address, slots: &[B256], ) -> reth_errors::ProviderResult { - self.0.proof(hashed_state, address, slots) + self.0.proof(input, address, slots) } fn witness( &self, - overlay: reth_trie::HashedPostState, + input: reth_trie::TrieInput, target: reth_trie::HashedPostState, ) -> reth_errors::ProviderResult> { - self.0.witness(overlay, target) + self.0.witness(input, target) } } diff --git a/crates/rpc/rpc/src/debug.rs b/crates/rpc/rpc/src/debug.rs index 95af81a7b..215843d1c 100644 --- a/crates/rpc/rpc/src/debug.rs +++ b/crates/rpc/rpc/src/debug.rs @@ -670,9 +670,8 @@ where // Generate an execution witness for the aggregated state of accessed accounts. // Destruct the cache database to retrieve the state provider. let state_provider = db.database.into_inner(); - let witness = state_provider - .witness(HashedPostState::default(), hashed_state) - .map_err(Into::into)?; + let witness = + state_provider.witness(Default::default(), hashed_state).map_err(Into::into)?; Ok(ExecutionWitness { witness, diff --git a/crates/storage/provider/src/providers/bundle_state_provider.rs b/crates/storage/provider/src/providers/bundle_state_provider.rs index 465d934f6..8a911d716 100644 --- a/crates/storage/provider/src/providers/bundle_state_provider.rs +++ b/crates/storage/provider/src/providers/bundle_state_provider.rs @@ -124,25 +124,23 @@ impl StateProofProvider { fn proof( &self, - hashed_state: HashedPostState, + mut input: TrieInput, address: Address, slots: &[B256], ) -> ProviderResult { let bundle_state = self.block_execution_data_provider.execution_outcome().state(); - let mut state = HashedPostState::from_bundle_state(&bundle_state.state); - state.extend(hashed_state); - self.state_provider.proof(state, address, slots) + input.prepend(HashedPostState::from_bundle_state(&bundle_state.state)); + self.state_provider.proof(input, address, slots) } fn witness( &self, - overlay: HashedPostState, + mut input: TrieInput, target: HashedPostState, ) -> ProviderResult> { let bundle_state = self.block_execution_data_provider.execution_outcome().state(); - let mut state = HashedPostState::from_bundle_state(&bundle_state.state); - state.extend(overlay); - self.state_provider.witness(state, target) + input.prepend(HashedPostState::from_bundle_state(&bundle_state.state)); + self.state_provider.witness(input, target) } } diff --git a/crates/storage/provider/src/providers/state/historical.rs b/crates/storage/provider/src/providers/state/historical.rs index b1caf66fa..6b4ee7fb8 100644 --- a/crates/storage/provider/src/providers/state/historical.rs +++ b/crates/storage/provider/src/providers/state/historical.rs @@ -335,25 +335,22 @@ impl<'b, TX: DbTx> StateProofProvider for HistoricalStateProviderRef<'b, TX> { /// Get account and storage proofs. fn proof( &self, - hashed_state: HashedPostState, + mut input: TrieInput, address: Address, slots: &[B256], ) -> ProviderResult { - let mut revert_state = self.revert_state()?; - revert_state.extend(hashed_state); - Proof::overlay_account_proof(self.tx, revert_state, address, slots) + input.prepend(self.revert_state()?); + Proof::overlay_account_proof(self.tx, input, address, slots) .map_err(Into::::into) } fn witness( &self, - overlay: HashedPostState, + mut input: TrieInput, target: HashedPostState, ) -> ProviderResult> { - let mut revert_state = self.revert_state()?; - revert_state.extend(overlay); - TrieWitness::overlay_witness(self.tx, revert_state, target) - .map_err(Into::::into) + input.prepend(self.revert_state()?); + TrieWitness::overlay_witness(self.tx, input, target).map_err(Into::::into) } } diff --git a/crates/storage/provider/src/providers/state/latest.rs b/crates/storage/provider/src/providers/state/latest.rs index 4ef98c5f8..0e6d7f69e 100644 --- a/crates/storage/provider/src/providers/state/latest.rs +++ b/crates/storage/provider/src/providers/state/latest.rs @@ -121,20 +121,20 @@ impl<'b, TX: DbTx> StorageRootProvider for LatestStateProviderRef<'b, TX> { impl<'b, TX: DbTx> StateProofProvider for LatestStateProviderRef<'b, TX> { fn proof( &self, - hashed_state: HashedPostState, + input: TrieInput, address: Address, slots: &[B256], ) -> ProviderResult { - Proof::overlay_account_proof(self.tx, hashed_state, address, slots) + Proof::overlay_account_proof(self.tx, input, address, slots) .map_err(Into::::into) } fn witness( &self, - overlay: HashedPostState, + input: TrieInput, target: HashedPostState, ) -> ProviderResult> { - TrieWitness::overlay_witness(self.tx, overlay, target).map_err(Into::::into) + TrieWitness::overlay_witness(self.tx, input, target).map_err(Into::::into) } } diff --git a/crates/storage/provider/src/providers/state/macros.rs b/crates/storage/provider/src/providers/state/macros.rs index 435d06a2a..161f89dd9 100644 --- a/crates/storage/provider/src/providers/state/macros.rs +++ b/crates/storage/provider/src/providers/state/macros.rs @@ -51,8 +51,8 @@ macro_rules! delegate_provider_impls { fn storage_root(&self, address: reth_primitives::Address, storage: reth_trie::HashedStorage) -> reth_storage_errors::provider::ProviderResult; } StateProofProvider $(where [$($generics)*])? { - fn proof(&self, state: reth_trie::HashedPostState, address: reth_primitives::Address, slots: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult; - fn witness(&self, state: reth_trie::HashedPostState, target: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult>; + fn proof(&self, input: reth_trie::TrieInput, address: reth_primitives::Address, slots: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult; + fn witness(&self, input: reth_trie::TrieInput, target: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult>; } ); } diff --git a/crates/storage/provider/src/test_utils/mock.rs b/crates/storage/provider/src/test_utils/mock.rs index b081dfed2..b00aa2ff0 100644 --- a/crates/storage/provider/src/test_utils/mock.rs +++ b/crates/storage/provider/src/test_utils/mock.rs @@ -617,7 +617,7 @@ impl StorageRootProvider for MockEthProvider { impl StateProofProvider for MockEthProvider { fn proof( &self, - _hashed_state: HashedPostState, + _input: TrieInput, address: Address, _slots: &[B256], ) -> ProviderResult { @@ -626,7 +626,7 @@ impl StateProofProvider for MockEthProvider { fn witness( &self, - _overlay: HashedPostState, + _input: TrieInput, _target: HashedPostState, ) -> ProviderResult> { Ok(HashMap::default()) diff --git a/crates/storage/provider/src/test_utils/noop.rs b/crates/storage/provider/src/test_utils/noop.rs index b78e153dd..e5ef41a91 100644 --- a/crates/storage/provider/src/test_utils/noop.rs +++ b/crates/storage/provider/src/test_utils/noop.rs @@ -355,7 +355,7 @@ impl StorageRootProvider for NoopProvider { impl StateProofProvider for NoopProvider { fn proof( &self, - _hashed_state: HashedPostState, + _input: TrieInput, address: Address, _slots: &[B256], ) -> ProviderResult { @@ -364,7 +364,7 @@ impl StateProofProvider for NoopProvider { fn witness( &self, - _overlay: HashedPostState, + _input: TrieInput, _target: HashedPostState, ) -> ProviderResult> { Ok(HashMap::default()) diff --git a/crates/storage/storage-api/src/trie.rs b/crates/storage/storage-api/src/trie.rs index eaac67cfe..5a3b74b30 100644 --- a/crates/storage/storage-api/src/trie.rs +++ b/crates/storage/storage-api/src/trie.rs @@ -51,7 +51,7 @@ pub trait StateProofProvider: Send + Sync { /// on top of the current state. fn proof( &self, - hashed_state: HashedPostState, + input: TrieInput, address: Address, slots: &[B256], ) -> ProviderResult; @@ -59,7 +59,7 @@ pub trait StateProofProvider: Send + Sync { /// Get trie witness for provided state. fn witness( &self, - overlay: HashedPostState, + input: TrieInput, target: HashedPostState, ) -> ProviderResult>; } diff --git a/crates/trie/db/src/proof.rs b/crates/trie/db/src/proof.rs index 8739dfc51..cffe0495b 100644 --- a/crates/trie/db/src/proof.rs +++ b/crates/trie/db/src/proof.rs @@ -2,7 +2,10 @@ use crate::{DatabaseHashedCursorFactory, DatabaseTrieCursorFactory}; use reth_db_api::transaction::DbTx; use reth_execution_errors::StateProofError; use reth_primitives::{Address, B256}; -use reth_trie::{hashed_cursor::HashedPostStateCursorFactory, proof::Proof, HashedPostState}; +use reth_trie::{ + hashed_cursor::HashedPostStateCursorFactory, proof::Proof, + trie_cursor::InMemoryTrieCursorFactory, TrieInput, +}; use reth_trie_common::AccountProof; /// Extends [`Proof`] with operations specific for working with a database transaction. @@ -10,10 +13,10 @@ pub trait DatabaseProof<'a, TX> { /// Create a new [Proof] from database transaction. fn from_tx(tx: &'a TX) -> Self; - /// Generates the state proof for target account and slots on top of this [`HashedPostState`]. + /// Generates the state proof for target account based on [`TrieInput`]. fn overlay_account_proof( tx: &'a TX, - post_state: HashedPostState, + input: TrieInput, address: Address, slots: &[B256], ) -> Result; @@ -29,17 +32,22 @@ impl<'a, TX: DbTx> DatabaseProof<'a, TX> fn overlay_account_proof( tx: &'a TX, - post_state: HashedPostState, + input: TrieInput, address: Address, slots: &[B256], ) -> Result { - let prefix_sets = post_state.construct_prefix_sets(); - let sorted = post_state.into_sorted(); - let hashed_cursor_factory = - HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &sorted); + let nodes_sorted = input.nodes.into_sorted(); + let state_sorted = input.state.into_sorted(); Self::from_tx(tx) - .with_hashed_cursor_factory(hashed_cursor_factory) - .with_prefix_sets_mut(prefix_sets) + .with_trie_cursor_factory(InMemoryTrieCursorFactory::new( + DatabaseTrieCursorFactory::new(tx), + &nodes_sorted, + )) + .with_hashed_cursor_factory(HashedPostStateCursorFactory::new( + DatabaseHashedCursorFactory::new(tx), + &state_sorted, + )) + .with_prefix_sets_mut(input.prefix_sets) .account_proof(address, slots) } } diff --git a/crates/trie/db/src/witness.rs b/crates/trie/db/src/witness.rs index deede885a..de5687191 100644 --- a/crates/trie/db/src/witness.rs +++ b/crates/trie/db/src/witness.rs @@ -3,7 +3,8 @@ use reth_db_api::transaction::DbTx; use reth_execution_errors::TrieWitnessError; use reth_primitives::{Bytes, B256}; use reth_trie::{ - hashed_cursor::HashedPostStateCursorFactory, witness::TrieWitness, HashedPostState, + hashed_cursor::HashedPostStateCursorFactory, trie_cursor::InMemoryTrieCursorFactory, + witness::TrieWitness, HashedPostState, TrieInput, }; use std::collections::HashMap; @@ -12,10 +13,10 @@ pub trait DatabaseTrieWitness<'a, TX> { /// Create a new [`TrieWitness`] from database transaction. fn from_tx(tx: &'a TX) -> Self; - /// Generates trie witness for target state on top of this [`HashedPostState`]. + /// Generates trie witness for target state based on [`TrieInput`]. fn overlay_witness( tx: &'a TX, - post_state: HashedPostState, + input: TrieInput, target: HashedPostState, ) -> Result, TrieWitnessError>; } @@ -29,16 +30,21 @@ impl<'a, TX: DbTx> DatabaseTrieWitness<'a, TX> fn overlay_witness( tx: &'a TX, - post_state: HashedPostState, + input: TrieInput, target: HashedPostState, ) -> Result, TrieWitnessError> { - let prefix_sets = post_state.construct_prefix_sets(); - let sorted = post_state.into_sorted(); - let hashed_cursor_factory = - HashedPostStateCursorFactory::new(DatabaseHashedCursorFactory::new(tx), &sorted); + let nodes_sorted = input.nodes.into_sorted(); + let state_sorted = input.state.into_sorted(); Self::from_tx(tx) - .with_hashed_cursor_factory(hashed_cursor_factory) - .with_prefix_sets_mut(prefix_sets) + .with_trie_cursor_factory(InMemoryTrieCursorFactory::new( + DatabaseTrieCursorFactory::new(tx), + &nodes_sorted, + )) + .with_hashed_cursor_factory(HashedPostStateCursorFactory::new( + DatabaseHashedCursorFactory::new(tx), + &state_sorted, + )) + .with_prefix_sets_mut(input.prefix_sets) .compute(target) } } diff --git a/crates/trie/trie/src/input.rs b/crates/trie/trie/src/input.rs index f8ca95361..3364f626a 100644 --- a/crates/trie/trie/src/input.rs +++ b/crates/trie/trie/src/input.rs @@ -1,7 +1,7 @@ use crate::{prefix_set::TriePrefixSetsMut, updates::TrieUpdates, HashedPostState}; /// Inputs for trie-related computations. -#[derive(Debug)] +#[derive(Default, Debug)] pub struct TrieInput { /// The collection of cached in-memory intermediate trie nodes that /// can be reused for computation. diff --git a/crates/trie/trie/src/proof.rs b/crates/trie/trie/src/proof.rs index 122518974..0a8b23f9c 100644 --- a/crates/trie/trie/src/proof.rs +++ b/crates/trie/trie/src/proof.rs @@ -42,6 +42,16 @@ impl Proof { } } + /// Set the trie cursor factory. + pub fn with_trie_cursor_factory(self, trie_cursor_factory: TF) -> Proof { + Proof { + trie_cursor_factory, + hashed_cursor_factory: self.hashed_cursor_factory, + prefix_sets: self.prefix_sets, + targets: self.targets, + } + } + /// Set the hashed cursor factory. pub fn with_hashed_cursor_factory(self, hashed_cursor_factory: HF) -> Proof { Proof { diff --git a/crates/trie/trie/src/witness.rs b/crates/trie/trie/src/witness.rs index fe6db2754..eb2862dd6 100644 --- a/crates/trie/trie/src/witness.rs +++ b/crates/trie/trie/src/witness.rs @@ -36,6 +36,16 @@ impl TrieWitness { } } + /// Set the trie cursor factory. + pub fn with_trie_cursor_factory(self, trie_cursor_factory: TF) -> TrieWitness { + TrieWitness { + trie_cursor_factory, + hashed_cursor_factory: self.hashed_cursor_factory, + prefix_sets: self.prefix_sets, + witness: self.witness, + } + } + /// Set the hashed cursor factory. pub fn with_hashed_cursor_factory(self, hashed_cursor_factory: HF) -> TrieWitness { TrieWitness {