feat: add execute_with_witness (#11074)

crates/ethereum/evm/src/execute.rs

@@ -26,18 +26,14 @@ use reth_primitives::{
 use reth_prune_types::PruneModes;
 use reth_revm::{
     batch::BlockBatchRecord,
-    db::{
-        states::{bundle_state::BundleRetention, StorageSlot},
-        BundleAccount, State,
-    },
+    db::{states::bundle_state::BundleRetention, State},
     state_change::post_block_balance_increments,
     Evm,
 };
 use revm_primitives::{
     db::{Database, DatabaseCommit},
     BlockEnv, CfgEnvWithHandlerCfg, EnvWithHandlerCfg, ResultAndState,
 };
-use std::collections::hash_map::Entry;
 
 /// Provides executors to execute regular ethereum blocks
 #[derive(Debug, Clone)]
@@ -378,90 +374,25 @@ where
 
         Ok(BlockExecutionOutput { state: self.state.take_bundle(), receipts, requests, gas_used })
     }
-}
-
-/// An executor that retains all cache state from execution in its bundle state.
-#[derive(Debug)]
-pub struct BlockAccessListExecutor<EvmConfig, DB> {
-    /// The executor used to execute single blocks
-    ///
-    /// All state changes are committed to the [State].
-    executor: EthBlockExecutor<EvmConfig, DB>,
-}
 
-impl<EvmConfig, DB> Executor<DB> for BlockAccessListExecutor<EvmConfig, DB>
-where
-    EvmConfig: ConfigureEvm<Header = Header>,
-    DB: Database<Error: Into<ProviderError> + Display>,
-{
-    type Input<'a> = BlockExecutionInput<'a, BlockWithSenders>;
-    type Output = BlockExecutionOutput<Receipt>;
-    type Error = BlockExecutionError;
-
-    /// Executes the block and commits the changes to the internal state.
-    ///
-    /// Returns the receipts of the transactions in the block.
-    ///
-    /// This also returns the accounts from the internal state cache in the bundle state, allowing
-    /// access to not only the state that changed during execution, but also the state accessed
-    /// during execution.
-    ///
-    /// Returns an error if the block could not be executed or failed verification.
-    fn execute(mut self, input: Self::Input<'_>) -> Result<Self::Output, Self::Error> {
+    fn execute_with_state_witness<F>(
+        mut self,
+        input: Self::Input<'_>,
+        mut witness: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: FnMut(&State<DB>),
+    {
         let BlockExecutionInput { block, total_difficulty } = input;
         let EthExecuteOutput { receipts, requests, gas_used } =
-            self.executor.execute_without_verification(block, total_difficulty)?;
+            self.execute_without_verification(block, total_difficulty)?;
 
         // NOTE: we need to merge keep the reverts for the bundle retention
-        self.executor.state.merge_transitions(BundleRetention::Reverts);
-
-        // now, ensure each account from the state is included in the bundle state
-        let mut bundle_state = self.executor.state.take_bundle();
-        for (address, account) in self.executor.state.cache.accounts {
-            // convert all slots, insert all slots
-            let account_info = account.account_info();
-            let account_storage = account.account.map(|a| a.storage).unwrap_or_default();
-
-            match bundle_state.state.entry(address) {
-                Entry::Vacant(entry) => {
-                    // we have to add the entire account here
-                    let extracted_storage = account_storage
-                        .into_iter()
-                        .map(|(k, v)| {
-                            (k, StorageSlot { previous_or_original_value: v, present_value: v })
-                        })
-                        .collect();
-
-                    let bundle_account = BundleAccount {
-                        info: account_info.clone(),
-                        original_info: account_info,
-                        storage: extracted_storage,
-                        status: account.status,
-                    };
-                    entry.insert(bundle_account);
-                }
-                Entry::Occupied(mut entry) => {
-                    // only add slots that are unchanged
-                    let current_account = entry.get_mut();
-
-                    // iterate over all storage slots, checking keys that are not in the bundle
-                    // state
-                    for (k, v) in account_storage {
-                        if let Entry::Vacant(storage_entry) = current_account.storage.entry(k) {
-                            storage_entry.insert(StorageSlot {
-                                previous_or_original_value: v,
-                                present_value: v,
-                            });
-                        }
-                    }
-                }
-            }
-        }
-
-        Ok(BlockExecutionOutput { state: bundle_state, receipts, requests, gas_used })
+        self.state.merge_transitions(BundleRetention::Reverts);
+        witness(&self.state);
+        Ok(BlockExecutionOutput { state: self.state.take_bundle(), receipts, requests, gas_used })
     }
 }
-
 /// An executor for a batch of blocks.
 ///
 /// State changes are tracked until the executor is finalized.

crates/evm/src/either.rs

@@ -12,6 +12,7 @@ use revm_primitives::db::Database;
 
 // re-export Either
 pub use futures_util::future::Either;
+use revm::State;
 
 impl<A, B> BlockExecutorProvider for Either<A, B>
 where
@@ -71,6 +72,20 @@ where
             Self::Right(b) => b.execute(input),
         }
     }
