feat(provider): hashed state methods for trie related providers (#9510)

crates/engine/tree/src/tree/memory_overlay.rs

@@ -4,8 +4,7 @@ use reth_primitives::{Account, Address, BlockNumber, Bytecode, StorageKey, Stora
 use reth_provider::{
     AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
 };
-use reth_trie::{updates::TrieUpdates, AccountProof};
-use revm::db::BundleState;
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
 
 /// A state provider that stores references to in-memory blocks along with their state as well as
 /// the historical state provider for fallback lookups.
@@ -79,13 +78,13 @@ impl<H> StateRootProvider for MemoryOverlayStateProvider<H>
 where
     H: StateRootProvider + Send,
 {
-    fn state_root(&self, bundle_state: &BundleState) -> ProviderResult<B256> {
+    fn hashed_state_root(&self, hashed_state: &HashedPostState) -> ProviderResult<B256> {
         todo!()
     }
 
-    fn state_root_with_updates(
+    fn hashed_state_root_with_updates(
         &self,
-        bundle_state: &BundleState,
+        hashed_state: &HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)> {
         todo!()
     }
@@ -95,9 +94,9 @@ impl<H> StateProofProvider for MemoryOverlayStateProvider<H>
 where
     H: StateProofProvider + Send,
 {
-    fn proof(
+    fn hashed_proof(
         &self,
-        state: &BundleState,
+        hashed_state: &HashedPostState,
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {

crates/revm/src/test_utils.rs

@@ -5,8 +5,7 @@ use reth_storage_api::{
     AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
 };
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof};
-use revm::db::BundleState;
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
 use std::collections::HashMap;
 
 /// Mock state for testing
@@ -66,22 +65,22 @@ impl BlockHashReader for StateProviderTest {
 }
 
 impl StateRootProvider for StateProviderTest {
-    fn state_root(&self, _bundle_state: &BundleState) -> ProviderResult<B256> {
+    fn hashed_state_root(&self, _hashed_state: &HashedPostState) -> ProviderResult<B256> {
         unimplemented!("state root computation is not supported")
     }
 
-    fn state_root_with_updates(
+    fn hashed_state_root_with_updates(
         &self,
-        _bundle_state: &BundleState,
+        _hashed_state: &HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)> {
         unimplemented!("state root computation is not supported")
     }
 }
 
 impl StateProofProvider for StateProviderTest {
-    fn proof(
+    fn hashed_proof(
         &self,
-        _state: &BundleState,
+        _hashed_state: &HashedPostState,
         _address: Address,
         _slots: &[B256],
     ) -> ProviderResult<AccountProof> {

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

@@ -16,29 +16,29 @@ pub type StateCacheDb<'a> = CacheDB<StateProviderDatabase<StateProviderTraitObjW
 pub struct StateProviderTraitObjWrapper<'a>(pub &'a dyn StateProvider);
 
 impl<'a> reth_provider::StateRootProvider for StateProviderTraitObjWrapper<'a> {
-    fn state_root(
+    fn hashed_state_root(
         &self,
-        bundle_state: &revm::db::BundleState,
+        hashed_state: &reth_trie::HashedPostState,
     ) -> reth_errors::ProviderResult<B256> {
-        self.0.state_root(bundle_state)
+        self.0.hashed_state_root(hashed_state)
     }
 
-    fn state_root_with_updates(
+    fn hashed_state_root_with_updates(
         &self,
-        bundle_state: &revm::db::BundleState,
+        hashed_state: &reth_trie::HashedPostState,
     ) -> reth_errors::ProviderResult<(B256, reth_trie::updates::TrieUpdates)> {
-        self.0.state_root_with_updates(bundle_state)
+        self.0.hashed_state_root_with_updates(hashed_state)
     }
 }
 
 impl<'a> reth_provider::StateProofProvider for StateProviderTraitObjWrapper<'a> {
-    fn proof(
+    fn hashed_proof(
         &self,
-        state: &revm::db::BundleState,
+        hashed_state: &reth_trie::HashedPostState,
         address: revm_primitives::Address,
         slots: &[B256],
     ) -> reth_errors::ProviderResult<reth_trie::AccountProof> {
-        self.0.proof(state, address, slots)
+        self.0.hashed_proof(hashed_state, address, slots)
     }
 }
 

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

@@ -4,7 +4,7 @@ use crate::{
 use reth_primitives::{Account, Address, BlockNumber, Bytecode, B256};
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
 use revm::db::BundleState;
 
 /// A state provider that resolves to data from either a wrapped [`crate::ExecutionOutcome`]
@@ -71,6 +71,13 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateRootProvider
         self.state_provider.state_root(&state)
     }
 
+    fn hashed_state_root(&self, hashed_state: &reth_trie::HashedPostState) -> ProviderResult<B256> {
+        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.clone());
+        self.state_provider.hashed_state_root(&state)
+    }
+
     fn state_root_with_updates(
         &self,
         bundle_state: &BundleState,
@@ -79,20 +86,31 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateRootProvider
         state.extend(bundle_state.clone());
         self.state_provider.state_root_with_updates(&state)
     }
+
+    fn hashed_state_root_with_updates(
+        &self,
+        hashed_state: &HashedPostState,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
+        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.clone());
+        self.state_provider.hashed_state_root_with_updates(&state)
+    }
 }
 
 impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProofProvider
     for BundleStateProvider<SP, EDP>
 {
-    fn proof(
+    fn hashed_proof(
         &self,
-        bundle_state: &BundleState,
+        hashed_state: &HashedPostState,
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {
-        let mut state = self.block_execution_data_provider.execution_outcome().state().clone();
-        state.extend(bundle_state.clone());
-        self.state_provider.proof(&state, address, slots)
+        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.clone());
+        self.state_provider.hashed_proof(&state, address, slots)
     }
 }
 

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

@@ -16,7 +16,6 @@ use reth_primitives::{
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
-use revm::db::BundleState;
 use std::fmt::Debug;
 
 /// State provider for a given block number which takes a tx reference.
@@ -257,15 +256,18 @@ impl<'b, TX: DbTx> BlockHashReader for HistoricalStateProviderRef<'b, TX> {
 }
 
 impl<'b, TX: DbTx> StateRootProvider for HistoricalStateProviderRef<'b, TX> {
-    fn state_root(&self, state: &BundleState) -> ProviderResult<B256> {
+    fn hashed_state_root(&self, hashed_state: &HashedPostState) -> ProviderResult<B256> {
         let mut revert_state = self.revert_state()?;
-        revert_state.extend(HashedPostState::from_bundle_state(&state.state));
+        revert_state.extend(hashed_state.clone());
         revert_state.state_root(self.tx).map_err(|err| ProviderError::Database(err.into()))
     }
 
-    fn state_root_with_updates(&self, state: &BundleState) -> ProviderResult<(B256, TrieUpdates)> {
+    fn hashed_state_root_with_updates(
+        &self,
+        hashed_state: &HashedPostState,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
         let mut revert_state = self.revert_state()?;
-        revert_state.extend(HashedPostState::from_bundle_state(&state.state));
+        revert_state.extend(hashed_state.clone());
         revert_state
             .state_root_with_updates(self.tx)
             .map_err(|err| ProviderError::Database(err.into()))
@@ -274,14 +276,14 @@ impl<'b, TX: DbTx> StateRootProvider for HistoricalStateProviderRef<'b, TX> {
 
 impl<'b, TX: DbTx> StateProofProvider for HistoricalStateProviderRef<'b, TX> {
     /// Get account and storage proofs.
-    fn proof(
+    fn hashed_proof(
         &self,
-        state: &BundleState,
+        hashed_state: &HashedPostState,
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {
         let mut revert_state = self.revert_state()?;
-        revert_state.extend(HashedPostState::from_bundle_state(&state.state));
+        revert_state.extend(hashed_state.clone());
         revert_state
             .account_proof(self.tx, address, slots)
             .map_err(|err| ProviderError::Database(err.into()))

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

@@ -13,7 +13,6 @@ use reth_primitives::{
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
 use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
-use revm::db::BundleState;
 
 /// State provider over latest state that takes tx reference.
 #[derive(Debug)]
@@ -75,30 +74,28 @@ impl<'b, TX: DbTx> BlockHashReader for LatestStateProviderRef<'b, TX> {
 }
 
 impl<'b, TX: DbTx> StateRootProvider for LatestStateProviderRef<'b, TX> {
-    fn state_root(&self, bundle_state: &BundleState) -> ProviderResult<B256> {
-        HashedPostState::from_bundle_state(&bundle_state.state)
-            .state_root(self.tx)
-            .map_err(|err| ProviderError::Database(err.into()))
+    fn hashed_state_root(&self, hashed_state: &HashedPostState) -> ProviderResult<B256> {
+        hashed_state.state_root(self.tx).map_err(|err| ProviderError::Database(err.into()))
     }
 
-    fn state_root_with_updates(
+    fn hashed_state_root_with_updates(
         &self,
-        bundle_state: &BundleState,
+        hashed_state: &HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        HashedPostState::from_bundle_state(&bundle_state.state)
+        hashed_state
             .state_root_with_updates(self.tx)
             .map_err(|err| ProviderError::Database(err.into()))
     }
 }
 
 impl<'b, TX: DbTx> StateProofProvider for LatestStateProviderRef<'b, TX> {
-    fn proof(
+    fn hashed_proof(
         &self,
-        bundle_state: &BundleState,
+        hashed_state: &HashedPostState,
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {
-        Ok(HashedPostState::from_bundle_state(&bundle_state.state)
+        Ok(hashed_state
             .account_proof(self.tx, address, slots)
             .map_err(Into::<reth_db::DatabaseError>::into)?)
     }

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

@@ -43,10 +43,13 @@ 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 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)>;
             }
             StateProofProvider $(where [$($generics)*])? {
                 fn proof(&self, state: &revm::db::BundleState, address: reth_primitives::Address, slots: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult<reth_trie::AccountProof>;
+                fn hashed_proof(&self, state: &reth_trie::HashedPostState, address: reth_primitives::Address, slots: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult<reth_trie::AccountProof>;
             }
         );
     }

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

@@ -17,11 +17,8 @@ use reth_primitives::{
 };
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
-use reth_trie::{updates::TrieUpdates, AccountProof};
-use revm::{
-    db::BundleState,
-    primitives::{BlockEnv, CfgEnvWithHandlerCfg},
-};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
+use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg};
 use std::{
     collections::{BTreeMap, HashMap},
     ops::{RangeBounds, RangeInclusive},
@@ -542,22 +539,22 @@ impl AccountReader for MockEthProvider {
 }
 
 impl StateRootProvider for MockEthProvider {
-    fn state_root(&self, _bundle_state: &BundleState) -> ProviderResult<B256> {
+    fn hashed_state_root(&self, _state: &HashedPostState) -> ProviderResult<B256> {
         Ok(B256::default())
     }
 
-    fn state_root_with_updates(
+    fn hashed_state_root_with_updates(
         &self,
-        _bundle_state: &BundleState,
+        _state: &HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)> {
         Ok((B256::default(), Default::default()))
     }
 }
 
 impl StateProofProvider for MockEthProvider {
-    fn proof(
+    fn hashed_proof(
         &self,
-        _state: &BundleState,
+        _hashed_state: &HashedPostState,
         address: Address,
         _slots: &[B256],
     ) -> ProviderResult<AccountProof> {

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

@@ -16,11 +16,8 @@ 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};
-use revm::{
-    db::BundleState,
-    primitives::{BlockEnv, CfgEnvWithHandlerCfg},
-};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
+use revm::primitives::{BlockEnv, CfgEnvWithHandlerCfg};
 use tokio::sync::broadcast;
 
 use crate::{
@@ -315,22 +312,22 @@ impl ChangeSetReader for NoopProvider {
 }
 
 impl StateRootProvider for NoopProvider {
-    fn state_root(&self, _state: &BundleState) -> ProviderResult<B256> {
+    fn hashed_state_root(&self, _state: &HashedPostState) -> ProviderResult<B256> {
         Ok(B256::default())
     }
 
-    fn state_root_with_updates(
+    fn hashed_state_root_with_updates(
         &self,
-        _bundle_state: &BundleState,
+        _state: &HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)> {
         Ok((B256::default(), TrieUpdates::default()))
     }
 }
 
 impl StateProofProvider for NoopProvider {
-    fn proof(
+    fn hashed_proof(
         &self,
-        _state: &BundleState,
+        _hashed_state: &HashedPostState,
         address: Address,
         _slots: &[B256],
     ) -> ProviderResult<AccountProof> {

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

@@ -1,6 +1,6 @@
 use reth_primitives::{Address, B256};
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{updates::TrieUpdates, AccountProof};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
 use revm::db::BundleState;
 
 /// A type that can compute the state root of a given post state.
@@ -13,13 +13,29 @@ pub trait StateRootProvider: Send + Sync {
     /// It is recommended to provide a different implementation from
     /// `state_root_with_updates` since it affects the memory usage during state root
     /// computation.
-    fn state_root(&self, bundle_state: &BundleState) -> ProviderResult<B256>;
+    fn state_root(&self, bundle_state: &BundleState) -> ProviderResult<B256> {
+        let hashed_state = HashedPostState::from_bundle_state(&bundle_state.state);
+        self.hashed_state_root(&hashed_state)
+    }
+
+    /// 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 BundleState on top of the current state with trie
     /// updates to be committed to the database.
     fn state_root_with_updates(
         &self,
         bundle_state: &BundleState,
+    ) -> ProviderResult<(B256, TrieUpdates)> {
+        let hashed_state = HashedPostState::from_bundle_state(&bundle_state.state);
+        self.hashed_state_root_with_updates(&hashed_state)
+    }
+
+    /// Returns the state root of the `HashedPostState` on top of the current state with trie
+    /// updates to be committed to the database.
+    fn hashed_state_root_with_updates(
+        &self,
+        hashed_state: &HashedPostState,
     ) -> ProviderResult<(B256, TrieUpdates)>;
 }
 
@@ -33,5 +49,17 @@ pub trait StateProofProvider: Send + Sync {
         state: &BundleState,
         address: Address,
         slots: &[B256],
+    ) -> ProviderResult<AccountProof> {
+        let hashed_state = HashedPostState::from_bundle_state(&state.state);
+        self.hashed_proof(&hashed_state, address, slots)
+    }
+
+    /// Get account and storage proofs of target keys in the `HashedPostState`
+    /// on top of the current state.
+    fn hashed_proof(
+        &self,
+        hashed_state: &HashedPostState,
+        address: Address,
+        slots: &[B256],
     ) -> ProviderResult<AccountProof>;
 }