feat: add helpers to obtain the execution witness for a payload

crates/optimism/payload/Cargo.toml

@@ -15,7 +15,7 @@ workspace = true
 # reth
 reth-chainspec.workspace = true
 reth-primitives.workspace = true
-reth-revm.workspace = true
+reth-revm = { workspace = true, features = ["witness"] }
 reth-transaction-pool.workspace = true
 reth-provider.workspace = true
 reth-rpc-types-compat.workspace = true
@@ -41,6 +41,7 @@ alloy-rlp.workspace = true
 op-alloy-rpc-types-engine.workspace = true
 op-alloy-consensus.workspace = true
 alloy-rpc-types-engine.workspace = true
+alloy-rpc-types-debug.workspace = true
 alloy-consensus.workspace = true
 
 # misc

crates/optimism/payload/src/builder.rs

@@ -5,6 +5,7 @@ use std::{fmt::Display, sync::Arc};
 use alloy_consensus::{Header, Transaction, EMPTY_OMMER_ROOT_HASH};
 use alloy_eips::merge::BEACON_NONCE;
 use alloy_primitives::{Address, Bytes, U256};
+use alloy_rpc_types_debug::ExecutionWitness;
 use alloy_rpc_types_engine::PayloadId;
 use reth_basic_payload_builder::*;
 use reth_chain_state::ExecutedBlock;
@@ -20,7 +21,7 @@ use reth_primitives::{
     revm_primitives::{BlockEnv, CfgEnvWithHandlerCfg},
     Block, BlockBody, Receipt, SealedHeader, TransactionSigned, TxType,
 };
