REVM: Support cheatcodes in setUp (#997)

* fix: support cheatcodes in `setUp`

* fix: subtract stipend without panic

* chore: rename test

* fix: set tx gas price to block basefee

* fix: use `CALLER` for `is_success` check

* chore: remove duplicate clap attribute

* fix: set chain id correctly in fork mode

* fix: separate evm block env from execution env

* chore: clippy

* refactor: block override without `block_env` fn

cli/src/cmd/forge/run.rs

@@ -433,7 +433,7 @@ impl<DB: DatabaseRef> Runner<DB> {
  self.executor.call_raw(self.sender, address, calldata.0, 0.into())?;
  Ok(RunResult {
  success: !reverted,
- gas: gas - stipend,
+ gas: gas.overflowing_sub(stipend).0,
  logs,
  traces: traces.map(|traces| vec![(TraceKind::Execution, traces)]).unwrap_or_default(),
  debug: vec![debug].into_iter().collect(),

cli/src/opts/evm.rs

@@ -39,7 +39,6 @@ pub struct EvmArgs {
  ///
  /// If you want to fetch state from a specific block number, see --fork-block-number.
  #[clap(long, short, alias = "rpc-url")]
- #[clap(alias = "rpc-url")]
  #[serde(rename = "eth_rpc_url", skip_serializing_if = "Option::is_none")]
  pub fork_url: Option<String>,
 

evm/src/executor/builder.rs

@@ -122,6 +122,7 @@ impl ExecutorBuilder {
  /// Configure the execution environment (gas limit, chain spec, ...)
  #[must_use]
  pub fn with_config(mut self, env: Env) -> Self {
+ self.inspector_config.block = env.block.clone();
  self.env = env;
  self
  }

evm/src/executor/fork/init.rs

@@ -1,5 +1,5 @@
 use ethers::{providers::Middleware, types::Address};
-use revm::{BlockEnv, Env, TxEnv};
+use revm::{BlockEnv, CfgEnv, Env, TxEnv};
 
 /// Initializes a REVM block environment based on a forked
 /// ethereum provider.
@@ -22,6 +22,10 @@ pub async fn environment<M: Middleware>(
  let block = block.expect("block not found");
 
  Ok(Env {
+ cfg: CfgEnv {
+ chain_id: override_chain_id.unwrap_or(rpc_chain_id.as_u64()).into(),
+ ..Default::default()
+ },
  block: BlockEnv {
  number: block.number.expect("block number not found").as_u64().into(),
  timestamp: block.timestamp,
@@ -37,6 +41,5 @@ pub async fn environment<M: Middleware>(
  gas_limit: block.gas_limit.as_u64(),
  ..Default::default()
  },
- ..Default::default()
  })
 }

evm/src/executor/inspector/cheatcodes/mod.rs

@@ -19,7 +19,8 @@ use ethers::{
  types::{Address, H256},
 };
 use revm::{
- opcode, CallInputs, CreateInputs, Database, EVMData, Gas, Inspector, Interpreter, Return,
+ opcode, BlockEnv, CallInputs, CreateInputs, Database, EVMData, Gas, Inspector, Interpreter,
+ Return,
 };
 use std::collections::BTreeMap;
 
@@ -32,6 +33,12 @@ pub struct Cheatcodes {
  /// Whether FFI is enabled or not
  ffi: bool,
 
+ /// The block environment
+ ///
+ /// Used in the cheatcode handler to overwrite the block environment separately from the
+ /// execution block environment.
+ pub block: Option<BlockEnv>,
+
  /// Address labels
  pub labels: BTreeMap<Address, String>,
 
@@ -55,8 +62,8 @@ pub struct Cheatcodes {
 }
 
 impl Cheatcodes {
- pub fn new(ffi: bool) -> Self {
- Self { ffi, ..Default::default() }
+ pub fn new(ffi: bool, block: BlockEnv) -> Self {
+ Self { ffi, block: Some(block), ..Default::default() }
  }
 
  fn apply_cheatcode<DB: Database>(
@@ -136,6 +143,21 @@ where
  }
  }
 
+ fn initialize_interp(
+ &mut self,
+ _: &mut Interpreter,
+ data: &mut EVMData<'_, DB>,
+ _: bool,
+ ) -> Return {
+ // When the first interpreter is initialized we've circumvented the balance and gas checks,
+ // so we apply our actual block data with the correct fees and all.
+ if let Some(block) = self.block.take() {
+ data.env.block = block;
+ }
+
+ Return::Continue
+ }
+
  fn step(&mut self, interpreter: &mut Interpreter, _: &mut EVMData<'_, DB>, _: bool) -> Return {
  // Record writes and reads if `record` has been called
  if let Some(storage_accesses) = &mut self.accesses {

evm/src/executor/inspector/mod.rs

@@ -16,12 +16,19 @@ pub use stack::InspectorStack;
 mod cheatcodes;
 pub use cheatcodes::Cheatcodes;
 
+use revm::BlockEnv;
+
 #[derive(Default, Clone, Debug)]
 pub struct InspectorStackConfig {
  /// Whether or not cheatcodes are enabled
  pub cheatcodes: bool,
  /// Whether or not the FFI cheatcode is enabled
  pub ffi: bool,
+ /// The block environment
+ ///
+ /// Used in the cheatcode handler to overwrite the block environment separately from the
+ /// execution block environment.
+ pub block: BlockEnv,
  /// Whether or not tracing is enabled
  pub tracing: bool,
  /// Whether or not the debugger is enabled
@@ -34,7 +41,7 @@ impl InspectorStackConfig {
 
  stack.logs = Some(LogCollector::new());
  if self.cheatcodes {
- stack.cheatcodes = Some(Cheatcodes::new(self.ffi));
+ stack.cheatcodes = Some(Cheatcodes::new(self.ffi, self.block.clone()));
  }
  if self.tracing {
  stack.tracer = Some(Tracer::new());

evm/src/executor/mod.rs

@@ -37,7 +37,7 @@ use hashbrown::HashMap;
 use inspector::InspectorStack;
 use revm::{
  db::{CacheDB, DatabaseCommit, EmptyDB},
- return_ok, Account, CreateScheme, Env, Return, TransactOut, TransactTo, TxEnv, EVM,
+ return_ok, Account, BlockEnv, CreateScheme, Env, Return, TransactOut, TransactTo, TxEnv, EVM,
 };
 use std::collections::BTreeMap;
 
@@ -57,7 +57,7 @@ pub enum EvmError {
  traces: Option<CallTraceArena>,
  debug: Option<DebugArena>,
  labels: BTreeMap<Address, String>,
- state_changeset: StateChangeset,
+ state_changeset: Option<StateChangeset>,
  },
  /// Error which occurred during ABI encoding/decoding
  #[error(transparent)]
@@ -105,7 +105,7 @@ pub struct CallResult<D: Detokenize> {
  ///
  /// This is only present if the changed state was not committed to the database (i.e. if you
  /// used `call` and `call_raw` not `call_committing` or `call_raw_committing`).
- pub state_changeset: StateChangeset,
+ pub state_changeset: Option<StateChangeset>,
 }
 
 /// The result of a raw call.
@@ -133,7 +133,7 @@ pub struct RawCallResult {
  ///
  /// This is only present if the changed state was not committed to the database (i.e. if you
  /// used `call` and `call_raw` not `call_committing` or `call_raw_committing`).
- pub state_changeset: StateChangeset,
+ pub state_changeset: Option<StateChangeset>,
 }
 
 impl Default for RawCallResult {
@@ -148,7 +148,7 @@ impl Default for RawCallResult {
  labels: BTreeMap::new(),
  traces: None,
  debug: None,
- state_changeset: StateChangeset::new(),
+ state_changeset: None,
  }
  }
 }
@@ -279,9 +279,37 @@ where
  calldata: Bytes,
  value: U256,
  ) -> Result<RawCallResult> {
- let result = self.call_raw(from, to, calldata, value)?;
- self.db.commit(result.state_changeset.clone());
- Ok(result)
+ let stipend = stipend(&calldata, self.env.cfg.spec_id);
+
+ // Build VM
+ let mut evm = EVM::new();
+ evm.env = self.build_env(from, TransactTo::Call(to), calldata, value);
+ evm.database(&mut self.db);
+
+ // Run the call
+ let mut inspector = self.inspector_config.stack();
+ let (status, out, gas, _) = evm.inspect_commit(&mut inspector);
+ let result = match out {
+ TransactOut::Call(data) => data,
+ _ => Bytes::default(),
+ };
+
+ // Persist the changed block environment
+ self.inspector_config.block = evm.env.block.clone();
+
+ let InspectorData { logs, labels, traces, debug } = collect_inspector_states(inspector);
+ Ok(RawCallResult {
+ status,
+ reverted: !matches!(status, return_ok!()),
+ result,
+ gas,
+ stipend,
+ logs: logs.to_vec(),
+ labels,
+ traces,
+ debug,
+ state_changeset: None,
+ })
  }
 
  /// Performs a call to an account on the current state of the VM.
@@ -379,7 +407,7 @@ where
  labels,
  traces,
  debug,
- state_changeset,
+ state_changeset: Some(state_changeset),
  })
  }
 
@@ -426,14 +454,8 @@ where
  let mut success = !reverted;
  if success {
  // Check if a DSTest assertion failed
- let call = executor.call::<bool, _, _>(
- Address::zero(),
- address,
- "failed()(bool)",
- (),
- 0.into(),
- None,
- );
+ let call =
+ executor.call::<bool, _, _>(*CALLER, address, "failed()(bool)", (), 0.into(), None);
 
  if let Ok(CallResult { result: failed, .. }) = call {
  success = !failed;
@@ -446,8 +468,20 @@ where
  fn build_env(&self, caller: Address, transact_to: TransactTo, data: Bytes, value: U256) -> Env {
  Env {
  cfg: self.env.cfg.clone(),
- block: self.env.block.clone(),
- tx: TxEnv { caller, transact_to, data, value, ..self.env.tx.clone() },
+ // We always set the gas price to 0 so we can execute the transaction regardless of
+ // network conditions - the actual gas price is kept in `self.block` and is applied by
+ // the cheatcode handler if it is enabled
+ block: BlockEnv { basefee: 0.into(), ..self.env.block.clone() },
+ tx: TxEnv {
+ caller,
+ transact_to,
+ data,
+ value,
+ // As above, we set the gas price to 0.
+ gas_price: 0.into(),
+ gas_priority_fee: None,
+ ..self.env.tx.clone()
+ },
  }
  }
 }

evm/src/fuzz/mod.rs

@@ -79,9 +79,11 @@ where
  .call_raw(self.sender, address, calldata.0.clone(), 0.into())
  .expect("could not make raw evm call");
  let call = call.borrow();
+ let state_changeset =
+ call.state_changeset.as_ref().expect("we should have a state changeset");
 
  // Build fuzzer state
- collect_state_from_call(&call.logs, &call.state_changeset, state.clone());
+ collect_state_from_call(&call.logs, state_changeset, state.clone());
 
  // When assume cheat code is triggered return a special string "FOUNDRY::ASSUME"
  if call.result.as_ref() == ASSUME_MAGIC_RETURN_CODE {
@@ -91,7 +93,7 @@ where
  let success = self.executor.is_success(
  address,
  call.reverted,
- call.state_changeset.clone(),
+ state_changeset.clone(),
  should_fail,
  );
 

forge/src/runner.rs

@@ -343,8 +343,12 @@ impl<'a, DB: DatabaseRef + Send + Sync> ContractRunner<'a, DB> {
  };
  traces.extend(execution_traces.map(|traces| (TraceKind::Execution, traces)).into_iter());
 
- let success =
- self.executor.is_success(setup.address, reverted, state_changeset, should_fail);
+ let success = self.executor.is_success(
+ setup.address,
+ reverted,
+ state_changeset.expect("we should have a state changeset"),
+ should_fail,
+ );
 
  // Record test execution time
  tracing::debug!(
@@ -358,7 +362,7 @@ impl<'a, DB: DatabaseRef + Send + Sync> ContractRunner<'a, DB> {
  reason,
  counterexample: None,
  logs,
- kind: TestKind::Standard(gas - stipend),
+ kind: TestKind::Standard(gas.overflowing_sub(stipend).0),
  traces,
  labeled_addresses,
  })

testdata/cheats/Setup.t.sol

@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: Unlicense
+pragma solidity >=0.8.0;
+
+import "ds-test/test.sol";
+import "./Cheats.sol";
+
+contract CheatsSetupTest is DSTest {
+ Cheats constant cheats = Cheats(HEVM_ADDRESS);
+
+ function setUp() public {
+ cheats.warp(10);
+ cheats.roll(100);
+ cheats.fee(1000);
+ }
+
+ function testCheatEnvironment() public {
+ assertEq(block.timestamp, 10, "block timestamp was not persisted from setup");
+ assertEq(block.number, 100, "block number was not persisted from setup");
+ assertEq(block.basefee, 1000, "basefee was not persisted from setup");
+ }
+}