feat(perf): integrate OnStateHook in executor

crates/ethereum/evm/src/execute.rs

@@ -14,7 +14,7 @@ use reth_evm::{
         BatchExecutor, BlockExecutionError, BlockExecutionInput, BlockExecutionOutput,
         BlockExecutorProvider, BlockValidationError, Executor, ProviderError,
     },
-    system_calls::SystemCaller,
+    system_calls::{NoopHook, OnStateHook, SystemCaller},
     ConfigureEvm,
 };
 use reth_execution_types::ExecutionOutcome;
@@ -130,16 +130,19 @@ where
     ///
     /// It does __not__ apply post-execution changes that do not require an [EVM](Evm), for that see
     /// [`EthBlockExecutor::post_execution`].
-    fn execute_state_transitions<Ext, DB>(
+    fn execute_state_transitions<Ext, DB, F>(
         &self,
         block: &BlockWithSenders,
         mut evm: Evm<'_, Ext, &mut State<DB>>,
+        state_hook: Option<F>,
     ) -> Result<EthExecuteOutput, BlockExecutionError>
     where
         DB: Database,
         DB::Error: Into<ProviderError> + Display,
+        F: OnStateHook,
     {
-        let mut system_caller = SystemCaller::new(&self.evm_config, &self.chain_spec);
+        let mut system_caller =
+            SystemCaller::new(&self.evm_config, &self.chain_spec).with_state_hook(state_hook);
 
         system_caller.apply_pre_execution_changes(block, &mut evm)?;
 
@@ -169,7 +172,8 @@ where
                     error: Box::new(new_err),
                 }
             })?;
-            evm.db_mut().commit(state);
+            evm.db_mut().commit(state.clone());
+            system_caller.on_state(&ResultAndState { result: result.clone(), state });
 
             // append gas used
             cumulative_gas_used += result.gas_used();
@@ -260,25 +264,37 @@ where
         EnvWithHandlerCfg::new_with_cfg_env(cfg, block_env, Default::default())
     }
 
+    fn execute_without_verification(
+        &mut self,
+        block: &BlockWithSenders,
+        total_difficulty: U256,
+    ) -> Result<EthExecuteOutput, BlockExecutionError> {
+        self.execute_without_verification_with_state_hook(block, total_difficulty, None::<NoopHook>)
+    }
+
     /// Execute a single block and apply the state changes to the internal state.
     ///
     /// Returns the receipts of the transactions in the block, the total gas used and the list of
     /// EIP-7685 [requests](Request).
     ///
     /// Returns an error if execution fails.
-    fn execute_without_verification(
+    fn execute_without_verification_with_state_hook<F>(
         &mut self,
         block: &BlockWithSenders,
         total_difficulty: U256,
-    ) -> Result<EthExecuteOutput, BlockExecutionError> {
+        state_hook: Option<F>,
+    ) -> Result<EthExecuteOutput, BlockExecutionError>
+    where
+        F: OnStateHook,
+    {
         // 1. prepare state on new block
         self.on_new_block(&block.header);
 
         // 2. configure the evm and execute
         let env = self.evm_env_for_block(&block.header, total_difficulty);
         let output = {
             let evm = self.executor.evm_config.evm_with_env(&mut self.state, env);
-            self.executor.execute_state_transitions(block, evm)
+            self.executor.execute_state_transitions(block, evm, state_hook)
         }?;
 
         // 3. apply post execution changes
@@ -368,6 +384,27 @@ where
         witness(&self.state);
         Ok(BlockExecutionOutput { state: self.state.take_bundle(), receipts, requests, gas_used })
     }
+
+    fn execute_with_state_hook<F>(
+        mut self,
+        input: Self::Input<'_>,
+        state_hook: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook,
+    {
+        let BlockExecutionInput { block, total_difficulty } = input;
+        let EthExecuteOutput { receipts, requests, gas_used } = self
+            .execute_without_verification_with_state_hook(
+                block,
+                total_difficulty,
+                Some(state_hook),
+            )?;
+
+        // NOTE: we need to merge keep the reverts for the bundle retention
+        self.state.merge_transitions(BundleRetention::Reverts);
+        Ok(BlockExecutionOutput { state: self.state.take_bundle(), receipts, requests, gas_used })
+    }
 }
 /// An executor for a batch of blocks.
 ///

crates/evm/src/either.rs

@@ -2,7 +2,10 @@
 
 use core::fmt::Display;
 
-use crate::execute::{BatchExecutor, BlockExecutorProvider, Executor};
+use crate::{
+    execute::{BatchExecutor, BlockExecutorProvider, Executor},
+    system_calls::OnStateHook,
+};
 use alloy_primitives::BlockNumber;
 use reth_execution_errors::BlockExecutionError;
 use reth_execution_types::{BlockExecutionInput, BlockExecutionOutput, ExecutionOutcome};
@@ -87,6 +90,20 @@ where
             Self::Right(b) => b.execute_with_state_witness(input, witness),
         }
     }