+
+    fn execute_with_state_witness<F>(
+        self,
+        input: Self::Input<'_>,
+        witness: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: FnMut(&State<DB>),
+    {
+        match self {
+            Self::Left(a) => a.execute_with_state_witness(input, witness),
+            Self::Right(b) => b.execute_with_state_witness(input, witness),
+        }
+    }
 }
 
 impl<A, B, DB> BatchExecutor<DB> for Either<A, B>

crates/evm/src/execute.rs

@@ -6,9 +6,9 @@ pub use reth_execution_types::{BlockExecutionInput, BlockExecutionOutput, Execut
 pub use reth_storage_errors::provider::ProviderError;
 
 use core::fmt::Display;
-
 use reth_primitives::{BlockNumber, BlockWithSenders, Receipt};
 use reth_prune_types::PruneModes;
+use revm::State;
 use revm_primitives::db::Database;
 
 /// A general purpose executor trait that executes an input (e.g. block) and produces an output
@@ -32,6 +32,16 @@ pub trait Executor<DB> {
     /// # Returns
     /// The output of the block execution.
     fn execute(self, input: Self::Input<'_>) -> Result<Self::Output, Self::Error>;
+
+    /// Executes the EVM with the given input and accepts a witness closure that is invoked with the
+    /// EVM state after execution.
+    fn execute_with_state_witness<F>(
+        self,
+        input: Self::Input<'_>,
+        witness: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: FnMut(&State<DB>);
 }
 
 /// A general purpose executor that can execute multiple inputs in sequence, validate the outputs,
@@ -178,6 +188,17 @@ mod tests {
         fn execute(self, _input: Self::Input<'_>) -> Result<Self::Output, Self::Error> {
             Err(BlockExecutionError::msg("execution unavailable for tests"))
         }
+
+        fn execute_with_state_witness<F>(
+            self,
+            _: Self::Input<'_>,
+            _: F,
+        ) -> Result<Self::Output, Self::Error>
+        where
+            F: FnMut(&State<DB>),
+        {
+            Err(BlockExecutionError::msg("execution unavailable for tests"))
+        }
     }
 
     impl<DB> BatchExecutor<DB> for TestExecutor<DB> {

crates/evm/src/noop.rs

@@ -1,12 +1,12 @@
 //! A no operation block executor implementation.
 
 use core::fmt::Display;
-
 use reth_execution_errors::BlockExecutionError;
 use reth_execution_types::{BlockExecutionInput, BlockExecutionOutput, ExecutionOutcome};
 use reth_primitives::{BlockNumber, BlockWithSenders, Receipt};
 use reth_prune_types::PruneModes;
 use reth_storage_errors::provider::ProviderError;
+use revm::State;
 use revm_primitives::db::Database;
 
 use crate::execute::{BatchExecutor, BlockExecutorProvider, Executor};
@@ -46,6 +46,17 @@ impl<DB> Executor<DB> for NoopBlockExecutorProvider {
     fn execute(self, _: Self::Input<'_>) -> Result<Self::Output, Self::Error> {
         Err(BlockExecutionError::msg(UNAVAILABLE_FOR_NOOP))
     }
+
+    fn execute_with_state_witness<F>(
+        self,
+        _: Self::Input<'_>,
+        _: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: FnMut(&State<DB>),
+    {
+        Err(BlockExecutionError::msg(UNAVAILABLE_FOR_NOOP))
+    }
 }
 
 impl<DB> BatchExecutor<DB> for NoopBlockExecutorProvider {

crates/evm/src/test_utils.rs

@@ -9,6 +9,7 @@ use reth_execution_types::ExecutionOutcome;
 use reth_primitives::{BlockNumber, BlockWithSenders, Receipt};
 use reth_prune_types::PruneModes;
 use reth_storage_errors::provider::ProviderError;
+use revm::State;
 use revm_primitives::db::Database;
 use std::{fmt::Display, sync::Arc};
 
@@ -60,6 +61,17 @@ impl<DB> Executor<DB> for MockExecutorProvider {
             gas_used: 0,
         })
     }
+
+    fn execute_with_state_witness<F>(
+        self,
+        _: Self::Input<'_>,
+        _: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: FnMut(&State<DB>),
+    {
+        unimplemented!()
+    }
 }
 
 impl<DB> BatchExecutor<DB> for MockExecutorProvider {

crates/optimism/evm/src/execute.rs

@@ -19,19 +19,14 @@ use reth_optimism_forks::OptimismHardfork;
 use reth_primitives::{BlockWithSenders, Header, Receipt, Receipts, TxType};
 use reth_prune_types::PruneModes;
 use reth_revm::{
-    batch::BlockBatchRecord,
-    db::{
-        states::{bundle_state::BundleRetention, StorageSlot},
-        BundleAccount,
-    },
-    state_change::post_block_balance_increments,
-    Evm, State,
+    batch::BlockBatchRecord, db::states::bundle_state::BundleRetention,
+    state_change::post_block_balance_increments, Evm, State,
 };
 use revm_primitives::{
     db::{Database, DatabaseCommit},
     BlockEnv, CfgEnvWithHandlerCfg, EnvWithHandlerCfg, ResultAndState,
 };
-use std::{collections::hash_map::Entry, fmt::Display, sync::Arc};
+use std::{fmt::Display, sync::Arc};
 use tracing::trace;
 
 /// Provides executors to execute regular optimism blocks
@@ -364,87 +359,28 @@ where
             gas_used,
         })
     }
-}
-
-/// An executor that retains all cache state from execution in its bundle state.
-#[derive(Debug)]
-pub struct OpBlockAccessListExecutor<EvmConfig, DB> {
-    /// The executor used to execute single blocks
-    ///
-    /// All state changes are committed to the [State].
-    executor: OpBlockExecutor<EvmConfig, DB>,
-}
-
-impl<EvmConfig, DB> Executor<DB> for OpBlockAccessListExecutor<EvmConfig, DB>
-where
-    EvmConfig: ConfigureEvm<Header = Header>,
-    DB: Database<Error: Into<ProviderError> + Display>,
-{
-    type Input<'a> = BlockExecutionInput<'a, BlockWithSenders>;
-    type Output = BlockExecutionOutput<Receipt>;
-    type Error = BlockExecutionError;
 
-    /// Executes the block and commits the changes to the internal state.
-    ///
-    /// Returns the receipts of the transactions in the block.
-    ///
-    /// This also returns the accounts from the internal state cache in the bundle state, allowing
-    /// access to not only the state that changed during execution, but also the state accessed
-    /// during execution.
-    ///
-    /// Returns an error if the block could not be executed or failed verification.
-    fn execute(mut self, input: Self::Input<'_>) -> Result<Self::Output, Self::Error> {
+    fn execute_with_state_witness<F>(
+        mut self,
+        input: Self::Input<'_>,
+        mut witness: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: FnMut(&State<DB>),
+    {
         let BlockExecutionInput { block, total_difficulty } = input;
-        let (receipts, gas_used) =
-            self.executor.execute_without_verification(block, total_difficulty)?;
+        let (receipts, gas_used) = self.execute_without_verification(block, total_difficulty)?;
 
         // NOTE: we need to merge keep the reverts for the bundle retention
-        self.executor.state.merge_transitions(BundleRetention::Reverts);
-
-        // now, ensure each account from the state is included in the bundle state
-        let mut bundle_state = self.executor.state.take_bundle();
-        for (address, account) in self.executor.state.cache.accounts {
-            // convert all slots, insert all slots
-            let account_info = account.account_info();
-            let account_storage = account.account.map(|a| a.storage).unwrap_or_default();
-
-            match bundle_state.state.entry(address) {
-                Entry::Vacant(entry) => {
-                    // we have to add the entire account here
-                    let extracted_storage = account_storage
-                        .into_iter()
-                        .map(|(k, v)| {
-                            (k, StorageSlot { previous_or_original_value: v, present_value: v })
-                        })
-                        .collect();
-
-                    let bundle_account = BundleAccount {
-                        info: account_info.clone(),
-                        original_info: account_info,
-                        storage: extracted_storage,
-                        status: account.status,
-                    };
-                    entry.insert(bundle_account);
-                }
-                Entry::Occupied(mut entry) => {
-                    // only add slots that are unchanged
-                    let current_account = entry.get_mut();
-
-                    // iterate over all storage slots, checking keys that are not in the bundle
-                    // state
-                    for (k, v) in account_storage {
-                        if let Entry::Vacant(storage_entry) = current_account.storage.entry(k) {
-                            storage_entry.insert(StorageSlot {
-                                previous_or_original_value: v,
-                                present_value: v,
-                            });
-                        }
-                    }
-                }
-            }
-        }
+        self.state.merge_transitions(BundleRetention::Reverts);
+        witness(&self.state);
 
-        Ok(BlockExecutionOutput { state: bundle_state, receipts, requests: vec![], gas_used })
+        Ok(BlockExecutionOutput {
+            state: self.state.take_bundle(),
+            receipts,
+            requests: vec![],
+            gas_used,
+        })
     }
 }