-use reth_provider::{ProviderError, StateProviderFactory, StateRootProvider};
+use reth_provider::{ProviderError, StateProofProvider, StateProviderFactory, StateRootProvider};
 use reth_revm::database::StateProviderDatabase;
 use reth_transaction_pool::{
     noop::NoopTransactionPool, BestTransactionsAttributes, PayloadTransactions, TransactionPool,
@@ -38,6 +39,7 @@ use crate::{
     payload::{OpBuiltPayload, OpPayloadBuilderAttributes},
 };
 use op_alloy_consensus::DepositTransaction;
+use reth_revm::witness::ExecutionWitnessRecord;
 use reth_transaction_pool::pool::BestPayloadTransactions;
 
 /// Optimism's payload builder
@@ -234,36 +236,33 @@ where
     Pool: TransactionPool,
     Txs: OpPayloadTransactions,
 {
-    /// Builds the payload on top of the state.
-    pub fn build<EvmConfig, DB, P>(
+    /// Executes the payload and returns the outcome.
+    pub fn execute<EvmConfig, DB>(
         self,
-        mut db: State<DB>,
-        ctx: OpPayloadBuilderCtx<EvmConfig>,
-    ) -> Result<BuildOutcomeKind<OpBuiltPayload>, PayloadBuilderError>
+        state: &mut State<DB>,
+        ctx: &OpPayloadBuilderCtx<EvmConfig>,
+    ) -> Result<BuildOutcomeKind<ExecutedPayload>, PayloadBuilderError>
     where
         EvmConfig: ConfigureEvm<Header = Header>,
-        DB: Database<Error = ProviderError> + AsRef<P>,
-        P: StateRootProvider,
+        DB: Database<Error = ProviderError>,
     {
         let Self { pool, best } = self;
         debug!(target: "payload_builder", id=%ctx.payload_id(), parent_header = ?ctx.parent().hash(), parent_number = ctx.parent().number, "building new payload");
 
         // 1. apply eip-4788 pre block contract call
-        ctx.apply_pre_beacon_root_contract_call(&mut db)?;
+        ctx.apply_pre_beacon_root_contract_call(state)?;
 
         // 2. ensure create2deployer is force deployed
-        ctx.ensure_create2_deployer(&mut db)?;
+        ctx.ensure_create2_deployer(state)?;
 
         // 3. execute sequencer transactions
-        let mut info = ctx.execute_sequencer_transactions(&mut db)?;
+        let mut info = ctx.execute_sequencer_transactions(state)?;
 
         // 4. if mem pool transactions are requested we execute them
         if !ctx.attributes().no_tx_pool {
             let best_txs = best.best_transactions(pool, ctx.best_transaction_attributes());
-            if let Some(cancelled) =
-                ctx.execute_best_transactions::<_, Pool>(&mut info, &mut db, best_txs)?
-            {
-                return Ok(cancelled)
+            if ctx.execute_best_transactions::<_, Pool>(&mut info, state, best_txs)?.is_some() {
+                return Ok(BuildOutcomeKind::Cancelled)
             }
 
             // check if the new payload is even more valuable
@@ -273,16 +272,38 @@ where
             }
         }
 
-        let WithdrawalsOutcome { withdrawals_root, withdrawals } =
-            ctx.commit_withdrawals(&mut db)?;
+        let withdrawals_outcome = ctx.commit_withdrawals(state)?;
 
         // merge all transitions into bundle state, this would apply the withdrawal balance changes
         // and 4788 contract call
-        db.merge_transitions(BundleRetention::Reverts);
+        state.merge_transitions(BundleRetention::Reverts);
+
+        Ok(BuildOutcomeKind::Better { payload: ExecutedPayload { info, withdrawals_outcome } })
+    }
+
+    /// Builds the payload on top of the state.
+    pub fn build<EvmConfig, DB, P>(
+        self,
+        mut state: State<DB>,
+        ctx: OpPayloadBuilderCtx<EvmConfig>,
+    ) -> Result<BuildOutcomeKind<OpBuiltPayload>, PayloadBuilderError>
+    where
+        EvmConfig: ConfigureEvm<Header = Header>,
+        DB: Database<Error = ProviderError> + AsRef<P>,
+        P: StateRootProvider,
+    {
+        let ExecutedPayload {
+            info,
+            withdrawals_outcome: WithdrawalsOutcome { withdrawals, withdrawals_root },
+        } = match self.execute(&mut state, &ctx)? {
+            BuildOutcomeKind::Better { payload } | BuildOutcomeKind::Freeze(payload) => payload,
+            BuildOutcomeKind::Cancelled => return Ok(BuildOutcomeKind::Cancelled),
+            BuildOutcomeKind::Aborted { fees } => return Ok(BuildOutcomeKind::Aborted { fees }),
+        };
 
         let block_number = ctx.block_number();
         let execution_outcome = ExecutionOutcome::new(
-            db.take_bundle(),
+            state.take_bundle(),
             vec![info.receipts.clone()].into(),
             block_number,
             Vec::new(),
@@ -302,7 +323,7 @@ where
         // // calculate the state root
         let hashed_state = HashedPostState::from_bundle_state(&execution_outcome.state().state);
         let (state_root, trie_output) = {
-            db.database.as_ref().state_root_with_updates(hashed_state.clone()).inspect_err(
+            state.database.as_ref().state_root_with_updates(hashed_state.clone()).inspect_err(
                 |err| {
                     warn!(target: "payload_builder",
                     parent_header=%ctx.parent().hash(),
@@ -388,6 +409,24 @@ where
             Ok(BuildOutcomeKind::Better { payload })
         }
     }
+
+    /// Builds the payload and returns its [`ExecutionWitness`] based on the state after execution.
+    pub fn witness<EvmConfig, DB, P>(
+        self,
+        state: &mut State<DB>,
+        ctx: &OpPayloadBuilderCtx<EvmConfig>,
+    ) -> Result<ExecutionWitness, PayloadBuilderError>
+    where
+        EvmConfig: ConfigureEvm<Header = Header>,
+        DB: Database<Error = ProviderError> + AsRef<P>,
+        P: StateProofProvider,
+    {
+        let _ = self.execute(state, ctx)?;
+        let ExecutionWitnessRecord { hashed_state, codes, keys } =
+            ExecutionWitnessRecord::from_executed_state(state);
+        let state = state.database.as_ref().witness(Default::default(), hashed_state)?;
+        Ok(ExecutionWitness { state: state.into_iter().collect(), codes, keys })
+    }
 }
 
 /// A type that returns a the [`PayloadTransactions`] that should be included in the pool.
@@ -411,6 +450,15 @@ impl OpPayloadTransactions for () {
     }
 }
 
+/// Holds the state after execution
+#[derive(Debug)]
+pub struct ExecutedPayload {
+    /// Tracked execution info
+    pub info: ExecutionInfo,
+    /// Outcome after committing withdrawals.
+    pub withdrawals_outcome: WithdrawalsOutcome,
+}
+
 /// This acts as the container for executed transactions and its byproducts (receipts, gas used)
 #[derive(Default, Debug)]
 pub struct ExecutionInfo {
@@ -725,13 +773,15 @@ where
         Ok(info)
     }
 
-    /// Executes the given best transactions and updates the execution info
+    /// Executes the given best transactions and updates the execution info.
+    ///
+    /// Returns `Ok(Some(())` if the job was cancelled.
     pub fn execute_best_transactions<DB, Pool>(
         &self,
         info: &mut ExecutionInfo,
         db: &mut State<DB>,
         mut best_txs: impl PayloadTransactions,
-    ) -> Result<Option<BuildOutcomeKind<OpBuiltPayload>>, PayloadBuilderError>
+    ) -> Result<Option<()>, PayloadBuilderError>
     where
         DB: Database<Error = ProviderError>,
         Pool: TransactionPool,
@@ -764,7 +814,7 @@ where
 
             // check if the job was cancelled, if so we can exit early
             if self.cancel.is_cancelled() {
-                return Ok(Some(BuildOutcomeKind::Cancelled))
+                return Ok(Some(()))
             }
 
             // Configure the environment for the tx.