feat(trie): pass state reference to StateProofProvider::proof (#9308)

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

@@ -95,7 +95,12 @@ impl<H> StateProofProvider for MemoryOverlayStateProvider<H>
 where
     H: StateProofProvider + Send,
 {
-    fn proof(&self, address: Address, slots: &[B256]) -> ProviderResult<AccountProof> {
+    fn proof(
+        &self,
+        state: &BundleState,
+        address: Address,
+        slots: &[B256],
+    ) -> ProviderResult<AccountProof> {
         todo!()
     }
 }

crates/revm/src/test_utils.rs

@@ -79,7 +79,12 @@ impl StateRootProvider for StateProviderTest {
 }
 
 impl StateProofProvider for StateProviderTest {
-    fn proof(&self, _address: Address, _slots: &[B256]) -> ProviderResult<AccountProof> {
+    fn proof(
+        &self,
+        _state: &BundleState,
+        _address: Address,
+        _slots: &[B256],
+    ) -> ProviderResult<AccountProof> {
         unimplemented!("proof generation is not supported")
     }
 }

crates/rpc/rpc-eth-api/src/helpers/state.rs

@@ -11,6 +11,7 @@ use reth_rpc_eth_types::{
 use reth_rpc_types::{serde_helpers::JsonStorageKey, EIP1186AccountProofResponse};
 use reth_rpc_types_compat::proof::from_primitive_account_proof;
 use reth_transaction_pool::{PoolTransaction, TransactionPool};
+use revm::db::BundleState;
 use revm_primitives::{BlockEnv, CfgEnvWithHandlerCfg, SpecId};
 
 use super::{EthApiSpec, LoadPendingBlock, SpawnBlocking};
@@ -104,7 +105,7 @@ pub trait EthState: LoadState + SpawnBlocking {
         Ok(self.spawn_blocking_io(move |this| {
             let state = this.state_at_block_id(block_id)?;
             let storage_keys = keys.iter().map(|key| key.0).collect::<Vec<_>>();
-            let proof = state.proof(address, &storage_keys)?;
+            let proof = state.proof(&BundleState::default(), address, &storage_keys)?;
             Ok(from_primitive_account_proof(proof))
         }))
     }

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

@@ -34,10 +34,11 @@ impl<'a> reth_provider::StateRootProvider for StateProviderTraitObjWrapper<'a> {
 impl<'a> reth_provider::StateProofProvider for StateProviderTraitObjWrapper<'a> {
     fn proof(
         &self,
+        state: &revm::db::BundleState,
         address: revm_primitives::Address,
         slots: &[B256],
     ) -> reth_errors::ProviderResult<reth_trie::AccountProof> {
-        self.0.proof(address, slots)
+        self.0.proof(state, address, slots)
     }
 }
 

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

@@ -84,7 +84,12 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateRootProvider
 impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProofProvider
     for BundleStateProvider<SP, EDP>
 {
-    fn proof(&self, _address: Address, _slots: &[B256]) -> ProviderResult<AccountProof> {
+    fn proof(
+        &self,
+        _bundle: &BundleState,
+        _address: Address,
+        _slots: &[B256],
+    ) -> ProviderResult<AccountProof> {
         Err(ProviderError::StateRootNotAvailableForHistoricalBlock)
     }
 }

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

@@ -274,7 +274,12 @@ 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(&self, _address: Address, _slots: &[B256]) -> ProviderResult<AccountProof> {
+    fn proof(
+        &self,
+        _state: &BundleState,
+        _address: Address,
+        _slots: &[B256],
+    ) -> ProviderResult<AccountProof> {
         Err(ProviderError::StateRootNotAvailableForHistoricalBlock)
     }
 }

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

@@ -12,7 +12,7 @@ use reth_primitives::{
 };
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
-use reth_trie::{proof::Proof, updates::TrieUpdates, AccountProof, HashedPostState};
+use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
 use revm::db::BundleState;
 
 /// State provider over latest state that takes tx reference.
@@ -92,9 +92,14 @@ impl<'b, TX: DbTx> StateRootProvider for LatestStateProviderRef<'b, TX> {
 }
 
 impl<'b, TX: DbTx> StateProofProvider for LatestStateProviderRef<'b, TX> {
-    fn proof(&self, address: Address, slots: &[B256]) -> ProviderResult<AccountProof> {
-        Ok(Proof::from_tx(self.tx)
-            .account_proof(address, slots)
+    fn proof(
+        &self,
+        bundle_state: &BundleState,
+        address: Address,
+        slots: &[B256],
+    ) -> ProviderResult<AccountProof> {
+        Ok(HashedPostState::from_bundle_state(&bundle_state.state)
+            .account_proof(self.tx, address, slots)
             .map_err(Into::<reth_db::DatabaseError>::into)?)
     }
 }

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

@@ -46,7 +46,7 @@ macro_rules! delegate_provider_impls {
                 fn state_root_with_updates(&self, state: &revm::db::BundleState) -> reth_storage_errors::provider::ProviderResult<(reth_primitives::B256, reth_trie::updates::TrieUpdates)>;
             }
             StateProofProvider $(where [$($generics)*])? {
-                fn proof(&self, address: reth_primitives::Address, slots: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult<reth_trie::AccountProof>;
+                fn proof(&self, state: &revm::db::BundleState, address: reth_primitives::Address, slots: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult<reth_trie::AccountProof>;
             }
         );
     }

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

@@ -555,7 +555,12 @@ impl StateRootProvider for MockEthProvider {
 }
 
 impl StateProofProvider for MockEthProvider {
-    fn proof(&self, address: Address, _slots: &[B256]) -> ProviderResult<AccountProof> {
+    fn proof(
+        &self,
+        _state: &BundleState,
+        address: Address,
+        _slots: &[B256],
+    ) -> ProviderResult<AccountProof> {
         Ok(AccountProof::new(address))
     }
 }

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

@@ -328,7 +328,12 @@ impl StateRootProvider for NoopProvider {
 }
 
 impl StateProofProvider for NoopProvider {
-    fn proof(&self, address: Address, _slots: &[B256]) -> ProviderResult<AccountProof> {
+    fn proof(
+        &self,
+        _state: &BundleState,
+        address: Address,
+        _slots: &[B256],
+    ) -> ProviderResult<AccountProof> {
         Ok(AccountProof::new(address))
     }
 }

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

@@ -26,6 +26,12 @@ pub trait StateRootProvider: Send + Sync {
 /// A type that can generate state proof on top of a given post state.
 #[auto_impl::auto_impl(&, Box, Arc)]
 pub trait StateProofProvider: Send + Sync {
-    /// Get account and storage proofs.
-    fn proof(&self, address: Address, slots: &[B256]) -> ProviderResult<AccountProof>;
+    /// Get account and storage proofs of target keys in the `BundleState`
+    /// on top of the current state.
+    fn proof(
+        &self,
+        state: &BundleState,
+        address: Address,
+        slots: &[B256],
+    ) -> ProviderResult<AccountProof>;
 }