feat: extract proof generation into StateProofProvider (#9303)

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

@@ -1,7 +1,9 @@
 use super::ExecutedBlock;
 use reth_errors::ProviderResult;
 use reth_primitives::{Account, Address, BlockNumber, Bytecode, StorageKey, StorageValue, B256};
-use reth_provider::{AccountReader, BlockHashReader, StateProvider, StateRootProvider};
+use reth_provider::{
+    AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
+};
 use reth_trie::{updates::TrieUpdates, AccountProof};
 use revm::db::BundleState;
 
@@ -89,6 +91,15 @@ where
     }
 }
 
+impl<H> StateProofProvider for MemoryOverlayStateProvider<H>
+where
+    H: StateProofProvider + Send,
+{
+    fn proof(&self, address: Address, slots: &[B256]) -> ProviderResult<AccountProof> {
+        todo!()
+    }
+}
+
 impl<H> StateProvider for MemoryOverlayStateProvider<H>
 where
     H: StateProvider + Send,
@@ -116,8 +127,4 @@ where
 
         self.historical.bytecode_by_hash(code_hash)
     }
-
-    fn proof(&self, address: Address, keys: &[B256]) -> ProviderResult<AccountProof> {
-        todo!()
-    }
 }

crates/revm/src/test_utils.rs

@@ -1,7 +1,9 @@
 use reth_primitives::{
     keccak256, Account, Address, BlockNumber, Bytecode, Bytes, StorageKey, B256, U256,
 };
-use reth_storage_api::{AccountReader, BlockHashReader, StateProvider, StateRootProvider};
+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;
@@ -76,6 +78,12 @@ impl StateRootProvider for StateProviderTest {
     }
 }
 
