feat: reintroduce backend.inspect (#802)

* reintroduce backend.inspect

* fix downcast

Author: nbaztec

crates/evm/core/src/backend/cow.rs

@@ -14,10 +14,13 @@ use alloy_rpc_types::TransactionRequest;
 use foundry_fork_db::DatabaseError;
 use revm::{
     db::DatabaseRef,
-    primitives::{Account, AccountInfo, Bytecode, Env, EnvWithHandlerCfg, HashMap as Map, SpecId},
+    primitives::{
+        Account, AccountInfo, Bytecode, Env, EnvWithHandlerCfg, HashMap as Map, ResultAndState,
+        SpecId,
+    },
     Database, DatabaseCommit, JournaledState,
 };
-use std::{borrow::Cow, collections::BTreeMap};
+use std::{any::Any, borrow::Cow, collections::BTreeMap};
 
 /// A wrapper around `Backend` that ensures only `revm::DatabaseRef` functions are called.
 ///
@@ -53,6 +56,30 @@ impl<'a> CowBackend<'a> {
         Self { backend: Cow::Borrowed(backend), is_initialized: false, spec_id: SpecId::LATEST }
     }
 
+    /// Executes the configured transaction of the `env` without committing state changes
+    ///
+    /// Note: in case there are any cheatcodes executed that modify the environment, this will
+    /// update the given `env` with the new values.
+    #[instrument(name = "inspect", level = "debug", skip_all)]
+    pub fn inspect<I: InspectorExt>(
+        &mut self,
+        env: &mut EnvWithHandlerCfg,
+        inspector: &mut I,
+        inspect_ctx: Box<dyn Any>,
+    ) -> eyre::Result<ResultAndState> {
+        // this is a new call to inspect with a new env, so even if we've cloned the backend
+        // already, we reset the initialized state
+        self.is_initialized = false;
+        self.spec_id = env.handler_cfg.spec_id;
+
+        self.backend.strategy.runner.clone().inspect(
+            self.backend.to_mut(),
+            env,
+            inspector,
+            inspect_ctx,
+        )
+    }
+
     pub fn new_borrowed(backend: &'a Backend) -> Self {
         Self { backend: Cow::Borrowed(backend), is_initialized: false, spec_id: SpecId::LATEST }
     }

crates/evm/core/src/backend/mod.rs

@@ -22,11 +22,12 @@ use revm::{
     precompile::{PrecompileSpecId, Precompiles},
     primitives::{
         Account, AccountInfo, BlobExcessGasAndPrice, Bytecode, Env, EnvWithHandlerCfg, EvmState,
-        EvmStorageSlot, HashMap as Map, Log, SpecId, KECCAK_EMPTY,
+        EvmStorageSlot, HashMap as Map, Log, ResultAndState, SpecId, KECCAK_EMPTY,
     },
     Database, DatabaseCommit, JournaledState,
 };
 use std::{
+    any::Any,
     collections::{BTreeMap, HashSet},
     time::Instant,
 };
@@ -773,11 +774,26 @@ impl Backend {
     /// Initializes settings we need to keep track of.
     ///
     /// We need to track these mainly to prevent issues when switching between different evms
-    pub fn initialize(&mut self, env: &EnvWithHandlerCfg) {
+    pub(crate) fn initialize(&mut self, env: &EnvWithHandlerCfg) {
         self.set_caller(env.tx.caller);
         self.set_spec_id(env.handler_cfg.spec_id);
     }
 
+    /// Executes the configured test call of the `env` without committing state changes.
+    ///
+    /// Note: in case there are any cheatcodes executed that modify the environment, this will
+    /// update the given `env` with the new values.
+    #[instrument(name = "inspect", level = "debug", skip_all)]
+    pub fn inspect<I: InspectorExt>(
+        &mut self,
+        env: &mut EnvWithHandlerCfg,
+        inspector: &mut I,
+        inspect_ctx: Box<dyn Any>,
+    ) -> eyre::Result<ResultAndState> {
+        self.initialize(env);
+        self.strategy.runner.clone().inspect(self, env, inspector, inspect_ctx)
+    }
+
     /// Returns the `EnvWithHandlerCfg` with the current `spec_id` set.
     fn env_with_handler_cfg(&self, env: Env) -> EnvWithHandlerCfg {
         EnvWithHandlerCfg::new_with_spec_id(Box::new(env), self.inner.spec_id)

crates/evm/core/src/backend/strategy.rs

@@ -1,8 +1,15 @@
 use std::{any::Any, fmt::Debug};
 
-use super::{BackendInner, Fork, ForkDB, ForkType, FoundryEvmInMemoryDB};
+use crate::InspectorExt;
+
+use super::{Backend, BackendInner, Fork, ForkDB, ForkType, FoundryEvmInMemoryDB};
 use alloy_primitives::{Address, U256};
-use revm::{db::CacheDB, primitives::HashSet, DatabaseRef, JournaledState};
+use eyre::{Context, Result};
+use revm::{
+    db::CacheDB,
+    primitives::{EnvWithHandlerCfg, HashSet, ResultAndState},
+    DatabaseRef, JournaledState,
+};
 use serde::{Deserialize, Serialize};
 
 pub struct BackendStrategyForkInfo<'a> {
@@ -62,6 +69,14 @@ pub trait BackendStrategyRunner: Debug + Send + Sync + BackendStrategyRunnerExt
 
     fn new_cloned(&self) -> Box<dyn BackendStrategyRunner>;
 
+    fn inspect(
+        &self,
+        backend: &mut Backend,
+        env: &mut EnvWithHandlerCfg,
+        inspector: &mut dyn InspectorExt,
+        inspect_ctx: Box<dyn Any>,
+    ) -> Result<ResultAndState>;
+
     /// When creating or switching forks, we update the AccountInfo of the contract
     fn update_fork_db(
         &self,
@@ -121,6 +136,22 @@ impl BackendStrategyRunner for EvmBackendStrategyRunner {
         Box::new(self.clone())
     }
 
+    fn inspect(
+        &self,
+        backend: &mut Backend,
+        env: &mut EnvWithHandlerCfg,
+        inspector: &mut dyn InspectorExt,
+        _inspect_ctx: Box<dyn Any>,
+    ) -> Result<ResultAndState> {
+        let mut evm = crate::utils::new_evm_with_inspector(backend, env.clone(), inspector);
+
+        let res = evm.transact().wrap_err("backend: failed while inspecting")?;
+
+        env.env = evm.context.evm.inner.env;
+
+        Ok(res)
+    }
+
     fn update_fork_db(
         &self,
         _ctx: &mut dyn BackendStrategyContext,

crates/evm/evm/src/executors/mod.rs

@@ -15,7 +15,6 @@ use alloy_primitives::{
     map::{AddressHashMap, HashMap},
     Address, Bytes, Log, U256,
 };
-use alloy_serde::OtherFields;
 use alloy_sol_types::{sol, SolCall};
 use foundry_evm_core::{
     backend::{Backend, BackendError, BackendResult, CowBackend, DatabaseExt, GLOBAL_FAIL_SLOT},
@@ -93,7 +92,7 @@ pub struct Executor {
     /// Whether `failed()` should be called on the test contract to determine if the test failed.
     legacy_assertions: bool,
 
-    strategy: ExecutorStrategy,
+    pub strategy: ExecutorStrategy,
 }
 
 impl Clone for Executor {
@@ -258,11 +257,6 @@ impl Executor {
         self
     }
 
-    #[inline]
-    pub fn set_transaction_other_fields(&mut self, other_fields: OtherFields) {
-        self.strategy.runner.set_inspect_context(self.strategy.context.as_mut(), other_fields);
-    }
-
     /// Deploys a contract and commits the new state to the underlying database.
     ///
     /// Executes a CREATE transaction with the contract `code` and persistent database state
@@ -438,15 +432,12 @@ impl Executor {
     pub fn call_with_env(&self, mut env: EnvWithHandlerCfg) -> eyre::Result<RawCallResult> {
         let mut inspector = self.inspector().clone();
         let mut backend = CowBackend::new_borrowed(self.backend());
-        // this is a new call to inspect with a new env, so even if we've cloned the backend
-        // already, we reset the initialized state
-        backend.is_initialized = false;
-        backend.spec_id = env.spec_id();
 
-        let result = self.strategy.runner.call_inspect(
+        let result = self.strategy.runner.call(
             self.strategy.context.as_ref(),
             &mut backend,
             &mut env,
+            &self.env,
             &mut inspector,
         )?;
 
@@ -463,9 +454,8 @@ impl Executor {
     pub fn transact_with_env(&mut self, mut env: EnvWithHandlerCfg) -> eyre::Result<RawCallResult> {
         let mut inspector = self.inspector.clone();
         let backend = &mut self.backend;
-        backend.initialize(&env);
 
-        let result_and_state = self.strategy.runner.transact_inspect(
+        let result_and_state = self.strategy.runner.transact(
             self.strategy.context.as_mut(),
             backend,
             &mut env,

crates/evm/evm/src/executors/strategy.rs

@@ -2,20 +2,19 @@ use std::{any::Any, fmt::Debug};
 
 use alloy_primitives::{Address, U256};
 use alloy_serde::OtherFields;
-use eyre::{Context, Result};
+use eyre::Result;
 use foundry_cheatcodes::strategy::{
     CheatcodeInspectorStrategy, EvmCheatcodeInspectorStrategyRunner,
 };
-use foundry_evm_core::{
-    backend::{strategy::BackendStrategy, BackendResult, DatabaseExt},
-    InspectorExt,
-};
+use foundry_evm_core::backend::{strategy::BackendStrategy, Backend, BackendResult, CowBackend};
 use foundry_zksync_compiler::DualCompiledContracts;
 use revm::{
     primitives::{Env, EnvWithHandlerCfg, ResultAndState},
     DatabaseRef,
 };
 
+use crate::inspectors::InspectorStack;
+
 use super::Executor;
 
 pub trait ExecutorStrategyContext: Debug + Send + Sync + Any {
@@ -76,24 +75,25 @@ pub trait ExecutorStrategyRunner: Debug + Send + Sync + ExecutorStrategyExt {
     fn set_nonce(&self, executor: &mut Executor, address: Address, nonce: u64)
         -> BackendResult<()>;
 
-    fn set_inspect_context(&self, ctx: &mut dyn ExecutorStrategyContext, other_fields: OtherFields);
-
-    fn call_inspect(
+    /// Execute a transaction and *WITHOUT* applying state changes.
+    fn call(
         &self,
         ctx: &dyn ExecutorStrategyContext,
-        db: &mut dyn DatabaseExt,
+        backend: &mut CowBackend<'_>,
         env: &mut EnvWithHandlerCfg,
-        inspector: &mut dyn InspectorExt,
-    ) -> eyre::Result<ResultAndState>;
+        executor_env: &EnvWithHandlerCfg,
+        inspector: &mut InspectorStack,
+    ) -> Result<ResultAndState>;
 
-    fn transact_inspect(
+    /// Execute a transaction and apply state changes.
+    fn transact(
         &self,
         ctx: &mut dyn ExecutorStrategyContext,
-        db: &mut dyn DatabaseExt,
+        backend: &mut Backend,
         env: &mut EnvWithHandlerCfg,
-        _executor_env: &EnvWithHandlerCfg,
-        inspector: &mut dyn InspectorExt,
-    ) -> eyre::Result<ResultAndState>;
+        executor_env: &EnvWithHandlerCfg,
+        inspector: &mut InspectorStack,
+    ) -> Result<ResultAndState>;
 
     fn new_backend_strategy(&self) -> BackendStrategy;
     fn new_cheatcode_inspector_strategy(
@@ -123,6 +123,19 @@ pub trait ExecutorStrategyExt {
     ) -> Result<()> {
         Ok(())
     }
+
+    /// Sets the transaction context for the next [ExecutorStrategyRunner::call] or
+    /// [ExecutorStrategyRunner::transact]. This selects whether to run the transaction on zkEVM
+    /// or the EVM.
+    /// This is based if the [OtherFields] contains
+    /// [foundry_zksync_core::ZKSYNC_TRANSACTION_OTHER_FIELDS_KEY] with
+    /// [foundry_zksync_core::ZkTransactionMetadata].
+    fn zksync_set_transaction_context(
+        &self,
+        _ctx: &mut dyn ExecutorStrategyContext,
+        _other_fields: OtherFields,
+    ) {
+    }
 }
 
 /// Implements [ExecutorStrategyRunner] for EVM.
@@ -138,55 +151,6 @@ impl ExecutorStrategyRunner for EvmExecutorStrategyRunner {
         Box::new(self.clone())
     }
 
-    fn set_inspect_context(
-        &self,
-        _ctx: &mut dyn ExecutorStrategyContext,
-        _other_fields: OtherFields,
-    ) {
-    }
-
-    /// Executes the configured test call of the `env` without committing state changes.
-    ///
-    /// Note: in case there are any cheatcodes executed that modify the environment, this will
-    /// update the given `env` with the new values.
-    fn call_inspect(
-        &self,
-        _ctx: &dyn ExecutorStrategyContext,
-        db: &mut dyn DatabaseExt,
-        env: &mut EnvWithHandlerCfg,
-        inspector: &mut dyn InspectorExt,
-    ) -> eyre::Result<ResultAndState> {
-        let mut evm = crate::utils::new_evm_with_inspector(db, env.clone(), inspector);
-
-        let res = evm.transact().wrap_err("backend: failed while inspecting")?;
-
-        env.env = evm.context.evm.inner.env;
-
-        Ok(res)
-    }
-
-    /// Executes the configured test call of the `env` without committing state changes.
-    /// Modifications to the state are however allowed.
-    ///
-    /// Note: in case there are any cheatcodes executed that modify the environment, this will
-    /// update the given `env` with the new values.
-    fn transact_inspect(
-        &self,
-        _ctx: &mut dyn ExecutorStrategyContext,
-        db: &mut dyn DatabaseExt,
-        env: &mut EnvWithHandlerCfg,
-        _executor_env: &EnvWithHandlerCfg,
-        inspector: &mut dyn InspectorExt,
-    ) -> eyre::Result<ResultAndState> {
-        let mut evm = crate::utils::new_evm_with_inspector(db, env.clone(), inspector);
-
-        let res = evm.transact().wrap_err("backend: failed while inspecting")?;
-
-        env.env = evm.context.evm.inner.env;
-
-        Ok(res)
-    }
-
     fn set_balance(
         &self,
         executor: &mut Executor,
@@ -214,6 +178,28 @@ impl ExecutorStrategyRunner for EvmExecutorStrategyRunner {
         Ok(())
     }
 
+    fn call(
+        &self,
+        _ctx: &dyn ExecutorStrategyContext,
+        backend: &mut CowBackend<'_>,
+        env: &mut EnvWithHandlerCfg,
+        _executor_env: &EnvWithHandlerCfg,
+        inspector: &mut InspectorStack,
+    ) -> Result<ResultAndState> {
+        backend.inspect(env, inspector, Box::new(()))
+    }
+
+    fn transact(
+        &self,
+        _ctx: &mut dyn ExecutorStrategyContext,
+        backend: &mut Backend,
+        env: &mut EnvWithHandlerCfg,
+        _executor_env: &EnvWithHandlerCfg,
+        inspector: &mut InspectorStack,
+    ) -> Result<ResultAndState> {
+        backend.inspect(env, inspector, Box::new(()))
+    }
+
     fn new_backend_strategy(&self) -> BackendStrategy {
         BackendStrategy::new_evm()
     }

crates/script/src/broadcast.rs

@@ -73,7 +73,7 @@ pub async fn send_transaction(
 ) -> Result<TxHash> {
     let zk_tx_meta =
         if let SendTransactionKind::Raw(tx, _) | SendTransactionKind::Unlocked(tx) = &mut kind {
-            foundry_strategy_zksync::get_zksync_transaction_metadata(&tx.other)
+            foundry_strategy_zksync::try_get_zksync_transaction_metadata(&tx.other)
         } else {
             None
         };

crates/script/src/runner.rs

@@ -257,7 +257,10 @@ impl ScriptRunner {
         other_fields: Option<OtherFields>,
     ) -> Result<ScriptResult> {
         if let Some(other_fields) = other_fields {
-            self.executor.set_transaction_other_fields(other_fields);
+            self.executor.strategy.runner.zksync_set_transaction_context(
+                self.executor.strategy.context.as_mut(),
+                other_fields,
+            );
         }
 
         if let Some(to) = to {