+
+    fn execute_with_state_hook<F>(
+        self,
+        input: Self::Input<'_>,
+        state_hook: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook,
+    {
+        match self {
+            Self::Left(a) => a.execute_with_state_hook(input, state_hook),
+            Self::Right(b) => b.execute_with_state_hook(input, state_hook),
+        }
+    }
 }
 
 impl<A, B, DB> BatchExecutor<DB> for Either<A, B>

crates/evm/src/execute.rs

@@ -12,6 +12,8 @@ use reth_prune_types::PruneModes;
 use revm::State;
 use revm_primitives::db::Database;
 
+use crate::system_calls::OnStateHook;
+
 /// A general purpose executor trait that executes an input (e.g. block) and produces an output
 /// (e.g. state changes and receipts).
 ///
@@ -43,6 +45,16 @@ pub trait Executor<DB> {
     ) -> Result<Self::Output, Self::Error>
     where
         F: FnMut(&State<DB>);
+
+    /// Executes the EVM with the given input and accepts a state hook closure that is invoked with
+    /// the EVM state after execution.
+    fn execute_with_state_hook<F>(
+        self,
+        input: Self::Input<'_>,
+        state_hook: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook;
 }
 
 /// A general purpose executor that can execute multiple inputs in sequence, validate the outputs,