+impl StateProofProvider for StateProviderTest {
+    fn proof(&self, _address: Address, _slots: &[B256]) -> ProviderResult<AccountProof> {
+        unimplemented!("proof generation is not supported")
+    }
+}
+
 impl StateProvider for StateProviderTest {
     fn storage(
         &self,
@@ -88,8 +96,4 @@ impl StateProvider for StateProviderTest {
     fn bytecode_by_hash(&self, code_hash: B256) -> ProviderResult<Option<Bytecode>> {
         Ok(self.contracts.get(&code_hash).cloned())
     }
-
-    fn proof(&self, _address: Address, _keys: &[B256]) -> ProviderResult<AccountProof> {
-        unimplemented!("proof generation is not supported")
-    }
 }

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

@@ -31,6 +31,16 @@ impl<'a> reth_provider::StateRootProvider for StateProviderTraitObjWrapper<'a> {
     }
 }
 
+impl<'a> reth_provider::StateProofProvider for StateProviderTraitObjWrapper<'a> {
+    fn proof(
+        &self,
+        address: revm_primitives::Address,
+        slots: &[B256],
+    ) -> reth_errors::ProviderResult<reth_trie::AccountProof> {
+        self.0.proof(address, slots)
+    }
+}
+
 impl<'a> reth_provider::AccountReader for StateProviderTraitObjWrapper<'a> {
     fn basic_account(
         &self,
@@ -93,14 +103,6 @@ impl<'a> StateProvider for StateProviderTraitObjWrapper<'a> {
         self.0.bytecode_by_hash(code_hash)
     }
 
-    fn proof(
-        &self,
-        address: revm_primitives::Address,
-        keys: &[B256],
-    ) -> reth_errors::ProviderResult<reth_trie::AccountProof> {
-        self.0.proof(address, keys)
-    }
-
     fn storage(
         &self,
         account: revm_primitives::Address,

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

@@ -2,6 +2,7 @@ use crate::{
     AccountReader, BlockHashReader, ExecutionDataProvider, StateProvider, StateRootProvider,
 };
 use reth_primitives::{Account, Address, BlockNumber, Bytecode, B256};
+use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
 use reth_trie::{updates::TrieUpdates, AccountProof};
 use revm::db::BundleState;
@@ -80,6 +81,14 @@ 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> {
+        Err(ProviderError::StateRootNotAvailableForHistoricalBlock)
+    }
+}
+
 impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProvider for BundleStateProvider<SP, EDP> {
     fn storage(
         &self,
@@ -107,8 +116,4 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProvider for BundleStat
 
         self.state_provider.bytecode_by_hash(code_hash)
     }
-
-    fn proof(&self, _address: Address, _keys: &[B256]) -> ProviderResult<AccountProof> {
-        Err(ProviderError::StateRootNotAvailableForHistoricalBlock)
-    }
 }

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

@@ -13,6 +13,7 @@ use reth_primitives::{
     constants::EPOCH_SLOTS, Account, Address, BlockNumber, Bytecode, StaticFileSegment, StorageKey,
     StorageValue, B256,
 };
+use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
 use revm::db::BundleState;
@@ -271,6 +272,13 @@ 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> {
+        Err(ProviderError::StateRootNotAvailableForHistoricalBlock)
+    }
+}
+
 impl<'b, TX: DbTx> StateProvider for HistoricalStateProviderRef<'b, TX> {
     /// Get storage.
     fn storage(
@@ -306,11 +314,6 @@ impl<'b, TX: DbTx> StateProvider for HistoricalStateProviderRef<'b, TX> {
     fn bytecode_by_hash(&self, code_hash: B256) -> ProviderResult<Option<Bytecode>> {
         self.tx.get::<tables::Bytecodes>(code_hash).map_err(Into::into)
     }
-
-    /// Get account and storage proofs.
-    fn proof(&self, _address: Address, _keys: &[B256]) -> ProviderResult<AccountProof> {
-        Err(ProviderError::StateRootNotAvailableForHistoricalBlock)
-    }
 }
 
 /// State provider for a given block number.

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

@@ -10,6 +10,7 @@ use reth_db_api::{
 use reth_primitives::{
     Account, Address, BlockNumber, Bytecode, StaticFileSegment, StorageKey, StorageValue, B256,
 };
+use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
 use reth_trie::{proof::Proof, updates::TrieUpdates, AccountProof, HashedPostState};
 use revm::db::BundleState;
@@ -90,6 +91,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::new(self.tx)
+            .account_proof(address, slots)
+            .map_err(Into::<reth_db::DatabaseError>::into)?)
+    }
+}
+
 impl<'b, TX: DbTx> StateProvider for LatestStateProviderRef<'b, TX> {
     /// Get storage.
     fn storage(
@@ -110,12 +119,6 @@ impl<'b, TX: DbTx> StateProvider for LatestStateProviderRef<'b, TX> {
     fn bytecode_by_hash(&self, code_hash: B256) -> ProviderResult<Option<Bytecode>> {
         self.tx.get::<tables::Bytecodes>(code_hash).map_err(Into::into)
     }
-
-    fn proof(&self, address: Address, slots: &[B256]) -> ProviderResult<AccountProof> {
-        Ok(Proof::new(self.tx)
-            .account_proof(address, slots)
-            .map_err(Into::<reth_db::DatabaseError>::into)?)
-    }
 }
 
 /// State provider for the latest state.

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

@@ -30,10 +30,6 @@ macro_rules! delegate_provider_impls {
     ($target:ty $(where [$($generics:tt)*])?) => {
         $crate::providers::state::macros::delegate_impls_to_as_ref!(
             for $target =>
-            StateRootProvider $(where [$($generics)*])? {
-                fn state_root(&self, state: &revm::db::BundleState) -> 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)>;
-            }
             AccountReader $(where [$($generics)*])? {
                 fn basic_account(&self, address: reth_primitives::Address) -> reth_storage_errors::provider::ProviderResult<Option<reth_primitives::Account>>;
             }
@@ -43,9 +39,15 @@ macro_rules! delegate_provider_impls {
             }
             StateProvider $(where [$($generics)*])? {
                 fn storage(&self, account: reth_primitives::Address, storage_key: reth_primitives::StorageKey) -> reth_storage_errors::provider::ProviderResult<Option<reth_primitives::StorageValue>>;
-                fn proof(&self, address: reth_primitives::Address, keys: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult<reth_trie::AccountProof>;
                 fn bytecode_by_hash(&self, code_hash: reth_primitives::B256) -> reth_storage_errors::provider::ProviderResult<Option<reth_primitives::Bytecode>>;
             }
+            StateRootProvider $(where [$($generics)*])? {
+                fn state_root(&self, state: &revm::db::BundleState) -> 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)>;
+            }
+            StateProofProvider $(where [$($generics)*])? {
+                fn proof(&self, address: reth_primitives::Address, slots: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult<reth_trie::AccountProof>;
+            }
         );
     }
 }

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

@@ -15,6 +15,7 @@ use reth_primitives::{
     SealedHeader, StorageKey, StorageValue, TransactionMeta, TransactionSigned,
     TransactionSignedNoHash, TxHash, TxNumber, Withdrawal, Withdrawals, B256, U256,
 };
+use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
 use reth_trie::{updates::TrieUpdates, AccountProof};
 use revm::{
@@ -553,6 +554,12 @@ impl StateRootProvider for MockEthProvider {
     }
 }
 
+impl StateProofProvider for MockEthProvider {
+    fn proof(&self, address: Address, _slots: &[B256]) -> ProviderResult<AccountProof> {
+        Ok(AccountProof::new(address))
+    }
+}
+
 impl StateProvider for MockEthProvider {
     fn storage(
         &self,
@@ -574,10 +581,6 @@ impl StateProvider for MockEthProvider {
             }
         }))
     }
-
-    fn proof(&self, address: Address, _keys: &[B256]) -> ProviderResult<AccountProof> {
-        Ok(AccountProof::new(address))
-    }
 }
 
 impl EvmEnvProvider for MockEthProvider {

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

@@ -14,6 +14,7 @@ use reth_primitives::{
 };
 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::{
@@ -326,6 +327,12 @@ impl StateRootProvider for NoopProvider {
     }
 }
 
+impl StateProofProvider for NoopProvider {
+    fn proof(&self, address: Address, _slots: &[B256]) -> ProviderResult<AccountProof> {
+        Ok(AccountProof::new(address))
+    }
+}
+
 impl StateProvider for NoopProvider {
     fn storage(
         &self,
@@ -338,10 +345,6 @@ impl StateProvider for NoopProvider {
     fn bytecode_by_hash(&self, _code_hash: B256) -> ProviderResult<Option<Bytecode>> {
         Ok(None)
     }
-
-    fn proof(&self, address: Address, _keys: &[B256]) -> ProviderResult<AccountProof> {
-        Ok(AccountProof::new(address))
-    }
 }
 
 impl EvmEnvProvider for NoopProvider {

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

@@ -1,4 +1,4 @@
-use super::{AccountReader, BlockHashReader, BlockIdReader, StateRootProvider};
+use super::{AccountReader, BlockHashReader, BlockIdReader, StateProofProvider, StateRootProvider};
 use auto_impl::auto_impl;
 use reth_execution_types::ExecutionOutcome;
 use reth_primitives::{
@@ -6,14 +6,15 @@ use reth_primitives::{
     StorageValue, B256, KECCAK_EMPTY, U256,
 };
 use reth_storage_errors::provider::{ProviderError, ProviderResult};
-use reth_trie::AccountProof;
 
 /// Type alias of boxed [`StateProvider`].
 pub type StateProviderBox = Box<dyn StateProvider>;
 
 /// An abstraction for a type that provides state data.
 #[auto_impl(&, Arc, Box)]
-pub trait StateProvider: BlockHashReader + AccountReader + StateRootProvider + Send + Sync {
+pub trait StateProvider:
+    BlockHashReader + AccountReader + StateRootProvider + StateProofProvider + Send + Sync
+{
     /// Get storage of given account.
     fn storage(
         &self,
@@ -24,9 +25,6 @@ pub trait StateProvider: BlockHashReader + AccountReader + StateRootProvider + S
     /// Get account code by its hash
     fn bytecode_by_hash(&self, code_hash: B256) -> ProviderResult<Option<Bytecode>>;
 
-    /// Get account and storage proofs.
-    fn proof(&self, address: Address, keys: &[B256]) -> ProviderResult<AccountProof>;
-
     /// Get account code by its address.
     ///
     /// Returns `None` if the account doesn't exist or account is not a contract

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

@@ -1,6 +1,6 @@
-use reth_primitives::B256;
+use reth_primitives::{Address, B256};
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::updates::TrieUpdates;
+use reth_trie::{updates::TrieUpdates, AccountProof};
 use revm::db::BundleState;
 
 /// A type that can compute the state root of a given post state.
@@ -22,3 +22,10 @@ pub trait StateRootProvider: Send + Sync {
         bundle_state: &BundleState,
     ) -> ProviderResult<(B256, TrieUpdates)>;
 }
+
+/// 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>;
+}