refactor contract simulation

Author: avalonche

crates/op-rbuilder/src/builders/builder_tx.rs

@@ -1,13 +1,15 @@
-use alloy_consensus::{Transaction, TxEip1559};
+use alloy_consensus::TxEip1559;
 use alloy_eips::{Encodable2718, eip7623::TOTAL_COST_FLOOR_PER_TOKEN};
 use alloy_evm::Database;
 use alloy_op_evm::OpEvm;
 use alloy_primitives::{
     Address, B256, Bytes, TxKind, U256,
     map::foldhash::{HashMap, HashSet, HashSetExt},
 };
+use alloy_sol_types::{ContractError, Error, Revert, SolCall, SolError, SolInterface};
 use core::fmt::Debug;
 use op_alloy_consensus::OpTypedTransaction;
+use op_alloy_rpc_types::OpTransactionRequest;
 use op_revm::{OpHaltReason, OpTransactionError};
 use reth_evm::{
     ConfigureEvm, Evm, EvmError, eth::receipt_builder::ReceiptBuilderCtx,
@@ -18,22 +20,27 @@ use reth_optimism_primitives::OpTransactionSigned;
 use reth_primitives::Recovered;
 use reth_provider::{ProviderError, StateProvider};
 use reth_revm::{State, database::StateProviderDatabase};
+use reth_rpc_api::eth::{EthTxEnvError, transaction::TryIntoTxEnv};
 use revm::{
-    DatabaseCommit,
-    context::result::{EVMError, ExecutionResult, ResultAndState},
+    DatabaseCommit, DatabaseRef,
+    context::{
+        ContextTr,
+        result::{EVMError, ExecutionResult, ResultAndState},
+    },
     inspector::NoOpInspector,
     state::Account,
 };
+use std::borrow::Cow;
 use tracing::warn;
 
 use crate::{
     builders::context::OpPayloadBuilderCtx, primitives::reth::ExecutionInfo, tx_signer::Signer,
 };
 
 #[derive(Debug, Default)]
-pub struct SimulationSuccessResult {
+pub struct SimulationSuccessResult<T: SolCall> {
     pub gas_used: u64,
-    pub output: Bytes,
+    pub output: T::Return,
     pub state_changes: HashMap<Address, Account>,
 }
 
@@ -61,8 +68,8 @@ impl BuilderTransactionCtx {
 
 #[derive(Debug, thiserror::Error)]
 pub enum InvalidContractDataError {
-    #[error("did not find expected log {0:?} in emitted logs")]
-    InvalidLogs(B256),
+    #[error("did not find expected logs expected {0:?} but got {1:?}")]
+    InvalidLogs(Vec<B256>, Vec<B256>),
     #[error("could not decode output from contract call")]
     OutputAbiDecodeError,
 }
@@ -80,11 +87,11 @@ pub enum BuilderTransactionError {
     #[error("contract {0} may be incorrect, invalid contract data: {1}")]
     InvalidContract(Address, InvalidContractDataError),
     /// Transaction halted execution
-    #[error("transaction halted {0:?}")]
-    TransactionHalted(OpHaltReason),
+    #[error("transaction to {0} halted {1:?}")]
+    TransactionHalted(Address, OpHaltReason),
     /// Transaction reverted
-    #[error("transaction reverted {0}")]
-    TransactionReverted(Box<dyn core::error::Error + Send + Sync>),
+    #[error("transaction to {0} reverted {1}")]
+    TransactionReverted(Address, Revert),
     /// Invalid tx errors during evm execution.
     #[error("invalid transaction error {0}")]
     InvalidTransactionError(Box<dyn core::error::Error + Send + Sync>),
@@ -108,6 +115,12 @@ impl From<EVMError<ProviderError, OpTransactionError>> for BuilderTransactionErr
     }
 }
 
+impl From<EthTxEnvError> for BuilderTransactionError {
+    fn from(error: EthTxEnvError) -> Self {
+        BuilderTransactionError::EvmExecutionError(Box::new(error))
+    }
+}
+
 impl From<BuilderTransactionError> for PayloadBuilderError {
     fn from(error: BuilderTransactionError) -> Self {
         match error {
@@ -130,15 +143,35 @@ impl BuilderTransactionError {
 }
 
 pub trait BuilderTransactions<ExtraCtx: Debug + Default = (), Extra: Debug + Default = ()> {
+    // Simulates and returns the signed builder transactions. The simulation modifies and commit
+    // changes to the db so call new_simulation_state to simulate on a new copy of the state
     fn simulate_builder_txs(
         &self,
         state_provider: impl StateProvider + Clone,
         info: &mut ExecutionInfo<Extra>,
         ctx: &OpPayloadBuilderCtx<ExtraCtx>,
-        db: &mut State<impl Database>,
+        db: &mut State<impl Database + DatabaseRef>,
         top_of_block: bool,
     ) -> Result<Vec<BuilderTransactionCtx>, BuilderTransactionError>;
 
+    fn simulate_builder_txs_with_new_state(
+        &self,
+        state_provider: impl StateProvider + Clone,
+        info: &mut ExecutionInfo<Extra>,
+        ctx: &OpPayloadBuilderCtx<ExtraCtx>,
+        db: &State<impl Database>,
+        top_of_block: bool,
+    ) -> Result<Vec<BuilderTransactionCtx>, BuilderTransactionError> {
+        let mut simulation_state = self.new_simulation_state(state_provider.clone(), db);
+        self.simulate_builder_txs(
+            state_provider,
+            info,
+            ctx,
+            &mut simulation_state,
+            top_of_block,
+        )
+    }
+
     fn add_builder_txs(
         &self,
         state_provider: impl StateProvider + Clone,
@@ -148,19 +181,20 @@ pub trait BuilderTransactions<ExtraCtx: Debug + Default = (), Extra: Debug + Def
         top_of_block: bool,
     ) -> Result<Vec<BuilderTransactionCtx>, BuilderTransactionError> {
         {
+            let builder_txs = self.simulate_builder_txs_with_new_state(
+                state_provider,
+                info,
+                builder_ctx,
+                db,
+                top_of_block,
+            )?;
+
             let mut evm = builder_ctx
                 .evm_config
                 .evm_with_env(&mut *db, builder_ctx.evm_env.clone());
 
             let mut invalid: HashSet<Address> = HashSet::new();
 
-            let builder_txs = self.simulate_builder_txs(
-                state_provider,
-                info,
-                builder_ctx,
-                evm.db_mut(),
-                top_of_block,
-            )?;
             for builder_tx in builder_txs.iter() {
                 if builder_tx.is_top_of_block != top_of_block {
                     // don't commit tx if the buidler tx is not being added in the intended
@@ -211,52 +245,37 @@ pub trait BuilderTransactions<ExtraCtx: Debug + Default = (), Extra: Debug + Def
         }
     }
 
-    fn simulate_builder_txs_state(
+    // Creates a copy of the state to simulate against
+    fn new_simulation_state(
         &self,
-        state_provider: impl StateProvider + Clone,
-        builder_txs: Vec<&BuilderTransactionCtx>,
-        ctx: &OpPayloadBuilderCtx<ExtraCtx>,
-        db: &mut State<impl Database>,
-    ) -> Result<State<StateProviderDatabase<impl StateProvider>>, BuilderTransactionError> {
-        let state = StateProviderDatabase::new(state_provider.clone());
-        let mut simulation_state = State::builder()
+        state_provider: impl StateProvider,
+        db: &State<impl Database>,
+    ) -> State<StateProviderDatabase<impl StateProvider>> {
+        let state = StateProviderDatabase::new(state_provider);
+
+        State::builder()
             .with_database(state)
             .with_cached_prestate(db.cache.clone())
             .with_bundle_update()
-            .build();
-        let mut evm = ctx
-            .evm_config
-            .evm_with_env(&mut simulation_state, ctx.evm_env.clone());
-
-        for builder_tx in builder_txs {
-            let ResultAndState { state, .. } = evm
-                .transact(&builder_tx.signed_tx)
-                .map_err(|err| BuilderTransactionError::EvmExecutionError(Box::new(err)))?;
-
-            evm.db_mut().commit(state);
-        }
-
-        Ok(simulation_state)
+            .build()
     }
 
     fn sign_tx(
         &self,
         to: Address,
         from: Signer,
-        gas_used: Option<u64>,
+        gas_used: u64,
         calldata: Bytes,
         ctx: &OpPayloadBuilderCtx<ExtraCtx>,
-        db: &mut State<impl Database>,
+        db: impl DatabaseRef,
     ) -> Result<Recovered<OpTransactionSigned>, BuilderTransactionError> {
         let nonce = get_nonce(db, from.address)?;
         // Create the EIP-1559 transaction
         let tx = OpTypedTransaction::Eip1559(TxEip1559 {
             chain_id: ctx.chain_id(),
             nonce,
             // Due to EIP-150, 63/64 of available gas is forwarded to external calls so need to add a buffer
-            gas_limit: gas_used
-                .map(|gas| gas * 64 / 63)
-                .unwrap_or(ctx.block_gas_limit()),
+            gas_limit: gas_used * 64 / 63,
             max_fee_per_gas: ctx.base_fee().into(),
             to: TxKind::Call(to),
             input: calldata,
@@ -265,18 +284,32 @@ pub trait BuilderTransactions<ExtraCtx: Debug + Default = (), Extra: Debug + Def
         Ok(from.sign_tx(tx)?)
     }
 
-    fn simulate_call(
+    fn commit_txs(
+        &self,
+        signed_txs: Vec<Recovered<OpTransactionSigned>>,
+        ctx: &OpPayloadBuilderCtx<ExtraCtx>,
+        db: &mut State<impl Database>,
+    ) -> Result<(), BuilderTransactionError> {
+        let mut evm = ctx.evm_config.evm_with_env(&mut *db, ctx.evm_env.clone());
+        for signed_tx in signed_txs {
+            let ResultAndState { state, .. } = evm
+                .transact(&signed_tx)
+                .map_err(|err| BuilderTransactionError::EvmExecutionError(Box::new(err)))?;
+            evm.db_mut().commit(state)
+        }
+        Ok(())
+    }
+
+    fn simulate_call<T: SolCall, E: SolInterface>(
         &self,
-        signed_tx: Recovered<OpTransactionSigned>,
-        expected_topic: Option<B256>,
-        revert_handler: impl FnOnce(Bytes) -> BuilderTransactionError,
-        evm: &mut OpEvm<
-            &mut State<StateProviderDatabase<impl StateProvider>>,
-            NoOpInspector,
-            PrecompilesMap,
-        >,
-    ) -> Result<SimulationSuccessResult, BuilderTransactionError> {
-        let ResultAndState { result, state } = match evm.transact(&signed_tx) {
+        tx: OpTransactionRequest,
+        expected_logs: Vec<B256>,
+        evm: &mut OpEvm<impl Database, NoOpInspector, PrecompilesMap>,
+    ) -> Result<SimulationSuccessResult<T>, BuilderTransactionError> {
+        let tx_env = tx.try_into_tx_env(evm.cfg(), evm.block())?;
+        let to = tx_env.base.kind.into_to().unwrap_or_default();
+
+        let ResultAndState { result, state } = match evm.transact(tx_env) {
             Ok(res) => res,
             Err(err) => {
                 if err.is_invalid_tx_err() {
@@ -291,28 +324,54 @@ pub trait BuilderTransactions<ExtraCtx: Debug + Default = (), Extra: Debug + Def
 
         match result {
             ExecutionResult::Success {
-                logs,
-                gas_used,
                 output,
+                gas_used,
+                logs,
                 ..
             } => {
-                if let Some(topic) = expected_topic
-                    && !logs.iter().any(|log| log.topics().first() == Some(&topic))
+                let topics: HashSet<B256> = logs
+                    .into_iter()
+                    .flat_map(|log| log.topics().to_vec())
+                    .collect();
+                if !expected_logs
+                    .iter()
+                    .all(|expected_topic| topics.contains(expected_topic))
                 {
                     return Err(BuilderTransactionError::InvalidContract(
-                        signed_tx.to().unwrap_or_default(),
-                        InvalidContractDataError::InvalidLogs(topic),
+                        to,
+                        InvalidContractDataError::InvalidLogs(
+                            expected_logs,
+                            topics.into_iter().collect(),
+                        ),
                     ));
                 }
-                Ok(SimulationSuccessResult {
+                let return_output = T::abi_decode_returns(&output.into_data()).map_err(|_| {
+                    BuilderTransactionError::InvalidContract(
+                        to,
+                        InvalidContractDataError::OutputAbiDecodeError,
+                    )
+                })?;
+                Ok(SimulationSuccessResult::<T> {
                     gas_used,
-                    output: output.into_data(),
+                    output: return_output,
                     state_changes: state,
                 })
             }
-            ExecutionResult::Revert { output, .. } => Err(revert_handler(output)),
+            ExecutionResult::Revert { output, .. } => {
+                let revert = ContractError::<E>::abi_decode(&output)
+                    .and_then(|reason| {
+                        reason
+                            .try_into()
+                            .map_err(|_| Error::Other(Cow::Borrowed("failed to convert to revert")))
+                    })
+                    .or_else(|_| Revert::abi_decode(&output))
+                    .unwrap_or_else(|_| {
+                        Revert::from(format!("unknown revert: {}", hex::encode(&output)))
+                    });
+                Err(BuilderTransactionError::TransactionReverted(to, revert))
+            }
             ExecutionResult::Halt { reason, .. } => Err(BuilderTransactionError::other(
-                BuilderTransactionError::TransactionHalted(reason),
+                BuilderTransactionError::TransactionHalted(to, reason),
             )),
         }
     }
@@ -335,7 +394,7 @@ impl<ExtraCtx: Debug + Default> BuilderTxBase<ExtraCtx> {
     pub(super) fn simulate_builder_tx(
         &self,
         ctx: &OpPayloadBuilderCtx<ExtraCtx>,
-        db: &mut State<impl Database>,
+        db: impl DatabaseRef,
     ) -> Result<Option<BuilderTransactionCtx>, BuilderTransactionError> {
         match self.signer {
             Some(signer) => {
@@ -380,15 +439,12 @@ impl<ExtraCtx: Debug + Default> BuilderTxBase<ExtraCtx> {
     fn signed_builder_tx(
         &self,
         ctx: &OpPayloadBuilderCtx<ExtraCtx>,
-        db: &mut State<impl Database>,
+        db: impl DatabaseRef,
         signer: Signer,
         gas_used: u64,
         message: Vec<u8>,
     ) -> Result<Recovered<OpTransactionSigned>, BuilderTransactionError> {
-        let nonce = db
-            .load_cache_account(signer.address)
-            .map(|acc| acc.account_info().unwrap_or_default().nonce)
-            .map_err(|_| BuilderTransactionError::AccountLoadFailed(signer.address))?;
+        let nonce = get_nonce(db, signer.address)?;
 
         // Create the EIP-1559 transaction
         let tx = OpTypedTransaction::Eip1559(TxEip1559 {
@@ -411,20 +467,17 @@ impl<ExtraCtx: Debug + Default> BuilderTxBase<ExtraCtx> {
     }
 }
 
-pub fn get_nonce(
-    db: &mut State<impl Database>,
-    address: Address,
-) -> Result<u64, BuilderTransactionError> {
-    db.load_cache_account(address)
-        .map(|acc| acc.account_info().unwrap_or_default().nonce)
+pub fn get_nonce(db: impl DatabaseRef, address: Address) -> Result<u64, BuilderTransactionError> {
+    db.basic_ref(address)
+        .map(|acc| acc.unwrap_or_default().nonce)
         .map_err(|_| BuilderTransactionError::AccountLoadFailed(address))
 }
 
 pub fn get_balance(
-    db: &mut State<impl Database>,
+    db: impl DatabaseRef,
     address: Address,
 ) -> Result<U256, BuilderTransactionError> {
-    db.load_cache_account(address)
-        .map(|acc| acc.account_info().unwrap_or_default().balance)
+    db.basic_ref(address)
+        .map(|acc| acc.unwrap_or_default().balance)
         .map_err(|_| BuilderTransactionError::AccountLoadFailed(address))
 }