feat(executor): Generic precompile overrides

bin/client/src/kona.rs

@@ -13,7 +13,7 @@ use kona_client::{
     BootInfo, CachingOracle,
 };
 use kona_common_proc::client_entry;
-use kona_executor::StatelessL2BlockExecutor;
+use kona_executor::{NoPrecompileOverride, StatelessL2BlockExecutor};
 use kona_primitives::L2AttributesWithParent;
 
 extern crate alloc;
@@ -61,6 +61,7 @@ fn main() -> Result<()> {
             .with_parent_header(driver.take_l2_safe_head_header())
             .with_fetcher(l2_provider)
             .with_hinter(TrieDBHintWriter)
+            .with_precompile_overrides(NoPrecompileOverride)
             .build()?;
         let Header { number, .. } = *executor.execute_payload(attributes)?;
         let output_root = executor.compute_output_root()?;

crates/executor/src/builder.rs

@@ -1,37 +1,47 @@
 //! Contains the builder pattern for the [StatelessL2BlockExecutor].
 
-use crate::StatelessL2BlockExecutor;
+use crate::{PrecompileOverride, StatelessL2BlockExecutor};
 use alloy_consensus::{Header, Sealable, Sealed};
 use anyhow::Result;
 use kona_derive::types::RollupConfig;
-use kona_mpt::{NoopTrieDBFetcher, NoopTrieDBHinter, TrieDB, TrieDBFetcher, TrieDBHinter};
+use kona_mpt::{TrieDB, TrieDBFetcher, TrieDBHinter};
 use revm::StateBuilder;
 
 /// The builder pattern for the [StatelessL2BlockExecutor].
 #[derive(Debug)]
-pub struct StatelessL2BlockExecutorBuilder<'a, F = NoopTrieDBFetcher, H = NoopTrieDBHinter>
+pub struct StatelessL2BlockExecutorBuilder<'a, F, H, PO>
 where
     F: TrieDBFetcher,
     H: TrieDBHinter,
+    PO: PrecompileOverride<F, H>,
 {
     /// The [RollupConfig].
     config: &'a RollupConfig,
     /// The parent [Header] to begin execution from.
     parent_header: Option<Sealed<Header>>,
+    /// The precompile overrides to use during execution.
+    precompile_overrides: Option<PO>,
     /// The [TrieDBFetcher] to fetch the state trie preimages.
     fetcher: Option<F>,
     /// The [TrieDBHinter] to hint the state trie preimages.
     hinter: Option<H>,
 }
 
-impl<'a, F, H> StatelessL2BlockExecutorBuilder<'a, F, H>
+impl<'a, F, H, PO> StatelessL2BlockExecutorBuilder<'a, F, H, PO>
 where
     F: TrieDBFetcher,
     H: TrieDBHinter,
+    PO: PrecompileOverride<F, H>,
 {
     /// Instantiate a new builder with the given [RollupConfig].
     pub fn with_config(config: &'a RollupConfig) -> Self {
-        Self { config, parent_header: None, fetcher: None, hinter: None }
+        Self {
+            config,
+            parent_header: None,
+            precompile_overrides: None,
+            fetcher: None,
+            hinter: None,
+        }
     }
 
     /// Set the [Header] to begin execution from.
@@ -52,8 +62,14 @@ where
         self
     }
 
+    /// Set the precompile overrides to use during execution.
+    pub fn with_precompile_overrides(mut self, precompile_overrides: PO) -> Self {
+        self.precompile_overrides = Some(precompile_overrides);
+        self
+    }
+
     /// Build the [StatelessL2BlockExecutor] from the builder configuration.
-    pub fn build(self) -> Result<StatelessL2BlockExecutor<'a, F, H>> {
+    pub fn build(self) -> Result<StatelessL2BlockExecutor<'a, F, H, PO>> {
         let fetcher = self.fetcher.ok_or(anyhow::anyhow!("Fetcher not set"))?;
         let hinter = self.hinter.ok_or(anyhow::anyhow!("Hinter not set"))?;
         let parent_header = self.parent_header.unwrap_or_else(|| {
@@ -64,6 +80,10 @@ where
         let trie_db = TrieDB::new(parent_header.state_root, parent_header, fetcher, hinter);
         let state = StateBuilder::new_with_database(trie_db).with_bundle_update().build();
 
-        Ok(StatelessL2BlockExecutor { config: self.config, state })
+        Ok(StatelessL2BlockExecutor {
+            config: self.config,
+            state,
+            _phantom: core::marker::PhantomData::<PO>,
+        })
     }
 }

crates/executor/src/lib.rs

@@ -12,10 +12,7 @@ use alloy_eips::eip2718::{Decodable2718, Encodable2718};
 use alloy_primitives::{address, keccak256, Address, Bytes, TxKind, B256, U256};
 use anyhow::{anyhow, Result};
 use kona_derive::types::{L2PayloadAttributes, RawTransaction, RollupConfig};
-use kona_mpt::{
-    ordered_trie_with_encoder, NoopTrieDBFetcher, NoopTrieDBHinter, TrieDB, TrieDBFetcher,
-    TrieDBHinter,
-};
+use kona_mpt::{ordered_trie_with_encoder, TrieDB, TrieDBFetcher, TrieDBHinter};
 use op_alloy_consensus::{OpReceiptEnvelope, OpTxEnvelope};
 use revm::{
     db::{states::bundle_state::BundleRetention, State},
@@ -30,6 +27,9 @@ use tracing::{debug, info};
 mod builder;
 pub use builder::StatelessL2BlockExecutorBuilder;
 
+mod precompile;
+pub use precompile::{NoPrecompileOverride, PrecompileOverride};
+
 mod eip4788;
 use eip4788::pre_block_beacon_root_contract_call;
 
@@ -42,24 +42,28 @@ use util::{extract_tx_gas_limit, is_system_transaction, logs_bloom, receipt_enve
 /// The block executor for the L2 client program. Operates off of a [TrieDB] backed [State],
 /// allowing for stateless block execution of OP Stack blocks.
 #[derive(Debug)]
-pub struct StatelessL2BlockExecutor<'a, F = NoopTrieDBFetcher, H = NoopTrieDBHinter>
+pub struct StatelessL2BlockExecutor<'a, F, H, PO>
 where
     F: TrieDBFetcher,
     H: TrieDBHinter,
+    PO: PrecompileOverride<F, H>,
 {
     /// The [RollupConfig].
     config: &'a RollupConfig,
     /// The inner state database component.
     state: State<TrieDB<F, H>>,
+    /// Phantom data for the precompile overrides.
+    _phantom: core::marker::PhantomData<PO>,
 }
 
-impl<'a, F, H> StatelessL2BlockExecutor<'a, F, H>
+impl<'a, F, H, PO> StatelessL2BlockExecutor<'a, F, H, PO>
 where
     F: TrieDBFetcher,
     H: TrieDBHinter,
+    PO: PrecompileOverride<F, H>,
 {
     /// Constructs a new [StatelessL2BlockExecutorBuilder] with the given [RollupConfig].
-    pub fn builder(config: &'a RollupConfig) -> StatelessL2BlockExecutorBuilder<'a, F, H> {
+    pub fn builder(config: &'a RollupConfig) -> StatelessL2BlockExecutorBuilder<'a, F, H, PO> {
         StatelessL2BlockExecutorBuilder::with_config(config)
     }
 
@@ -114,12 +118,22 @@ where
         // Ensure that the create2 contract is deployed upon transition to the Canyon hardfork.
         ensure_create2_deployer_canyon(&mut self.state, self.config, payload.timestamp)?;
 
-        // Construct the EVM with the given configuration.
-        // TODO(clabby): Accelerate precompiles w/ custom precompile handler.
         let mut cumulative_gas_used = 0u64;
         let mut receipts: Vec<OpReceiptEnvelope> = Vec::with_capacity(payload.transactions.len());
         let is_regolith = self.config.is_regolith_active(payload.timestamp);
 
+        // Construct the block-scoped EVM with the given configuration.
+        // The transaction environment is set within the loop for each transaction.
+        let mut evm = Evm::builder()
+            .with_db(&mut self.state)
+            .with_env_with_handler_cfg(EnvWithHandlerCfg::new_with_cfg_env(
+                initialized_cfg.clone(),
+                initialized_block_env.clone(),
+                Default::default(),
+            ))
+            .append_handler_register(PO::set_precompiles)
+            .build();
+
         // Execute the transactions in the payload.
         let transactions = payload
             .transactions
@@ -144,13 +158,10 @@ where
                 anyhow::bail!("EIP-4844 transactions are not supported");
             }
 
-            let mut evm = Evm::builder()
-                .with_db(&mut self.state)
-                .with_env_with_handler_cfg(EnvWithHandlerCfg::new_with_cfg_env(
-                    initialized_cfg.clone(),
-                    initialized_block_env.clone(),
-                    Self::prepare_tx_env(&transaction, raw_transaction)?,
-                ))
+            // Modify the transaction environment with the current transaction.
+            evm = evm
+                .modify()
+                .with_tx_env(Self::prepare_tx_env(&transaction, raw_transaction)?)
                 .build();
 
             // If the transaction is a deposit, cache the depositor account.
@@ -208,6 +219,9 @@ where
             cumulative_gas_used = cumulative_gas_used
         );
 
+        // Drop the EVM to rid the exclusive reference to the database.
+        drop(evm);
+
         // Merge all state transitions into the cache state.
         debug!(target: "client_executor", "Merging state transitions");
         self.state.merge_transitions(BundleRetention::Reverts);
@@ -702,6 +716,7 @@ mod test {
             .with_parent_header(header.seal_slow())
             .with_fetcher(TestdataTrieDBFetcher::new("block_120794432_exec"))
             .with_hinter(NoopTrieDBHinter)
+            .with_precompile_overrides(NoPrecompileOverride)
             .build()
             .unwrap();
 
@@ -755,6 +770,7 @@ mod test {
             .with_parent_header(parent_header.seal_slow())
             .with_fetcher(TestdataTrieDBFetcher::new("block_121049889_exec"))
             .with_hinter(NoopTrieDBHinter)
+            .with_precompile_overrides(NoPrecompileOverride)
             .build()
             .unwrap();
 
@@ -812,6 +828,7 @@ mod test {
             .with_parent_header(parent_header.seal_slow())
             .with_fetcher(TestdataTrieDBFetcher::new("block_121003241_exec"))
             .with_hinter(NoopTrieDBHinter)
+            .with_precompile_overrides(NoPrecompileOverride)
             .build()
             .unwrap();
 
@@ -876,6 +893,7 @@ mod test {
             .with_parent_header(parent_header.seal_slow())
             .with_fetcher(TestdataTrieDBFetcher::new("block_121057303_exec"))
             .with_hinter(NoopTrieDBHinter)
+            .with_precompile_overrides(NoPrecompileOverride)
             .build()
             .unwrap();
 
@@ -934,6 +952,7 @@ mod test {
             .with_parent_header(parent_header.seal_slow())
             .with_fetcher(TestdataTrieDBFetcher::new("block_121065789_exec"))
             .with_hinter(NoopTrieDBHinter)
+            .with_precompile_overrides(NoPrecompileOverride)
             .build()
             .unwrap();
 
@@ -1001,6 +1020,7 @@ mod test {
             .with_parent_header(parent_header.seal_slow())
             .with_fetcher(TestdataTrieDBFetcher::new("block_121135704_exec"))
             .with_hinter(NoopTrieDBHinter)
+            .with_precompile_overrides(NoPrecompileOverride)
             .build()
             .unwrap();
 

crates/executor/src/precompile.rs

@@ -0,0 +1,28 @@
+//! Contains the [PrecompileOverride] trait.
+
+use kona_mpt::{TrieDB, TrieDBFetcher, TrieDBHinter};
+use revm::{db::State, handler::register::EvmHandler};
+
+/// A trait for defining precompile overrides during execution.
+pub trait PrecompileOverride<F, H>
+where
+    F: TrieDBFetcher,
+    H: TrieDBHinter,
+{
+    /// Set the precompiles to use during execution.
+    fn set_precompiles(handler: &mut EvmHandler<'_, (), &mut State<TrieDB<F, H>>>);
+}
+
+/// Default implementation of [PrecompileOverride], which does not override any precompiles.
+#[derive(Debug, Default)]
+pub struct NoPrecompileOverride;
+
+impl<F, H> PrecompileOverride<F, H> for NoPrecompileOverride
+where
+    F: TrieDBFetcher,
+    H: TrieDBHinter,
+{
+    fn set_precompiles(_: &mut EvmHandler<'_, (), &mut State<TrieDB<F, H>>>) {
+        // Do nothing
+    }
+}