From 6e00e5842669c98063d55931650c81d7a0ce3da1 Mon Sep 17 00:00:00 2001
From: Matthias Seitz <matthias.seitz@outlook.de>
Date: Fri, 15 Nov 2024 18:33:37 +0100
Subject: [PATCH] feat: add payload witness fn (#12579)

---
 crates/optimism/payload/src/builder.rs | 77 ++++++++++++++++++++------
 1 file changed, 61 insertions(+), 16 deletions(-)

diff --git a/crates/optimism/payload/src/builder.rs b/crates/optimism/payload/src/builder.rs
index 2542ff527df0..8fb569f3dac7 100644
--- a/crates/optimism/payload/src/builder.rs
+++ b/crates/optimism/payload/src/builder.rs
@@ -39,6 +39,7 @@ use crate::{
     payload::{OpBuiltPayload, OpPayloadBuilderAttributes},
 };
 use op_alloy_consensus::DepositTransaction;
+use op_alloy_rpc_types_engine::OpPayloadAttributes;
 use reth_revm::witness::ExecutionWitnessRecord;
 use reth_transaction_pool::pool::BestPayloadTransactions;
 
@@ -94,21 +95,6 @@ where
     EvmConfig: ConfigureEvm<Header = Header>,
     Txs: OpPayloadTransactions,
 {
-    /// Returns the configured [`CfgEnvWithHandlerCfg`] and [`BlockEnv`] for the targeted payload
-    /// (that has the `parent` as its parent).
-    pub fn cfg_and_block_env(
-        &self,
-        config: &PayloadConfig<OpPayloadBuilderAttributes>,
-        parent: &Header,
-    ) -> Result<(CfgEnvWithHandlerCfg, BlockEnv), EvmConfig::Error> {
-        let next_attributes = NextBlockEnvAttributes {
-            timestamp: config.attributes.timestamp(),
-            suggested_fee_recipient: config.attributes.suggested_fee_recipient(),
-            prev_randao: config.attributes.prev_randao(),
-        };
-        self.evm_config.next_cfg_and_block_env(parent, next_attributes)
-    }
-
     /// Constructs an Optimism payload from the transactions sent via the
     /// Payload attributes by the sequencer. If the `no_tx_pool` argument is passed in
     /// the payload attributes, the transaction pool will be ignored and the only transactions
@@ -126,7 +112,7 @@ where
         Pool: TransactionPool,
     {
         let (initialized_cfg, initialized_block_env) = self
-            .cfg_and_block_env(&args.config, &args.config.parent_header)
+            .cfg_and_block_env(&args.config.attributes, &args.config.parent_header)
             .map_err(PayloadBuilderError::other)?;
 
         let BuildArguments { client, pool, mut cached_reads, config, cancel, best_payload } = args;
@@ -161,6 +147,65 @@ where
     }
 }
 
+impl<EvmConfig, Txs> OpPayloadBuilder<EvmConfig, Txs>
+where
+    EvmConfig: ConfigureEvm<Header = Header>,
+{
+    /// Returns the configured [`CfgEnvWithHandlerCfg`] and [`BlockEnv`] for the targeted payload
+    /// (that has the `parent` as its parent).
+    pub fn cfg_and_block_env(
+        &self,
+        attributes: &OpPayloadBuilderAttributes,
+        parent: &Header,
+    ) -> Result<(CfgEnvWithHandlerCfg, BlockEnv), EvmConfig::Error> {
+        let next_attributes = NextBlockEnvAttributes {
+            timestamp: attributes.timestamp(),
+            suggested_fee_recipient: attributes.suggested_fee_recipient(),
+            prev_randao: attributes.prev_randao(),
+        };
+        self.evm_config.next_cfg_and_block_env(parent, next_attributes)
+    }
+
+    /// Computes the witness for the payload.
+    pub fn payload_witness<Client>(
+        &self,
+        client: &Client,
+        parent: SealedHeader,
+        attributes: OpPayloadAttributes,
+    ) -> Result<ExecutionWitness, PayloadBuilderError>
+    where
+        Client: StateProviderFactory + ChainSpecProvider<ChainSpec = OpChainSpec>,
+    {
+        let attributes = OpPayloadBuilderAttributes::try_new(parent.hash(), attributes, 3)
+            .map_err(PayloadBuilderError::other)?;
+
+        let (initialized_cfg, initialized_block_env) =
+            self.cfg_and_block_env(&attributes, &parent).map_err(PayloadBuilderError::other)?;
+
+        let config = PayloadConfig {
+            parent_header: Arc::new(parent),
+            attributes,
+            extra_data: Default::default(),
+        };
+        let ctx = OpPayloadBuilderCtx {
+            evm_config: self.evm_config.clone(),
+            chain_spec: client.chain_spec(),
+            config,
+            initialized_cfg,
+            initialized_block_env,
+            cancel: Default::default(),
+            best_payload: Default::default(),
+        };
+
+        let state_provider = client.state_by_block_hash(ctx.parent().hash())?;
+        let state = StateProviderDatabase::new(state_provider);
+        let mut state = State::builder().with_database(state).with_bundle_update().build();
+
+        let builder = OpBuilder { pool: NoopTransactionPool::default(), best: () };
+        builder.witness(&mut state, &ctx)
+    }
+}
+
 /// Implementation of the [`PayloadBuilder`] trait for [`OpPayloadBuilder`].
 impl<Pool, Client, EvmConfig, Txs> PayloadBuilder<Pool, Client> for OpPayloadBuilder<EvmConfig, Txs>
 where