feat(trie): integrate TrieInput into StateProofProvider methods (#10929)

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<AccountProof> {
@@ -954,7 +954,7 @@ mod tests {
 
         fn witness(
             &self,
-            _overlay: HashedPostState,
+            _input: TrieInput,
             _target: HashedPostState,
         ) -> ProviderResult<HashMap<B256, Bytes>> {
             Ok(HashMap::default())

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<AccountProof> {
-        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<HashMap<B256, Bytes>> {
-        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)
     }
 }
 

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()

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<AccountProof> {
@@ -112,7 +112,7 @@ impl StateProofProvider for StateProviderTest {
 
     fn witness(
         &self,
-        _overlay: HashedPostState,
+        _input: TrieInput,
         _target: HashedPostState,
     ) -> ProviderResult<HashMap<B256, Bytes>> {
         unimplemented!("witness generation is not supported")

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<reth_trie::AccountProof> {
-        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<std::collections::HashMap<B256, alloy_primitives::Bytes>> {
-        self.0.witness(overlay, target)
+        self.0.witness(input, target)
     }
 }
 

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,

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

@@ -124,25 +124,23 @@ impl<SP: StateProvider, EDP: ExecutionDataProvider> StateProofProvider
 {
     fn proof(
         &self,
-        hashed_state: HashedPostState,
+        mut input: TrieInput,
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {
         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<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)
+        input.prepend(HashedPostState::from_bundle_state(&bundle_state.state));
+        self.state_provider.witness(input, target)
     }
 }
 

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<AccountProof> {
-        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::<ProviderError>::into)
     }
 
     fn witness(
         &self,
-        overlay: HashedPostState,
+        mut input: TrieInput,
         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)
+        input.prepend(self.revert_state()?);
+        TrieWitness::overlay_witness(self.tx, input, target).map_err(Into::<ProviderError>::into)
     }
 }
 

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<AccountProof> {
-        Proof::overlay_account_proof(self.tx, hashed_state, address, slots)
+        Proof::overlay_account_proof(self.tx, input, address, slots)
             .map_err(Into::<ProviderError>::into)
     }
 
     fn witness(
         &self,
-        overlay: HashedPostState,
+        input: TrieInput,
         target: HashedPostState,
     ) -> ProviderResult<HashMap<B256, Bytes>> {
-        TrieWitness::overlay_witness(self.tx, overlay, target).map_err(Into::<ProviderError>::into)
+        TrieWitness::overlay_witness(self.tx, input, target).map_err(Into::<ProviderError>::into)
     }
 }
 

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<reth_primitives::B256>;
             }
             StateProofProvider $(where [$($generics)*])? {
-                fn 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>>;
+                fn proof(&self, input: reth_trie::TrieInput, address: reth_primitives::Address, slots: &[reth_primitives::B256]) -> reth_storage_errors::provider::ProviderResult<reth_trie::AccountProof>;
+                fn witness(&self, input: reth_trie::TrieInput, 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

@@ -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<AccountProof> {
@@ -626,7 +626,7 @@ impl StateProofProvider for MockEthProvider {
 
     fn witness(
         &self,
-        _overlay: HashedPostState,
+        _input: TrieInput,
         _target: HashedPostState,
     ) -> ProviderResult<HashMap<B256, Bytes>> {
         Ok(HashMap::default())

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<AccountProof> {
@@ -364,7 +364,7 @@ impl StateProofProvider for NoopProvider {
 
     fn witness(
         &self,
-        _overlay: HashedPostState,
+        _input: TrieInput,
         _target: HashedPostState,
     ) -> ProviderResult<HashMap<B256, Bytes>> {
         Ok(HashMap::default())

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<AccountProof>;
@@ -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<HashMap<B256, Bytes>>;
 }

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<AccountProof, StateProofError>;
@@ -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<AccountProof, StateProofError> {
-        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)
     }
 }

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<HashMap<B256, Bytes>, 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<HashMap<B256, Bytes>, 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)
     }
 }

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.

crates/trie/trie/src/proof.rs

@@ -42,6 +42,16 @@ impl<T, H> Proof<T, H> {
         }
     }
 
+    /// Set the trie cursor factory.
+    pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> Proof<TF, H> {
+        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<HF>(self, hashed_cursor_factory: HF) -> Proof<T, HF> {
         Proof {

crates/trie/trie/src/witness.rs

@@ -36,6 +36,16 @@ impl<T, H> TrieWitness<T, H> {
         }
     }
 
+    /// Set the trie cursor factory.
+    pub fn with_trie_cursor_factory<TF>(self, trie_cursor_factory: TF) -> TrieWitness<TF, H> {
+        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<HF>(self, hashed_cursor_factory: HF) -> TrieWitness<T, HF> {
         TrieWitness {