feat(trie): expose trie witness on proof provider (#9915)

crates/chain-state/src/in_memory.rs

@@ -667,7 +667,9 @@ mod tests {
     use crate::test_utils::TestBlockBuilder;
     use rand::Rng;
     use reth_errors::ProviderResult;
-    use reth_primitives::{Account, BlockNumber, Bytecode, Receipt, StorageKey, StorageValue};
+    use reth_primitives::{
+        Account, BlockNumber, Bytecode, Bytes, Receipt, StorageKey, StorageValue,
+    };
     use reth_storage_api::{
         AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
     };
@@ -762,6 +764,14 @@ mod tests {
         ) -> ProviderResult<AccountProof> {
             Ok(AccountProof::new(Address::random()))
         }
+
+        fn witness(
+            &self,
+            _overlay: HashedPostState,
+            _target: HashedPostState,
+        ) -> ProviderResult<HashMap<B256, Bytes>> {
+            Ok(HashMap::default())
+        }
     }
 
     #[test]

crates/chain-state/src/memory_overlay.rs

@@ -1,6 +1,10 @@
+use std::collections::HashMap;
+
 use super::ExecutedBlock;
 use reth_errors::ProviderResult;
-use reth_primitives::{Account, Address, BlockNumber, Bytecode, StorageKey, StorageValue, B256};
+use reth_primitives::{
+    Account, Address, BlockNumber, Bytecode, Bytes, StorageKey, StorageValue, B256,
+};
 use reth_storage_api::{
     AccountReader, BlockHashReader, StateProofProvider, StateProvider, StateRootProvider,
 };
@@ -111,6 +115,17 @@ impl StateProofProvider for MemoryOverlayStateProvider {
         state.extend(hashed_state);
         self.historical.hashed_proof(state, address, slots)
     }
+
+    // TODO: Currently this does not reuse available in-memory trie nodes.
+    fn witness(
+        &self,
+        overlay: HashedPostState,
+        target: HashedPostState,
+    ) -> ProviderResult<HashMap<B256, Bytes>> {
+        let mut state = self.hashed_post_state.clone();
+        state.extend(overlay);
+        self.historical.witness(state, target)
+    }
 }
 
 impl StateProvider for MemoryOverlayStateProvider {

crates/evm/execution-errors/src/trie.rs

@@ -7,6 +7,34 @@ use thiserror_no_std::Error;
 
 /// State root errors.
 #[derive(Error, Debug, PartialEq, Eq, Clone)]
+pub enum StateRootError {
+    /// Internal database error.
+    #[error(transparent)]
+    Database(#[from] DatabaseError),
+    /// Storage root error.
+    #[error(transparent)]
+    StorageRootError(#[from] StorageRootError),
+}
+
+impl From<StateRootError> for DatabaseError {
+    fn from(err: StateRootError) -> Self {
+        match err {
+            StateRootError::Database(err) |
+            StateRootError::StorageRootError(StorageRootError::Database(err)) => err,
+        }
+    }
+}
+
+/// Storage root error.
+#[derive(Error, PartialEq, Eq, Clone, Debug)]
+pub enum StorageRootError {
+    /// Internal database error.
+    #[error(transparent)]
+    Database(#[from] DatabaseError),
+}
+
+/// State proof errors.
+#[derive(Error, Debug, PartialEq, Eq, Clone)]
 pub enum StateProofError {
     /// Internal database error.
     #[error(transparent)]
@@ -45,30 +73,8 @@ pub enum TrieWitnessError {
     MissingTargetNode(Nibbles),
 }
 
-/// State root errors.
-#[derive(Error, Debug, PartialEq, Eq, Clone)]
-pub enum StateRootError {
-    /// Internal database error.
-    #[error(transparent)]
-    Database(#[from] DatabaseError),
-    /// Storage root error.
-    #[error(transparent)]
-    StorageRootError(#[from] StorageRootError),
-}
-
-impl From<StateRootError> for DatabaseError {
-    fn from(err: StateRootError) -> Self {
-        match err {
-            StateRootError::Database(err) |
-            StateRootError::StorageRootError(StorageRootError::Database(err)) => err,
-        }
+impl From<TrieWitnessError> for ProviderError {
+    fn from(value: TrieWitnessError) -> Self {
+        Self::TrieWitnessError(value.to_string())
     }
 }
-
-/// Storage root error.
-#[derive(Error, PartialEq, Eq, Clone, Debug)]
-pub enum StorageRootError {
-    /// Internal database error.
-    #[error(transparent)]
-    Database(#[from] DatabaseError),
-}

crates/revm/src/test_utils.rs

@@ -89,6 +89,14 @@ impl StateProofProvider for StateProviderTest {
     ) -> ProviderResult<AccountProof> {
         unimplemented!("proof generation is not supported")
     }
+
+    fn witness(
+        &self,
+        _overlay: HashedPostState,
+        _target: HashedPostState,
+    ) -> ProviderResult<HashMap<B256, Bytes>> {
+        unimplemented!("witness generation is not supported")
+    }
 }
 
 impl StateProvider for StateProviderTest {

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

@@ -40,6 +40,14 @@ impl<'a> reth_provider::StateProofProvider for StateProviderTraitObjWrapper<'a>
     ) -> reth_errors::ProviderResult<reth_trie::AccountProof> {
         self.0.hashed_proof(hashed_state, address, slots)
     }
+
+    fn witness(
+        &self,
+        overlay: reth_trie::HashedPostState,
+        target: reth_trie::HashedPostState,
+    ) -> reth_errors::ProviderResult<std::collections::HashMap<B256, reth_primitives::Bytes>> {
+        self.0.witness(overlay, target)
+    }
 }
 
 impl<'a> reth_provider::AccountReader for StateProviderTraitObjWrapper<'a> {

crates/storage/errors/src/provider.rs

@@ -30,6 +30,9 @@ pub enum ProviderError {
     /// Nippy jar error.
     #[error("nippy jar error: {0}")]
     NippyJar(String),
+    /// Trie witness error.
+    #[error("trie witness error: {0}")]
+    TrieWitnessError(String),
     /// Error when recovering the sender for a transaction
     #[error("failed to recover sender for transaction")]
     SenderRecoveryError,

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

@@ -1,7 +1,9 @@
+use std::collections::HashMap;
+
 use crate::{
     AccountReader, BlockHashReader, ExecutionDataProvider, StateProvider, StateRootProvider,
 };
-use reth_primitives::{Account, Address, BlockNumber, Bytecode, B256};
+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};
@@ -112,6 +114,17 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProofProvider
         state.extend(hashed_state);
         self.state_provider.hashed_proof(state, address, slots)
     }
+
+    fn witness(
+        &self,
+        overlay: HashedPostState,
+        target: HashedPostState,
+    ) -> ProviderResult<HashMap<B256, Bytes>> {
+        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)
+    }
 }
 
 impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProvider for BundleStateProvider<SP, EDP> {

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

@@ -10,14 +10,17 @@ use reth_db_api::{
     transaction::DbTx,
 };
 use reth_primitives::{
-    constants::EPOCH_SLOTS, Account, Address, BlockNumber, Bytecode, StaticFileSegment, StorageKey,
-    StorageValue, B256,
+    constants::EPOCH_SLOTS, Account, Address, BlockNumber, Bytecode, Bytes, StaticFileSegment,
+    StorageKey, StorageValue, B256,
 };
 use reth_storage_api::StateProofProvider;
 use reth_storage_errors::provider::ProviderResult;
-use reth_trie::{proof::Proof, updates::TrieUpdates, AccountProof, HashedPostState, StateRoot};
-use reth_trie_db::{DatabaseProof, DatabaseStateRoot};
-use std::fmt::Debug;
+use reth_trie::{
+    proof::Proof, updates::TrieUpdates, witness::TrieWitness, AccountProof, HashedPostState,
+    StateRoot,
+};
+use reth_trie_db::{DatabaseProof, DatabaseStateRoot, DatabaseTrieWitness};
+use std::{collections::HashMap, fmt::Debug};
 
 /// State provider for a given block number which takes a tx reference.
 ///
@@ -288,6 +291,17 @@ impl<'b, TX: DbTx> StateProofProvider for HistoricalStateProviderRef<'b, TX> {
         Proof::overlay_account_proof(self.tx, revert_state, address, slots)
             .map_err(Into::<ProviderError>::into)
     }
+
+    fn witness(
+        &self,
+        overlay: HashedPostState,
+        target: HashedPostState,
+    ) -> ProviderResult<HashMap<B256, Bytes>> {
+        let mut revert_state = self.revert_state()?;
+        revert_state.extend(overlay);
+        TrieWitness::overlay_witness(self.tx, revert_state, target)
+            .map_err(Into::<ProviderError>::into)
+    }
 }
 
 impl<'b, TX: DbTx> StateProvider for HistoricalStateProviderRef<'b, TX> {

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

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use crate::{
     providers::{state::macros::delegate_provider_impls, StaticFileProvider},
     AccountReader, BlockHashReader, StateProvider, StateRootProvider,
@@ -8,12 +10,16 @@ use reth_db_api::{
     transaction::DbTx,
 };
 use reth_primitives::{
-    Account, Address, BlockNumber, Bytecode, StaticFileSegment, StorageKey, StorageValue, B256,
+    Account, Address, BlockNumber, Bytecode, Bytes, 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, StateRoot};
-use reth_trie_db::{DatabaseProof, DatabaseStateRoot};
+use reth_trie::{
+    proof::Proof, updates::TrieUpdates, witness::TrieWitness, AccountProof, HashedPostState,
+    StateRoot,
+};
+use reth_trie_db::{DatabaseProof, DatabaseStateRoot, DatabaseTrieWitness};
 
 /// State provider over latest state that takes tx reference.
 #[derive(Debug)]
@@ -99,6 +105,14 @@ impl<'b, TX: DbTx> StateProofProvider for LatestStateProviderRef<'b, TX> {
         Proof::overlay_account_proof(self.tx, hashed_state, address, slots)
             .map_err(Into::<ProviderError>::into)
     }
+
+    fn witness(
+        &self,
+        overlay: HashedPostState,
+        target: HashedPostState,
+    ) -> ProviderResult<HashMap<B256, Bytes>> {
+        TrieWitness::overlay_witness(self.tx, overlay, target).map_err(Into::<ProviderError>::into)
+    }
 }
 
 impl<'b, TX: DbTx> StateProvider for LatestStateProviderRef<'b, TX> {

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

@@ -50,6 +50,7 @@ macro_rules! delegate_provider_impls {
             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>;
+                fn witness(&self, state: reth_trie::HashedPostState, target: reth_trie::HashedPostState) -> reth_storage_errors::provider::ProviderResult<std::collections::HashMap<reth_primitives::B256, reth_primitives::Bytes>>;
             }
         );
     }

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

@@ -560,6 +560,14 @@ impl StateProofProvider for MockEthProvider {
     ) -> ProviderResult<AccountProof> {
         Ok(AccountProof::new(address))
     }
+
+    fn witness(
+        &self,
+        _overlay: HashedPostState,
+        _target: HashedPostState,
+    ) -> ProviderResult<HashMap<B256, Bytes>> {
+        Ok(HashMap::default())
+    }
 }
 
 impl StateProvider for MockEthProvider {

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

@@ -1,4 +1,5 @@
 use std::{
+    collections::HashMap,
     ops::{RangeBounds, RangeInclusive},
     sync::Arc,
 };
@@ -9,9 +10,9 @@ use reth_db_api::models::{AccountBeforeTx, StoredBlockBodyIndices};
 use reth_evm::ConfigureEvmEnv;
 use reth_primitives::{
     Account, Address, Block, BlockHash, BlockHashOrNumber, BlockId, BlockNumber, BlockWithSenders,
-    Bytecode, Header, Receipt, SealedBlock, SealedBlockWithSenders, SealedHeader, StorageKey,
-    StorageValue, TransactionMeta, TransactionSigned, TransactionSignedNoHash, TxHash, TxNumber,
-    Withdrawal, Withdrawals, B256, U256,
+    Bytecode, Bytes, Header, Receipt, SealedBlock, SealedBlockWithSenders, SealedHeader,
+    StorageKey, StorageValue, TransactionMeta, TransactionSigned, TransactionSignedNoHash, TxHash,
+    TxNumber, Withdrawal, Withdrawals, B256, U256,
 };
 use reth_prune_types::{PruneCheckpoint, PruneSegment};
 use reth_stages_types::{StageCheckpoint, StageId};
@@ -333,6 +334,14 @@ impl StateProofProvider for NoopProvider {
     ) -> ProviderResult<AccountProof> {
         Ok(AccountProof::new(address))
     }
+
+    fn witness(
+        &self,
+        _overlay: HashedPostState,
+        _target: HashedPostState,
+    ) -> ProviderResult<HashMap<B256, Bytes>> {
+        Ok(HashMap::default())
+    }
 }
 
 impl StateProvider for NoopProvider {

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

@@ -1,7 +1,8 @@
-use reth_primitives::{Address, B256};
+use reth_primitives::{Address, Bytes, B256};
 use reth_storage_errors::provider::ProviderResult;
 use reth_trie::{updates::TrieUpdates, AccountProof, HashedPostState};
 use revm::db::BundleState;
+use std::collections::HashMap;
 
 /// A type that can compute the state root of a given post state.
 #[auto_impl::auto_impl(&, Box, Arc)]
@@ -60,4 +61,11 @@ pub trait StateProofProvider: Send + Sync {
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof>;
+
+    /// Get trie witness for provided state.
+    fn witness(
+        &self,
+        overlay: HashedPostState,
+        target: HashedPostState,
+    ) -> ProviderResult<HashMap<B256, Bytes>>;
 }