fix: Respect --override-* flags on cast call with --trace flag

Author: Yen

Cargo.lock

@@ -262,6 +262,7 @@ impl CallArgs {
                 trace_mode,
                 odyssey,
                 create2_deployer,
+                state_overrides,
             )?;
 
             let value = tx.value.unwrap_or_default();

crates/cast/src/cmd/call.rs

@@ -186,6 +186,7 @@ impl RunArgs {
             trace_mode,
             odyssey,
             create2_deployer,
+            None,
         )?;
         let mut env = Env::new_with_spec_id(
             env.evm_env.cfg_env.clone(),

crates/cast/src/cmd/run.rs

@@ -32,6 +32,7 @@ alloy-primitives = { workspace = true, features = [
     "arbitrary",
     "rlp",
 ] }
+alloy-rpc-types.workspace = true
 alloy-sol-types.workspace = true
 revm = { workspace = true, default-features = false, features = [
     "std",

crates/evm/evm/Cargo.toml

@@ -258,6 +258,35 @@ impl Executor {
         Ok(self.backend().basic_ref(address)?.map(|acc| acc.nonce).unwrap_or_default())
     }
 
+    /// Set the code of an account.
+    pub fn set_code(&mut self, address: Address, code: Bytecode) -> BackendResult<()> {
+        let mut account = self.backend().basic_ref(address)?.unwrap_or_default();
+        account.code = Some(code);
+        self.backend_mut().insert_account_info(address, account);
+        Ok(())
+    }
+
+    /// Set the storage of an account.
+    pub fn set_storage(
+        &mut self,
+        address: Address,
+        storage: HashMap<U256, U256>,
+    ) -> BackendResult<()> {
+        self.backend_mut().replace_account_storage(address, storage)?;
+        Ok(())
+    }
+
+    /// Set a storage slot of an account.
+    pub fn set_storage_slot(
+        &mut self,
+        address: Address,
+        slot: U256,
+        value: U256,
+    ) -> BackendResult<()> {
+        self.backend_mut().insert_account_storage(address, slot, value)?;
+        Ok(())
+    }
+
     /// Returns `true` if the account has no code.
     pub fn is_empty_code(&self, address: Address) -> BackendResult<bool> {
         Ok(self.backend().basic_ref(address)?.map(|acc| acc.is_empty_code_hash()).unwrap_or(true))

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

@@ -2,12 +2,14 @@ use crate::{
     executors::{Executor, ExecutorBuilder},
     Env,
 };
-use alloy_primitives::Address;
+use alloy_primitives::{map::HashMap, Address, U256};
+use alloy_rpc_types::state::StateOverride;
+use eyre::Context;
 use foundry_compilers::artifacts::EvmVersion;
 use foundry_config::{utils::evm_spec_id, Chain, Config};
 use foundry_evm_core::{backend::Backend, fork::CreateFork, opts::EvmOpts};
 use foundry_evm_traces::TraceMode;
-use revm::primitives::hardfork::SpecId;
+use revm::{primitives::hardfork::SpecId, state::Bytecode};
 use std::ops::{Deref, DerefMut};
 
 /// A default executor with tracing enabled
@@ -23,18 +25,48 @@ impl TracingExecutor {
         trace_mode: TraceMode,
         odyssey: bool,
         create2_deployer: Address,
+        state_overrides: Option<StateOverride>,
     ) -> eyre::Result<Self> {
         let db = Backend::spawn(fork)?;
-        Ok(Self {
-            // configures a bare version of the evm executor: no cheatcode inspector is enabled,
-            // tracing will be enabled only for the targeted transaction
-            executor: ExecutorBuilder::new()
-                .inspectors(|stack| {
-                    stack.trace_mode(trace_mode).odyssey(odyssey).create2_deployer(create2_deployer)
-                })
-                .spec_id(evm_spec_id(version.unwrap_or_default(), odyssey))
-                .build(env, db),
-        })
+        // configures a bare version of the evm executor: no cheatcode inspector is enabled,
+        // tracing will be enabled only for the targeted transaction
+        let mut executor = ExecutorBuilder::new()
+            .inspectors(|stack| {
+                stack.trace_mode(trace_mode).odyssey(odyssey).create2_deployer(create2_deployer)
+            })
+            .spec_id(evm_spec_id(version.unwrap_or_default(), odyssey))
+            .build(env, db);
+
+        // Apply the state overrides.
+        if let Some(state_overrides) = state_overrides {
+            for (address, overrides) in state_overrides {
+                if let Some(balance) = overrides.balance {
+                    executor.set_balance(address, balance)?;
+                }
+                if let Some(nonce) = overrides.nonce {
+                    executor.set_nonce(address, nonce)?;
+                }
+                if let Some(code) = overrides.code {
+                    let bytecode = Bytecode::new_raw_checked(code)
+                        .wrap_err("invalid bytecode in state override")?;
+                    executor.set_code(address, bytecode)?;
+                }
+                if let Some(state) = overrides.state {
+                    let state: HashMap<U256, U256> = state
+                        .into_iter()
+                        .map(|(slot, value)| (slot.into(), value.into()))
+                        .collect();
+                    executor.set_storage(address, state)?;
+                }
+                if let Some(state_diff) = overrides.state_diff {
+                    for (slot, value) in state_diff {
+                        executor.set_storage_slot(address, slot.into(), value.into())?;
+                    }
+                }
+            }
+        }
+
+        Ok(Self { executor })
     }
 
     /// Returns the spec id of the executor

crates/evm/evm/src/executors/trace.rs

@@ -321,6 +321,7 @@ pub async fn get_tracing_executor(
         TraceMode::Call,
         is_odyssey,
         create2_deployer,
+        None,
     )?;
 
     Ok((env, executor))