@@ -199,6 +211,17 @@ mod tests {
         {
             Err(BlockExecutionError::msg("execution unavailable for tests"))
         }
+
+        fn execute_with_state_hook<F>(
+            self,
+            _: Self::Input<'_>,
+            _: F,
+        ) -> Result<Self::Output, Self::Error>
+        where
+            F: OnStateHook,
+        {
+            Err(BlockExecutionError::msg("execution unavailable for tests"))
+        }
     }
 
     impl<DB> BatchExecutor<DB> for TestExecutor<DB> {

crates/evm/src/noop.rs

@@ -10,7 +10,10 @@ use reth_storage_errors::provider::ProviderError;
 use revm::State;
 use revm_primitives::db::Database;
 
-use crate::execute::{BatchExecutor, BlockExecutorProvider, Executor};
+use crate::{
+    execute::{BatchExecutor, BlockExecutorProvider, Executor},
+    system_calls::OnStateHook,
+};
 
 const UNAVAILABLE_FOR_NOOP: &str = "execution unavailable for noop";
 
@@ -58,6 +61,17 @@ impl<DB> Executor<DB> for NoopBlockExecutorProvider {
     {
         Err(BlockExecutionError::msg(UNAVAILABLE_FOR_NOOP))
     }
+
+    fn execute_with_state_hook<F>(
+        self,
+        _: Self::Input<'_>,
+        _: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook,
+    {
+        Err(BlockExecutionError::msg(UNAVAILABLE_FOR_NOOP))
+    }
 }
 
 impl<DB> BatchExecutor<DB> for NoopBlockExecutorProvider {

crates/evm/src/system_calls/mod.rs

@@ -49,22 +49,19 @@ pub struct SystemCaller<'a, EvmConfig, Chainspec, Hook = NoopHook> {
     hook: Option<Hook>,
 }
 
-impl<'a, EvmConfig, Chainspec> SystemCaller<'a, EvmConfig, Chainspec> {
+impl<'a, EvmConfig, Chainspec> SystemCaller<'a, EvmConfig, Chainspec, NoopHook> {
     /// Create a new system caller with the given EVM config, database, and chain spec, and creates
     /// the EVM with the given initialized config and block environment.
     pub const fn new(evm_config: &'a EvmConfig, chain_spec: Chainspec) -> Self {
         Self { evm_config, chain_spec, hook: None }
     }
-}
-
-impl<'a, EvmConfig, Chainspec, Hook> SystemCaller<'a, EvmConfig, Chainspec, Hook> {
     /// Installs a custom hook to be called after each state change.
     pub fn with_state_hook<H: OnStateHook>(
         self,
-        hook: H,
+        hook: Option<H>,
     ) -> SystemCaller<'a, EvmConfig, Chainspec, H> {
         let Self { evm_config, chain_spec, .. } = self;
-        SystemCaller { evm_config, chain_spec, hook: Some(hook) }
+        SystemCaller { evm_config, chain_spec, hook }
     }
     /// Convenience method to consume the type and drop borrowed fields
     pub fn finish(self) {}
@@ -321,4 +318,11 @@ where
 
         eip7251::post_commit(result_and_state.result)
     }
+
+    /// Delegate to stored `OnStateHook`, noop if hook is `None`.
+    pub fn on_state(&mut self, state: &ResultAndState) {
+        if let Some(ref mut hook) = &mut self.hook {
+            hook.on_state(state);
+        }
+    }
 }

crates/evm/src/test_utils.rs

@@ -1,7 +1,10 @@
 //! Helpers for testing.
 
-use crate::execute::{
-    BatchExecutor, BlockExecutionInput, BlockExecutionOutput, BlockExecutorProvider, Executor,
+use crate::{
+    execute::{
+        BatchExecutor, BlockExecutionInput, BlockExecutionOutput, BlockExecutorProvider, Executor,
+    },
+    system_calls::OnStateHook,
 };
 use alloy_primitives::BlockNumber;
 use parking_lot::Mutex;
@@ -73,6 +76,17 @@ impl<DB> Executor<DB> for MockExecutorProvider {
     {
         unimplemented!()
     }
+
+    fn execute_with_state_hook<F>(
+        self,
+        _: Self::Input<'_>,
+        _: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook,
+    {
+        unimplemented!()
+    }
 }
 
 impl<DB> BatchExecutor<DB> for MockExecutorProvider {

crates/optimism/evm/src/execute.rs

@@ -10,7 +10,7 @@ use reth_evm::{
         BatchExecutor, BlockExecutionError, BlockExecutionInput, BlockExecutionOutput,
         BlockExecutorProvider, BlockValidationError, Executor, ProviderError,
     },
-    system_calls::SystemCaller,
+    system_calls::{NoopHook, OnStateHook, SystemCaller},
     ConfigureEvm,
 };
 use reth_execution_types::ExecutionOutcome;
@@ -111,15 +111,18 @@ where
     /// # Note
     ///
     /// It does __not__ apply post-execution changes.
-    fn execute_pre_and_transactions<Ext, DB>(
+    fn execute_pre_and_transactions<Ext, DB, F>(
         &self,
         block: &BlockWithSenders,
         mut evm: Evm<'_, Ext, &mut State<DB>>,
+        state_hook: Option<F>,
     ) -> Result<(Vec<Receipt>, u64), BlockExecutionError>
     where
         DB: Database<Error: Into<ProviderError> + Display>,
+        F: OnStateHook,
     {
-        let mut system_caller = SystemCaller::new(&self.evm_config, &self.chain_spec);
+        let mut system_caller =
+            SystemCaller::new(&self.evm_config, &self.chain_spec).with_state_hook(state_hook);
 
         // apply pre execution changes
         system_caller.apply_beacon_root_contract_call(
@@ -193,7 +196,8 @@ where
                 "Executed transaction"
             );
 
-            evm.db_mut().commit(state);
+            evm.db_mut().commit(state.clone());
+            system_caller.on_state(&ResultAndState { result: result.clone(), state });
 
             // append gas used
             cumulative_gas_used += result.gas_used();
@@ -278,16 +282,28 @@ where
         EnvWithHandlerCfg::new_with_cfg_env(cfg, block_env, Default::default())
     }
 
+    fn execute_without_verification(
+        &mut self,
+        block: &BlockWithSenders,
+        total_difficulty: U256,
+    ) -> Result<(Vec<Receipt>, u64), BlockExecutionError> {
+        self.execute_without_verification_with_state_hook(block, total_difficulty, None::<NoopHook>)
+    }
+
     /// Execute a single block and apply the state changes to the internal state.
     ///
     /// Returns the receipts of the transactions in the block and the total gas used.
     ///
     /// Returns an error if execution fails.
-    fn execute_without_verification(
+    fn execute_without_verification_with_state_hook<F>(
         &mut self,
         block: &BlockWithSenders,
         total_difficulty: U256,
-    ) -> Result<(Vec<Receipt>, u64), BlockExecutionError> {
+        state_hook: Option<F>,
+    ) -> Result<(Vec<Receipt>, u64), BlockExecutionError>
+    where
+        F: OnStateHook,
+    {
         // 1. prepare state on new block
         self.on_new_block(&block.header);
 
@@ -296,7 +312,7 @@ where
 
         let (receipts, gas_used) = {
             let evm = self.executor.evm_config.evm_with_env(&mut self.state, env);
-            self.executor.execute_pre_and_transactions(block, evm)
+            self.executor.execute_pre_and_transactions(block, evm, state_hook)
         }?;
 
         // 3. apply post execution changes
@@ -383,6 +399,32 @@ where
             gas_used,
         })
     }
+
+    fn execute_with_state_hook<F>(
+        mut self,
+        input: Self::Input<'_>,
+        state_hook: F,
+    ) -> Result<Self::Output, Self::Error>
+    where
+        F: OnStateHook,
+    {
+        let BlockExecutionInput { block, total_difficulty } = input;
+        let (receipts, gas_used) = self.execute_without_verification_with_state_hook(
+            block,
+            total_difficulty,
+            Some(state_hook),
+        )?;
+
+        // NOTE: we need to merge keep the reverts for the bundle retention
+        self.state.merge_transitions(BundleRetention::Reverts);
+
+        Ok(BlockExecutionOutput {
+            state: self.state.take_bundle(),
+            receipts,
+            requests: vec![],
+            gas_used,
+        })
+    }
 }
 
 /// An executor for a batch of blocks.