Skip to content

Commit

Permalink
feat(forge): add vm.lastCallGas cheatcode (#7573)
Browse files Browse the repository at this point in the history
* add `gasUsed and lastGasUsed methods to Vm

* reorder

* basic sketching of idea, planning to use Gas struct to deliver as much gas related information rather than a uint256

* add transformation

* to prevent recording gas by default, only enable after recordGas is enabled

* update struct layout, implementation builds, now connecting to cheatcodes

* fix cheatcodes

* refactor to use simple u64 as Gas struct doesnt have a default

* change from Gas struct to simple u64 as I ran into issues with cache being reset to 0

* it appears cheatcodes are resolved before the actual function calls are therefore it doesnt actually remember the correct value from the previous execution but only of the previous executed cheatcode

* still not working

* finally works, stupid me didnt realize i had to cross call frames

* emit gas record

* test convenient single field access

* add gas record struct back

* pass down isolate bool, only enable gas tracing if enabled

* raise error if cheatcode is used outside of isolation mode

* mark as view

* show gas refund and memory expansion

* improve example

* add isolation test, currently does not run as expected

* fix fmt

* avoid formatting changes

* avoid commiting formatting changes, editor now configured correctly

* lastGasUsed -> lastCallGas

* small name fix

* remove separate isolation profile, just configure on the runner

* fix forge fmt

* note on why path should never happen

* remove separate isolated param, inherit from config

* add support for non-isolation mode

* remove isolate standalone, create additional entry in cheats and document subset of cheats that require to be tested in isolation mode as well

* improve tests, use asserts and add option to exclude contracts from test filter, not just individual tests or paths

* typo, no need to define path exclusion of forks in isolated tests as it is not relevant
  • Loading branch information
zerosnacks authored Apr 8, 2024
1 parent 04e2263 commit b88d167
Show file tree
Hide file tree
Showing 10 changed files with 293 additions and 21 deletions.
51 changes: 51 additions & 0 deletions crates/cheatcodes/assets/cheatcodes.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/cheatcodes/spec/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ impl Cheatcodes<'static> {
Vm::ChainInfo::STRUCT.clone(),
Vm::AccountAccess::STRUCT.clone(),
Vm::StorageAccess::STRUCT.clone(),
Vm::Gas::STRUCT.clone(),
]),
enums: Cow::Owned(vec![
Vm::CallerMode::ENUM.clone(),
Expand Down
52 changes: 36 additions & 16 deletions crates/cheatcodes/spec/src/vm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,19 @@ interface Vm {
address emitter;
}

struct Gas {
// The gas limit of the call.
uint64 gasLimit;
// The total gas used.
uint64 gasTotalUsed;
// The amount of gas used for memory expansion.
uint64 gasMemoryUsed;
// The amount of gas refunded.
int64 gasRefunded;
// The amount of gas remaining.
uint64 gasRemaining;
}

/// An RPC URL and its alias. Returned by `rpcUrlStructs`.
struct Rpc {
/// The alias of the RPC URL.
Expand Down Expand Up @@ -169,6 +182,22 @@ interface Vm {
uint256 chainId;
}

/// The storage accessed during an `AccountAccess`.
struct StorageAccess {
/// The account whose storage was accessed.
address account;
/// The slot that was accessed.
bytes32 slot;
/// If the access was a write.
bool isWrite;
/// The previous value of the slot.
bytes32 previousValue;
/// The new value of the slot.
bytes32 newValue;
/// If the access was reverted.
bool reverted;
}

/// The result of a `stopAndReturnStateDiff` call.
struct AccountAccess {
/// The chain and fork the access occurred.
Expand Down Expand Up @@ -207,22 +236,6 @@ interface Vm {
uint64 depth;
}

/// The storage accessed during an `AccountAccess`.
struct StorageAccess {
/// The account whose storage was accessed.
address account;
/// The slot that was accessed.
bytes32 slot;
/// If the access was a write.
bool isWrite;
/// The previous value of the slot.
bytes32 previousValue;
/// The new value of the slot.
bytes32 newValue;
/// If the access was reverted.
bool reverted;
}

// ======== EVM ========

/// Gets the address for a given private key.
Expand Down Expand Up @@ -594,6 +607,7 @@ interface Vm {
function getRecordedLogs() external returns (Log[] memory logs);

// -------- Gas Metering --------

// It's recommend to use the `noGasMetering` modifier included with forge-std, instead of
// using these functions directly.

Expand All @@ -605,6 +619,12 @@ interface Vm {
#[cheatcode(group = Evm, safety = Safe)]
function resumeGasMetering() external;

// -------- Gas Measurement --------

/// Gets the gas used in the last call.
#[cheatcode(group = Evm, safety = Safe)]
function lastCallGas() external view returns (Gas memory gas);

// ======== Test Assertions and Utilities ========

/// If the condition is false, discard this run's fuzz inputs and generate new ones.
Expand Down
13 changes: 13 additions & 0 deletions crates/cheatcodes/src/evm.rs
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,19 @@ impl Cheatcode for resumeGasMeteringCall {
}
}

impl Cheatcode for lastCallGasCall {
fn apply(&self, state: &mut Cheatcodes) -> Result {
let Self {} = self;
ensure!(state.last_call_gas.is_some(), "`lastCallGas` is only available after a call");
Ok(state
.last_call_gas
.as_ref()
// This should never happen, as we ensure `last_call_gas` is `Some` above.
.expect("`lastCallGas` is only available after a call")
.abi_encode())
}
}

impl Cheatcode for chainIdCall {
fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
let Self { newChainId } = self;
Expand Down
22 changes: 21 additions & 1 deletion crates/cheatcodes/src/inspector.rs
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,10 @@ pub struct Cheatcodes {
/// Recorded logs
pub recorded_logs: Option<Vec<crate::Vm::Log>>,

/// Cache of the amount of gas used in previous call.
/// This is used by the `lastCallGas` cheatcode.
pub last_call_gas: Option<crate::Vm::Gas>,

/// Mocked calls
// **Note**: inner must a BTreeMap because of special `Ord` impl for `MockCallDataContext`
pub mocked_calls: HashMap<Address, BTreeMap<MockCallDataContext, MockCallReturnData>>,
Expand Down Expand Up @@ -1033,9 +1037,25 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
// Exit early for calls to cheatcodes as other logic is not relevant for cheatcode
// invocations
if cheatcode_call {
return outcome;
return outcome
}

// Record the gas usage of the call, this allows the `lastCallGas` cheatcode to
// retrieve the gas usage of the last call.
let gas = outcome.result.gas;
self.last_call_gas = Some(crate::Vm::Gas {
// The gas limit of the call.
gasLimit: gas.limit(),
// The total gas used.
gasTotalUsed: gas.spend(),
// The amount of gas used for memory expansion.
gasMemoryUsed: gas.memory(),
// The amount of gas refunded.
gasRefunded: gas.refunded(),
// The amount of gas remaining.
gasRemaining: gas.remaining(),
});

// If `startStateDiffRecording` has been called, update the `reverted` status of the
// previous call depth's recorded accesses, if any
if let Some(recorded_account_diffs_stack) = &mut self.recorded_account_diffs_stack {
Expand Down
23 changes: 20 additions & 3 deletions crates/forge/tests/it/cheats.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,11 @@ use crate::{
use foundry_config::{fs_permissions::PathPermission, FsPermissions};
use foundry_test_utils::Filter;

/// Executes all cheat code tests but not fork cheat codes
/// Executes all cheat code tests but not fork cheat codes or tests that require isolation mode
async fn test_cheats_local(test_data: &ForgeTestData) {
let mut filter =
Filter::new(".*", ".*", &format!(".*cheats{RE_PATH_SEPARATOR}*")).exclude_paths("Fork");
let mut filter = Filter::new(".*", ".*", &format!(".*cheats{RE_PATH_SEPARATOR}*"))
.exclude_paths("Fork")
.exclude_contracts("Isolated");

// Exclude FFI tests on Windows because no `echo`, and file tests that expect certain file paths
if cfg!(windows) {
Expand All @@ -24,11 +25,27 @@ async fn test_cheats_local(test_data: &ForgeTestData) {
TestConfig::with_filter(runner, filter).run().await;
}

/// Executes subset of all cheat code tests in isolation mode
async fn test_cheats_local_isolated(test_data: &ForgeTestData) {
let filter = Filter::new(".*", ".*(Isolated)", &format!(".*cheats{RE_PATH_SEPARATOR}*"));

let mut config = test_data.config.clone();
config.isolate = true;
let runner = test_data.runner_with_config(config);

TestConfig::with_filter(runner, filter).run().await;
}

#[tokio::test(flavor = "multi_thread")]
async fn test_cheats_local_default() {
test_cheats_local(&TEST_DATA_DEFAULT).await
}

#[tokio::test(flavor = "multi_thread")]
async fn test_cheats_local_default_isolated() {
test_cheats_local_isolated(&TEST_DATA_DEFAULT).await
}

#[tokio::test(flavor = "multi_thread")]
async fn test_cheats_local_multi_version() {
test_cheats_local(&TEST_DATA_MULTI_VERSION).await
Expand Down
8 changes: 7 additions & 1 deletion crates/forge/tests/it/test_helpers.rs
Original file line number Diff line number Diff line change
Expand Up @@ -203,7 +203,12 @@ impl ForgeTestData {
config.prompt_timeout = 0;

let root = self.project.root();
let opts = self.evm_opts.clone();
let mut opts = self.evm_opts.clone();

if config.isolate {
opts.isolate = true;
}

let env = opts.local_evm_env();
let output = self.output.clone();
let artifact_ids = output.artifact_ids().map(|(id, _)| id).collect();
Expand All @@ -215,6 +220,7 @@ impl ForgeTestData {
None,
None,
))
.enable_isolation(opts.isolate)
.sender(config.sender)
.with_test_options(self.test_opts.clone())
.build(root, output, env, opts.clone())
Expand Down
17 changes: 17 additions & 0 deletions crates/test-utils/src/filter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ pub struct Filter {
contract_regex: Regex,
path_regex: Regex,
exclude_tests: Option<Regex>,
exclude_contracts: Option<Regex>,
exclude_paths: Option<Regex>,
}

Expand All @@ -21,6 +22,7 @@ impl Filter {
path_regex: Regex::new(path_pattern)
.unwrap_or_else(|_| panic!("Failed to parse path pattern: `{path_pattern}`")),
exclude_tests: None,
exclude_contracts: None,
exclude_paths: None,
}
}
Expand All @@ -41,6 +43,14 @@ impl Filter {
self
}

/// All contracts to also exclude
///
/// This is a workaround since regex does not support negative look aheads
pub fn exclude_contracts(mut self, pattern: &str) -> Self {
self.exclude_contracts = Some(Regex::new(pattern).unwrap());
self
}

/// All paths to also exclude
///
/// This is a workaround since regex does not support negative look aheads
Expand All @@ -55,6 +65,7 @@ impl Filter {
contract_regex: Regex::new(".*").unwrap(),
path_regex: Regex::new(".*").unwrap(),
exclude_tests: None,
exclude_contracts: None,
exclude_paths: None,
}
}
Expand All @@ -71,6 +82,12 @@ impl TestFilter for Filter {
}

fn matches_contract(&self, contract_name: &str) -> bool {
if let Some(exclude) = &self.exclude_contracts {
if exclude.is_match(contract_name) {
return false;
}
}

self.contract_regex.is_match(contract_name)
}

Expand Down
2 changes: 2 additions & 0 deletions testdata/cheats/Vm.sol

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit b88d167

Please sign in to comment.