diff --git a/.gitignore b/.gitignore index 19f666e451bf9..5b61e32022994 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ .DS_STORE /target out/ +snapshots/ out.json .idea .vscode diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index 5da7ed0dc80ea..90525f81b15d5 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -8423,6 +8423,46 @@ }, "safety": "unsafe" }, + { + "func": { + "id": "snapshotGasLastCall_0", + "description": "Snapshot capture the gas usage of the last call by name.", + "declaration": "function snapshotGasLastCall(string calldata name) external;", + "visibility": "external", + "mutability": "", + "signature": "snapshotGasLastCall(string)", + "selector": "0xdd9fca12", + "selectorBytes": [ + 221, + 159, + 202, + 18 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, + { + "func": { + "id": "snapshotGasLastCall_1", + "description": "Snapshot capture the gas usage of the last call by name in a group.", + "declaration": "function snapshotGasLastCall(string calldata group, string calldata name) external;", + "visibility": "external", + "mutability": "", + "signature": "snapshotGasLastCall(string,string)", + "selector": "0x200c6772", + "selectorBytes": [ + 32, + 12, + 103, + 114 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, { "func": { "id": "snapshotState", @@ -8443,6 +8483,46 @@ "status": "stable", "safety": "unsafe" }, + { + "func": { + "id": "snapshotValue_0", + "description": "Snapshot capture an arbitrary numerical value by name.\nThe group name is derived from the contract name.", + "declaration": "function snapshotValue(string calldata name, uint256 value) external;", + "visibility": "external", + "mutability": "", + "signature": "snapshotValue(string,uint256)", + "selector": "0x51db805a", + "selectorBytes": [ + 81, + 219, + 128, + 90 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, + { + "func": { + "id": "snapshotValue_1", + "description": "Snapshot capture an arbitrary numerical value by name in a group.", + "declaration": "function snapshotValue(string calldata group, string calldata name, uint256 value) external;", + "visibility": "external", + "mutability": "", + "signature": "snapshotValue(string,string,uint256)", + "selector": "0x6d2b27d8", + "selectorBytes": [ + 109, + 43, + 39, + 216 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, { "func": { "id": "split", @@ -8583,6 +8663,46 @@ "status": "stable", "safety": "unsafe" }, + { + "func": { + "id": "startSnapshotGas_0", + "description": "Start a snapshot capture of the current gas usage by name.\nThe group name is derived from the contract name.", + "declaration": "function startSnapshotGas(string calldata name) external;", + "visibility": "external", + "mutability": "", + "signature": "startSnapshotGas(string)", + "selector": "0x3cad9d7b", + "selectorBytes": [ + 60, + 173, + 157, + 123 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, + { + "func": { + "id": "startSnapshotGas_1", + "description": "Start a snapshot capture of the current gas usage by name in a group.", + "declaration": "function startSnapshotGas(string calldata group, string calldata name) external;", + "visibility": "external", + "mutability": "", + "signature": "startSnapshotGas(string,string)", + "selector": "0x6cd0cc53", + "selectorBytes": [ + 108, + 208, + 204, + 83 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, { "func": { "id": "startStateDiffRecording", @@ -8703,6 +8823,66 @@ "status": "stable", "safety": "unsafe" }, + { + "func": { + "id": "stopSnapshotGas_0", + "description": "", + "declaration": "function stopSnapshotGas() external returns (uint256 gasUsed);", + "visibility": "external", + "mutability": "", + "signature": "stopSnapshotGas()", + "selector": "0xf6402eda", + "selectorBytes": [ + 246, + 64, + 46, + 218 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, + { + "func": { + "id": "stopSnapshotGas_1", + "description": "Stop the snapshot capture of the current gas usage by name, capturing the gas used since the start.\nThe group name is derived from the contract name.", + "declaration": "function stopSnapshotGas(string calldata name) external returns (uint256 gasUsed);", + "visibility": "external", + "mutability": "", + "signature": "stopSnapshotGas(string)", + "selector": "0x773b2805", + "selectorBytes": [ + 119, + 59, + 40, + 5 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, + { + "func": { + "id": "stopSnapshotGas_2", + "description": "Stop the snapshot capture of the current gas usage by name in a group, capturing the gas used since the start.", + "declaration": "function stopSnapshotGas(string calldata group, string calldata name) external returns (uint256 gasUsed);", + "visibility": "external", + "mutability": "", + "signature": "stopSnapshotGas(string,string)", + "selector": "0x0c9db707", + "selectorBytes": [ + 12, + 157, + 183, + 7 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, { "func": { "id": "store", diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index a2973e55bb66e..550c887636449 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -508,6 +508,49 @@ interface Vm { #[cheatcode(group = Evm, safety = Unsafe)] function readCallers() external returns (CallerMode callerMode, address msgSender, address txOrigin); + // ----- Arbitrary Snapshots ----- + + /// Snapshot capture an arbitrary numerical value by name. + /// The group name is derived from the contract name. + #[cheatcode(group = Evm, safety = Unsafe)] + function snapshotValue(string calldata name, uint256 value) external; + + /// Snapshot capture an arbitrary numerical value by name in a group. + #[cheatcode(group = Evm, safety = Unsafe)] + function snapshotValue(string calldata group, string calldata name, uint256 value) external; + + // -------- Gas Snapshots -------- + + /// Snapshot capture the gas usage of the last call by name. + #[cheatcode(group = Evm, safety = Unsafe)] + function snapshotGasLastCall(string calldata name) external; + + /// Snapshot capture the gas usage of the last call by name in a group. + #[cheatcode(group = Evm, safety = Unsafe)] + function snapshotGasLastCall(string calldata group, string calldata name) external; + + /// Start a snapshot capture of the current gas usage by name. + /// The group name is derived from the contract name. + #[cheatcode(group = Evm, safety = Unsafe)] + function startSnapshotGas(string calldata name) external; + + /// Start a snapshot capture of the current gas usage by name in a group. + #[cheatcode(group = Evm, safety = Unsafe)] + function startSnapshotGas(string calldata group, string calldata name) external; + + // Stop the snapshot capture of the current gas by latest snapshot name, capturing the gas used since the start. + #[cheatcode(group = Evm, safety = Unsafe)] + function stopSnapshotGas() external returns (uint256 gasUsed); + + /// Stop the snapshot capture of the current gas usage by name, capturing the gas used since the start. + /// The group name is derived from the contract name. + #[cheatcode(group = Evm, safety = Unsafe)] + function stopSnapshotGas(string calldata name) external returns (uint256 gasUsed); + + /// Stop the snapshot capture of the current gas usage by name in a group, capturing the gas used since the start. + #[cheatcode(group = Evm, safety = Unsafe)] + function stopSnapshotGas(string calldata group, string calldata name) external returns (uint256 gasUsed); + // -------- State Snapshots -------- /// `snapshot` is being deprecated in favor of `snapshotState`. It will be removed in future versions. diff --git a/crates/cheatcodes/src/config.rs b/crates/cheatcodes/src/config.rs index 715e26dc57bb9..caf83162bd4bd 100644 --- a/crates/cheatcodes/src/config.rs +++ b/crates/cheatcodes/src/config.rs @@ -50,6 +50,8 @@ pub struct CheatsConfig { /// If Some, `vm.getDeployedCode` invocations are validated to be in scope of this list. /// If None, no validation is performed. pub available_artifacts: Option, + /// Name of the script/test contract which is currently running. + pub running_contract: Option, /// Version of the script/test contract which is currently running. pub running_version: Option, /// Whether to enable legacy (non-reverting) assertions. @@ -65,6 +67,7 @@ impl CheatsConfig { evm_opts: EvmOpts, available_artifacts: Option, script_wallets: Option, + running_contract: Option, running_version: Option, ) -> Self { let mut allowed_paths = vec![config.root.0.clone()]; @@ -93,6 +96,7 @@ impl CheatsConfig { labels: config.labels.clone(), script_wallets, available_artifacts, + running_contract, running_version, assertions_revert: config.assertions_revert, seed: config.fuzz.seed, @@ -222,6 +226,7 @@ impl Default for CheatsConfig { labels: Default::default(), script_wallets: None, available_artifacts: Default::default(), + running_contract: Default::default(), running_version: Default::default(), assertions_revert: true, seed: None, @@ -241,6 +246,7 @@ mod tests { None, None, None, + None, ) } diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index 7ed1ce1a4edf0..59bfdf2f68fcc 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -52,6 +52,19 @@ impl RecordAccess { } } +/// Records the `snapshotGas*` cheatcodes. +#[derive(Clone, Debug)] +pub struct GasRecord { + /// The group name of the gas snapshot. + pub group: String, + /// The name of the gas snapshot. + pub name: String, + /// The total gas used in the gas snapshot. + pub gas_used: u64, + /// Depth at which the gas snapshot was taken. + pub depth: u64, +} + /// Records `deal` cheatcodes #[derive(Clone, Debug)] pub struct DealRecord { @@ -506,6 +519,85 @@ impl Cheatcode for readCallersCall { } } +impl Cheatcode for snapshotValue_0Call { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { name, value } = self; + inner_create_value_snapshot(ccx, None, Some(name.clone()), value.to_string()) + } +} + +impl Cheatcode for snapshotValue_1Call { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { group, name, value } = self; + inner_create_value_snapshot(ccx, Some(group.clone()), Some(name.clone()), value.to_string()) + } +} + +impl Cheatcode for snapshotGasLastCall_0Call { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { name } = self; + let Some(last_call_gas) = &ccx.state.gas_metering.last_call_gas else { + bail!("no external call was made yet"); + }; + inner_create_value_snapshot( + ccx, + None, + Some(name.clone()), + last_call_gas.gasTotalUsed.to_string(), + ) + } +} + +impl Cheatcode for snapshotGasLastCall_1Call { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { name, group } = self; + let Some(last_call_gas) = &ccx.state.gas_metering.last_call_gas else { + bail!("no external call was made yet"); + }; + inner_create_value_snapshot( + ccx, + Some(group.clone()), + Some(name.clone()), + last_call_gas.gasTotalUsed.to_string(), + ) + } +} + +impl Cheatcode for startSnapshotGas_0Call { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { name } = self; + inner_start_gas_snapshot(ccx, None, Some(name.clone())) + } +} + +impl Cheatcode for startSnapshotGas_1Call { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { group, name } = self; + inner_start_gas_snapshot(ccx, Some(group.clone()), Some(name.clone())) + } +} + +impl Cheatcode for stopSnapshotGas_0Call { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self {} = self; + inner_stop_gas_snapshot(ccx, None, None) + } +} + +impl Cheatcode for stopSnapshotGas_1Call { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { name } = self; + inner_stop_gas_snapshot(ccx, None, Some(name.clone())) + } +} + +impl Cheatcode for stopSnapshotGas_2Call { + fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { group, name } = self; + inner_stop_gas_snapshot(ccx, Some(group.clone()), Some(name.clone())) + } +} + // Deprecated in favor of `snapshotStateCall` impl Cheatcode for snapshotCall { fn apply_stateful(&self, ccx: &mut CheatsCtxt) -> Result { @@ -695,6 +787,127 @@ fn inner_delete_state_snapshots(ccx: &mut CheatsCtxt) -> Re Ok(Default::default()) } +fn inner_create_value_snapshot( + ccx: &mut CheatsCtxt, + group: Option, + name: Option, + value: String, +) -> Result { + let cheatcodes = ccx.state.clone(); + let group = group + .as_deref() + .unwrap_or(cheatcodes.config.running_contract.as_ref().expect("expected running contract")) + .to_string(); + let name = name.as_deref().unwrap_or("default").to_string(); + + ccx.state.gas_snapshots.entry(group).or_default().insert(name, value); + + Ok(Default::default()) +} + +fn inner_start_gas_snapshot( + ccx: &mut CheatsCtxt, + group: Option, + name: Option, +) -> Result { + // Revert if there is an active gas snapshot as we can only have one active snapshot at a time. + if ccx.state.gas_metering.last_snapshot_group.is_some() || + ccx.state.gas_metering.last_snapshot_name.is_some() + { + bail!( + "gas snapshot was already started with group: {:?} and name: {:?}", + ccx.state.gas_metering.last_snapshot_group, + ccx.state.gas_metering.last_snapshot_name + ); + } + + let group = group.as_deref().unwrap_or( + ccx.state.config.running_contract.as_deref().expect("expected running contract"), + ); + let name = name.as_deref().unwrap_or("default").to_string(); + + ccx.state.gas_metering.gas_records.push(GasRecord { + group: group.to_string(), + name: name.clone(), + gas_used: 0, + depth: ccx.ecx.journaled_state.depth() - 1, + }); + + ccx.state.gas_metering.last_snapshot_group = Some(group.to_string()); + ccx.state.gas_metering.last_snapshot_name = Some(name); + + ccx.state.gas_metering.start(); + + Ok(Default::default()) +} + +fn inner_stop_gas_snapshot( + ccx: &mut CheatsCtxt, + group: Option, + name: Option, +) -> Result { + // If group and name are not provided, use the last snapshot group and name. + let group = group + .as_deref() + .unwrap_or( + ccx.state + .gas_metering + .last_snapshot_group + .as_deref() + .expect("no gas snapshot was started with this group"), + ) + .to_string(); + let name = name + .as_deref() + .unwrap_or( + ccx.state + .gas_metering + .last_snapshot_name + .as_deref() + .expect("no gas snapshot was started with this name"), + ) + .to_string(); + + if let Some(record) = ccx + .state + .gas_metering + .gas_records + .iter_mut() + .find(|record| record.group == group && record.name == name) + { + // Calculate the gas used since the snapshot was started. + // We subtract 151 from the gas used to account for gas used by the snapshot itself. + let value = record.gas_used - 151; + + ccx.state + .gas_snapshots + .entry(group.to_string()) + .or_default() + .insert(name.clone(), value.to_string()); + + // Stop the gas metering. + ccx.state.gas_metering.stop(); + + // Remove the gas record. + ccx.state + .gas_metering + .gas_records + .retain(|record| record.group != group && record.name != name); + + // Clear last snapshot cache. + if ccx.state.gas_metering.last_snapshot_group == Some(group.to_string()) && + ccx.state.gas_metering.last_snapshot_name == Some(name) + { + ccx.state.gas_metering.last_snapshot_group = None; + ccx.state.gas_metering.last_snapshot_name = None; + } + + Ok(value.abi_encode()) + } else { + bail!("no gas snapshot was started with the name: {name} in group: {group}"); + } +} + /// Reads the current caller information and returns the current [CallerMode], `msg.sender` and /// `tx.origin`. /// diff --git a/crates/cheatcodes/src/inspector.rs b/crates/cheatcodes/src/inspector.rs index 401966c4626b0..bd2398c801179 100644 --- a/crates/cheatcodes/src/inspector.rs +++ b/crates/cheatcodes/src/inspector.rs @@ -5,7 +5,7 @@ use crate::{ mapping::{self, MappingSlots}, mock::{MockCallDataContext, MockCallReturnData}, prank::Prank, - DealRecord, RecordAccess, + DealRecord, GasRecord, RecordAccess, }, inspector::utils::CommonCreateInput, script::{Broadcast, ScriptWallets}, @@ -17,8 +17,8 @@ use crate::{ }, }, utils::IgnoredTraces, - CheatsConfig, CheatsCtxt, DynCheatcode, Error, Result, Vm, - Vm::AccountAccess, + CheatsConfig, CheatsCtxt, DynCheatcode, Error, Result, + Vm::{self, AccountAccess}, }; use alloy_primitives::{hex, Address, Bytes, Log, TxKind, B256, U256}; use alloy_rpc_types::request::{TransactionInput, TransactionRequest}; @@ -230,12 +230,34 @@ pub struct GasMetering { /// Stores frames paused gas. pub paused_frames: Vec, + /// The group of the last snapshot taken. + pub last_snapshot_group: Option, + /// The name of the last snapshot taken. + pub last_snapshot_name: Option, + /// Cache of the amount of gas used in previous call. /// This is used by the `lastCallGas` cheatcode. pub last_call_gas: Option, + + /// True if gas recording is enabled. + pub recording: bool, + /// The gas used in the last frame. + pub last_gas_used: u64, + /// Gas records for the active snapshots. + pub gas_records: Vec, } impl GasMetering { + /// Start the gas recording. + pub fn start(&mut self) { + self.recording = true; + } + + /// Stop the gas recording. + pub fn stop(&mut self) { + self.recording = false; + } + /// Resume paused gas metering. pub fn resume(&mut self) { if self.paused { @@ -430,6 +452,10 @@ pub struct Cheatcodes { /// Gas metering state. pub gas_metering: GasMetering, + /// Contains gas snapshots made over the course of a test suite. + // **Note**: both must a BTreeMap to ensure the order of the keys is deterministic. + pub gas_snapshots: BTreeMap>, + /// Mapping slots. pub mapping_slots: Option>, @@ -488,6 +514,7 @@ impl Cheatcodes { serialized_jsons: Default::default(), eth_deals: Default::default(), gas_metering: Default::default(), + gas_snapshots: Default::default(), mapping_slots: Default::default(), pc: Default::default(), breakpoints: Default::default(), @@ -753,6 +780,13 @@ impl Cheatcodes { } } + // Store the total gas used for all active gas records started by `startSnapshotGas`. + self.gas_metering.gas_records.iter_mut().for_each(|record| { + if ecx.journaled_state.depth() == record.depth + 1 { + record.gas_used = record.gas_used.saturating_add(outcome.result.gas.spent()); + } + }); + // 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 { @@ -1152,6 +1186,10 @@ impl Inspector for Cheatcodes { #[inline] fn step_end(&mut self, interpreter: &mut Interpreter, ecx: &mut EvmContext) { + if self.gas_metering.recording { + self.meter_gas_record(interpreter, ecx); + } + if self.gas_metering.paused { self.meter_gas_end(interpreter); } @@ -1309,6 +1347,13 @@ impl Inspector for Cheatcodes { gasRemaining: gas.remaining(), }); + // Store the total gas used for all active gas records started by `startSnapshotGas`. + self.gas_metering.gas_records.iter_mut().for_each(|record| { + if ecx.journaled_state.depth() == record.depth + 1 { + record.gas_used = record.gas_used.saturating_add(gas.spent()); + } + }); + // 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 { @@ -1556,6 +1601,59 @@ impl Cheatcodes { } } + #[cold] + fn meter_gas_record( + &mut self, + interpreter: &mut Interpreter, + ecx: &mut EvmContext, + ) { + if matches!(interpreter.instruction_result, InstructionResult::Continue) { + match interpreter.current_opcode() { + op::CREATE | + op::CALL | + op::CALLCODE | + op::DELEGATECALL | + op::CREATE2 | + op::STATICCALL | + op::EXTSTATICCALL | + op::EXTDELEGATECALL => { + // Reset gas used when entering a new frame. + self.gas_metering.last_gas_used = 0; + } + _ => { + self.gas_metering.gas_records.iter_mut().for_each(|record| { + if ecx.journaled_state.depth() == record.depth + 1 { + // Initialize after new frame, use this as the starting point. + if self.gas_metering.last_gas_used == 0 { + self.gas_metering.last_gas_used = interpreter.gas.spent(); + return; + } + + // Calculate the gas difference between the last and current frame. + let gas_diff = interpreter + .gas + .spent() + .saturating_sub(self.gas_metering.last_gas_used); + + // Update the gas record. + record.gas_used = record.gas_used.saturating_add(gas_diff); + + // Update for next iteration. + self.gas_metering.last_gas_used = interpreter.gas.spent(); + } + }); + } + } + } + + println!( + "RECORD [{:?}]: {:?} @ {:?}", + ecx.journaled_state.depth(), + revm::interpreter::OpCode::new(interpreter.current_opcode()).unwrap().as_str(), + self.gas_metering.gas_records + ); + } + #[cold] fn meter_gas_end(&mut self, interpreter: &mut Interpreter) { // Remove recorded gas if we exit frame. diff --git a/crates/chisel/src/executor.rs b/crates/chisel/src/executor.rs index a0b9fc391e9b9..3e09f1fd14e0a 100644 --- a/crates/chisel/src/executor.rs +++ b/crates/chisel/src/executor.rs @@ -308,6 +308,7 @@ impl SessionSource { self.config.evm_opts.clone(), None, None, + None, Some(self.solc.version.clone()), ) .into(), diff --git a/crates/common/src/fs.rs b/crates/common/src/fs.rs index 45c21eba69e2e..71a62d13a7ae7 100644 --- a/crates/common/src/fs.rs +++ b/crates/common/src/fs.rs @@ -56,6 +56,15 @@ pub fn write_json_file(path: &Path, obj: &T) -> Result<()> { writer.flush().map_err(|e| FsPathError::write(e, path)) } +/// Writes the object as a pretty JSON object. +pub fn write_pretty_json_file(path: &Path, obj: &T) -> Result<()> { + let file = create_file(path)?; + let mut writer = BufWriter::new(file); + serde_json::to_writer_pretty(&mut writer, obj) + .map_err(|source| FsPathError::WriteJson { source, path: path.into() })?; + writer.flush().map_err(|e| FsPathError::write(e, path)) +} + /// Wrapper for `std::fs::write` pub fn write(path: impl AsRef, contents: impl AsRef<[u8]>) -> Result<()> { let path = path.as_ref(); diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index d77e30492fc5b..758cc29dd31f4 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -176,6 +176,8 @@ pub struct Config { pub cache: bool, /// where the cache is stored if enabled pub cache_path: PathBuf, + /// where the gas snapshots are stored + pub snapshots: PathBuf, /// where the broadcast logs are stored pub broadcast: PathBuf, /// additional solc allow paths for `--allow-paths` @@ -719,6 +721,7 @@ impl Config { self.out = p(&root, &self.out); self.broadcast = p(&root, &self.broadcast); self.cache_path = p(&root, &self.cache_path); + self.snapshots = p(&root, &self.snapshots); if let Some(build_info_path) = self.build_info_path { self.build_info_path = Some(p(&root, &build_info_path)); @@ -2087,6 +2090,7 @@ impl Default for Config { cache: true, cache_path: "cache".into(), broadcast: "broadcast".into(), + snapshots: "snapshots".into(), allow_paths: vec![], include_paths: vec![], force: false, diff --git a/crates/evm/evm/src/executors/fuzz/mod.rs b/crates/evm/evm/src/executors/fuzz/mod.rs index 54480181ff1cd..3d4aa8bb5ec06 100644 --- a/crates/evm/evm/src/executors/fuzz/mod.rs +++ b/crates/evm/evm/src/executors/fuzz/mod.rs @@ -17,7 +17,10 @@ use foundry_evm_fuzz::{ use foundry_evm_traces::SparsedTraceArena; use indicatif::ProgressBar; use proptest::test_runner::{TestCaseError, TestError, TestRunner}; -use std::{cell::RefCell, collections::HashMap}; +use std::{ + cell::RefCell, + collections::{BTreeMap, HashMap}, +}; mod types; pub use types::{CaseOutcome, CounterExampleOutcome, FuzzOutcome}; @@ -39,6 +42,8 @@ pub struct FuzzTestData { pub coverage: Option, // Stores logs for all fuzz cases pub logs: Vec, + // Stores gas snapshots for all fuzz cases + pub gas_snapshots: BTreeMap>, // Deprecated cheatcodes mapped to their replacements. pub deprecated_cheatcodes: HashMap<&'static str, Option<&'static str>>, } @@ -108,9 +113,11 @@ impl FuzzedExecutor { FuzzOutcome::Case(case) => { let mut data = execution_data.borrow_mut(); data.gas_by_case.push((case.case.gas, case.case.stipend)); + if data.first_case.is_none() { data.first_case.replace(case.case); } + if let Some(call_traces) = case.traces { if data.traces.len() == max_traces_to_collect { data.traces.pop(); @@ -118,14 +125,25 @@ impl FuzzedExecutor { data.traces.push(call_traces); data.breakpoints.replace(case.breakpoints); } + if show_logs { data.logs.extend(case.logs); } + + // Collect gas snapshots. + for (group, new_snapshots) in case.gas_snapshots.iter() { + data.gas_snapshots + .entry(group.clone()) + .or_default() + .extend(new_snapshots.clone()); + } + // Collect and merge coverage if `forge snapshot` context. match &mut data.coverage { Some(prev) => prev.merge(case.coverage.unwrap()), opt => *opt = case.coverage, } + data.deprecated_cheatcodes = case.deprecated_cheatcodes; Ok(()) @@ -171,6 +189,7 @@ impl FuzzedExecutor { breakpoints: last_run_breakpoints, gas_report_traces: traces.into_iter().map(|a| a.arena).collect(), coverage: fuzz_result.coverage, + gas_snapshots: fuzz_result.gas_snapshots, deprecated_cheatcodes: fuzz_result.deprecated_cheatcodes, }; @@ -239,6 +258,11 @@ impl FuzzedExecutor { (cheats.breakpoints.clone(), cheats.deprecated.clone()) }); + let gas_snapshots = call + .cheatcodes + .as_ref() + .map_or_else(Default::default, |cheats| cheats.gas_snapshots.clone()); + let success = self.executor.is_raw_call_mut_success(address, &mut call, should_fail); if success { Ok(FuzzOutcome::Case(CaseOutcome { @@ -247,6 +271,7 @@ impl FuzzedExecutor { coverage: call.coverage, breakpoints, logs: call.logs, + gas_snapshots, deprecated_cheatcodes, })) } else { diff --git a/crates/evm/evm/src/executors/fuzz/types.rs b/crates/evm/evm/src/executors/fuzz/types.rs index 081ff91129cf4..439bdb8ebf569 100644 --- a/crates/evm/evm/src/executors/fuzz/types.rs +++ b/crates/evm/evm/src/executors/fuzz/types.rs @@ -5,7 +5,7 @@ use foundry_evm_coverage::HitMaps; use foundry_evm_fuzz::FuzzCase; use foundry_evm_traces::SparsedTraceArena; use revm::interpreter::InstructionResult; -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; /// Returned by a single fuzz in the case of a successful run #[derive(Debug)] @@ -20,6 +20,8 @@ pub struct CaseOutcome { pub breakpoints: Breakpoints, /// logs of a single fuzz test case. pub logs: Vec, + /// Gas snapshots of a single fuzz test case + pub gas_snapshots: BTreeMap>, // Deprecated cheatcodes mapped to their replacements. pub deprecated_cheatcodes: HashMap<&'static str, Option<&'static str>>, } diff --git a/crates/evm/fuzz/src/lib.rs b/crates/evm/fuzz/src/lib.rs index 2930a6aa020ca..ef472fd296e3d 100644 --- a/crates/evm/fuzz/src/lib.rs +++ b/crates/evm/fuzz/src/lib.rs @@ -15,7 +15,11 @@ use foundry_evm_coverage::HitMaps; use foundry_evm_traces::{CallTraceArena, SparsedTraceArena}; use itertools::Itertools; use serde::{Deserialize, Serialize}; -use std::{collections::HashMap, fmt, sync::Arc}; +use std::{ + collections::{BTreeMap, HashMap}, + fmt, + sync::Arc, +}; pub use proptest::test_runner::{Config as FuzzConfig, Reason}; @@ -184,6 +188,9 @@ pub struct FuzzTestResult { /// Breakpoints for debugger. Correspond to the same fuzz case as `traces`. pub breakpoints: Option, + /// Any captured gas snapshots along the test's execution which should be accumulated. + pub gas_snapshots: BTreeMap>, + // Deprecated cheatcodes mapped to their replacements. pub deprecated_cheatcodes: HashMap<&'static str, Option<&'static str>>, } diff --git a/crates/forge/bin/cmd/test/mod.rs b/crates/forge/bin/cmd/test/mod.rs index 0eebbb9a0ee7a..7941fa9434072 100644 --- a/crates/forge/bin/cmd/test/mod.rs +++ b/crates/forge/bin/cmd/test/mod.rs @@ -20,7 +20,13 @@ use foundry_cli::{ opts::CoreBuildArgs, utils::{self, LoadConfig}, }; -use foundry_common::{cli_warn, compile::ProjectCompiler, evm::EvmArgs, fs, shell}; +use foundry_common::{ + cli_warn, + compile::ProjectCompiler, + evm::EvmArgs, + fs::{self, create_dir_all, read_json_file, remove_dir_all, write_pretty_json_file}, + shell, +}; use foundry_compilers::{ artifacts::output_selection::OutputSelection, compilers::{multi::MultiCompilerLanguage, CompilerSettings, Language}, @@ -546,6 +552,8 @@ impl TestArgs { .gas_report .then(|| GasReport::new(config.gas_reports.clone(), config.gas_reports_ignore.clone())); + let mut gas_snapshots = BTreeMap::>::new(); + let mut outcome = TestOutcome::empty(self.allow_failure); let mut any_test_failed = false; @@ -655,8 +663,77 @@ impl TestArgs { } } } + + // Collect and merge gas snapshots. + for (group, new_snapshots) in result.gas_snapshots.iter() { + gas_snapshots.entry(group.clone()).or_default().extend(new_snapshots.clone()); + } } + // Check for differences in gas snapshots if `FORGE_SNAPSHOT_CHECK` is set. + // Exiting early with code 1 if differences are found. + if std::env::var("FORGE_SNAPSHOT_CHECK").is_ok() { + let differences_found = gas_snapshots.clone().into_iter().fold( + false, + |mut found, (group, snapshots)| { + let previous_snapshots: BTreeMap = + read_json_file(&config.snapshots.join(format!("{group}.json"))) + .expect("Failed to read snapshots from disk"); + + let diff: BTreeMap<_, _> = snapshots + .iter() + .filter_map(|(k, v)| { + previous_snapshots.get(k).and_then(|previous_snapshot| { + if previous_snapshot != v { + Some((k.clone(), (previous_snapshot.clone(), v.clone()))) + } else { + None + } + }) + }) + .collect(); + + if !diff.is_empty() { + println!( + "{}", + format!("\n[{group}] Failed to match snapshots:").red().bold() + ); + + for (key, (previous_snapshot, snapshot)) in &diff { + println!( + "{}", + format!("- [{key}] {previous_snapshot} → {snapshot}").red() + ); + } + + found = true; + } + + found + }, + ); + + if differences_found { + println!(); + eyre::bail!("Snapshots differ from previous run"); + } + } + + // Remove any existing gas snapshots. + if config.snapshots.exists() { + remove_dir_all(&config.snapshots) + .expect("Failed to remove gas snapshots directory"); + } + + // Create `snapshots` directory if it doesn't exist. + create_dir_all(&config.snapshots)?; + + // Write gas snapshots to disk per group. + gas_snapshots.clone().into_iter().for_each(|(group, snapshots)| { + write_pretty_json_file(&config.snapshots.join(format!("{group}.json")), &snapshots) + .expect("Failed to write gas snapshots to disk"); + }); + // Print suite summary. shell::println(suite_result.summary())?; diff --git a/crates/forge/src/multi_runner.rs b/crates/forge/src/multi_runner.rs index 43aade0ff65c1..802ad3884cc65 100644 --- a/crates/forge/src/multi_runner.rs +++ b/crates/forge/src/multi_runner.rs @@ -243,6 +243,7 @@ impl MultiContractRunner { self.evm_opts.clone(), Some(self.known_contracts.clone()), None, + Some(artifact_id.name.clone()), Some(artifact_id.version.clone()), ); diff --git a/crates/forge/src/result.rs b/crates/forge/src/result.rs index 171a234a5ee02..93f2a38ce1b44 100644 --- a/crates/forge/src/result.rs +++ b/crates/forge/src/result.rs @@ -409,6 +409,9 @@ pub struct TestResult { /// pc breakpoint char map pub breakpoints: Breakpoints, + /// Any captured gas snapshots along the test's execution which should be accumulated. + pub gas_snapshots: BTreeMap>, + /// Deprecated cheatcodes (mapped to their replacements, if any) used in current test. #[serde(skip)] pub deprecated_cheatcodes: HashMap<&'static str, Option<&'static str>>, @@ -528,6 +531,7 @@ impl TestResult { if let Some(cheatcodes) = raw_call_result.cheatcodes { self.breakpoints = cheatcodes.breakpoints; + self.gas_snapshots = cheatcodes.gas_snapshots; self.deprecated_cheatcodes = cheatcodes.deprecated; } @@ -562,6 +566,7 @@ impl TestResult { self.duration = Duration::default(); self.gas_report_traces = result.gas_report_traces.into_iter().map(|t| vec![t]).collect(); self.breakpoints = result.breakpoints.unwrap_or_default(); + self.gas_snapshots = result.gas_snapshots; self.deprecated_cheatcodes = result.deprecated_cheatcodes; self diff --git a/crates/forge/tests/cli/config.rs b/crates/forge/tests/cli/config.rs index 1ddd402722dd4..0b2b2adbb720f 100644 --- a/crates/forge/tests/cli/config.rs +++ b/crates/forge/tests/cli/config.rs @@ -37,6 +37,7 @@ forgetest!(can_extract_config_values, |prj, cmd| { libs: vec!["lib-test".into()], cache: true, cache_path: "test-cache".into(), + snapshots "snapshots".into(), broadcast: "broadcast".into(), force: true, evm_version: EvmVersion::Byzantium, diff --git a/crates/script/src/lib.rs b/crates/script/src/lib.rs index a20a884f04b3d..d9a18d80f9649 100644 --- a/crates/script/src/lib.rs +++ b/crates/script/src/lib.rs @@ -609,6 +609,7 @@ impl ScriptConfig { self.evm_opts.clone(), Some(known_contracts), Some(script_wallets), + Some(target.name), Some(target.version), ) .into(), diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index 993f722f814b3..31fbd8656e7d8 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -416,7 +416,11 @@ interface Vm { function skip(bool skipTest, string calldata reason) external; function sleep(uint256 duration) external; function snapshot() external returns (uint256 snapshotId); + function snapshotGasLastCall(string calldata name) external; + function snapshotGasLastCall(string calldata group, string calldata name) external; function snapshotState() external returns (uint256 snapshotId); + function snapshotValue(string calldata name, uint256 value) external; + function snapshotValue(string calldata group, string calldata name, uint256 value) external; function split(string calldata input, string calldata delimiter) external pure returns (string[] memory outputs); function startBroadcast() external; function startBroadcast(address signer) external; @@ -424,12 +428,17 @@ interface Vm { function startMappingRecording() external; function startPrank(address msgSender) external; function startPrank(address msgSender, address txOrigin) external; + function startSnapshotGas(string calldata name) external; + function startSnapshotGas(string calldata group, string calldata name) external; function startStateDiffRecording() external; function stopAndReturnStateDiff() external returns (AccountAccess[] memory accountAccesses); function stopBroadcast() external; function stopExpectSafeMemory() external; function stopMappingRecording() external; function stopPrank() external; + function stopSnapshotGas() external returns (uint256 gasUsed); + function stopSnapshotGas(string calldata name) external returns (uint256 gasUsed); + function stopSnapshotGas(string calldata group, string calldata name) external returns (uint256 gasUsed); function store(address target, bytes32 slot, bytes32 value) external; function toBase64URL(bytes calldata data) external pure returns (string memory); function toBase64URL(string calldata data) external pure returns (string memory); diff --git a/testdata/default/cheats/GasSnapshots.t.sol b/testdata/default/cheats/GasSnapshots.t.sol new file mode 100644 index 0000000000000..3a91493a50d45 --- /dev/null +++ b/testdata/default/cheats/GasSnapshots.t.sol @@ -0,0 +1,289 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +pragma solidity 0.8.18; + +import "ds-test/test.sol"; +import "cheats/Vm.sol"; + +contract GasSnapshotTest is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + + Flare public flare; + uint256 public slot; + + function setUp() public { + flare = new Flare(); + } + + function testGasExternal() public { + vm.startSnapshotGas("testAssertGasExternal"); + + flare.update(2); + + vm.stopSnapshotGas(); + } + + function testGasInternal() public { + vm.startSnapshotGas("testAssertGasInternalA"); + + slot = 1; + + vm.stopSnapshotGas(); + + vm.startSnapshotGas("testAssertGasInternalB"); + + slot = 2; + + vm.stopSnapshotGas(); + + vm.startSnapshotGas("testAssertGasInternalC"); + + slot = 0; + + vm.stopSnapshotGas(); + + vm.startSnapshotGas("testAssertGasInternalD"); + + slot = 1; + + vm.stopSnapshotGas(); + + vm.startSnapshotGas("testAssertGasInternalE"); + + slot = 2; + + vm.stopSnapshotGas(); + } + + // Writes to `GasSnapshotTest` group with custom names. + function testSnapshotValueDefaultGroup1() public { + uint256 a = 123; + uint256 b = 456; + uint256 c = 789; + + vm.snapshotValue("a", a); + vm.snapshotValue("b", b); + vm.snapshotValue("c", c); + } + + // Writes to same `GasSnapshotTest` group with custom names. + function testSnapshotValueDefaultGroup2() public { + uint256 d = 123; + uint256 e = 456; + uint256 f = 789; + + vm.snapshotValue("d", d); + vm.snapshotValue("e", e); + vm.snapshotValue("f", f); + } + + // Writes to `CustomGroup` group with custom names. + // Asserts that the order of the values is alphabetical. + function testSnapshotValueCustomGroup1() public { + uint256 o = 123; + uint256 i = 456; + uint256 q = 789; + + vm.snapshotValue("CustomGroup", "q", q); + vm.snapshotValue("CustomGroup", "i", i); + vm.snapshotValue("CustomGroup", "o", o); + } + + // Writes to `CustomGroup` group with custom names. + // Asserts that the order of the values is alphabetical. + function testSnapshotValueCustomGroup2() public { + uint256 x = 123; + uint256 e = 456; + uint256 z = 789; + + vm.snapshotValue("CustomGroup", "z", z); + vm.snapshotValue("CustomGroup", "x", x); + vm.snapshotValue("CustomGroup", "e", e); + } + + // Writes to `GasSnapshotTest` group with `testSnapshotGasDefault` name. + function testSnapshotGasSectionDefaultGroupStop() public { + vm.startSnapshotGas("testSnapshotGasSection"); + + flare.run(256); + + // vm.stopSnapshotGas() will use the last snapshot name. + uint256 gasUsed = vm.stopSnapshotGas(); + assertGt(gasUsed, 0); + } + + // Writes to `GasSnapshotTest` group with `testSnapshotGasCustom` name. + function testSnapshotGasSectionCustomGroupStop() public { + vm.startSnapshotGas("CustomGroup", "testSnapshotGasSection"); + + flare.run(256); + + // vm.stopSnapshotGas() will use the last snapshot name, even with custom group. + uint256 gasUsed = vm.stopSnapshotGas(); + assertGt(gasUsed, 0); + } + + // Writes to `GasSnapshotTest` group with `testSnapshotGasSection` name. + function testSnapshotGasSectionName() public { + vm.startSnapshotGas("testSnapshotGasSectionName"); + + flare.run(256); + + uint256 gasUsed = vm.stopSnapshotGas("testSnapshotGasSectionName"); + assertGt(gasUsed, 0); + } + + // Writes to `CustomGroup` group with `testSnapshotGasSection` name. + function testSnapshotGasSectionGroupName() public { + vm.startSnapshotGas("CustomGroup", "testSnapshotGasSectionGroupName"); + + flare.run(256); + + uint256 gasUsed = vm.stopSnapshotGas("CustomGroup", "testSnapshotGasSectionGroupName"); + assertGt(gasUsed, 0); + } + + // Writes to `GasSnapshotTest` group with `testSnapshotGas` name. + function testSnapshotGasLastCallName() public { + flare.run(1); + + vm.snapshotGasLastCall("testSnapshotGasName"); + } + + // Writes to `CustomGroup` group with `testSnapshotGas` name. + function testSnapshotGasLastCallGroupName() public { + flare.run(1); + + vm.snapshotGasLastCall("CustomGroup", "testSnapshotGasGroupName"); + } +} + +contract GasComparisonTest is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + + uint256 public slotA; + uint256 public slotB; + uint256 public cachedGas; + + function testGasComparisonEmpty() public { + // Start a cheatcode snapshot. + vm.startSnapshotGas("ComparisonGroup", "testGasComparisonEmptyA"); + vm.stopSnapshotGas(); + + // Start a comparitive Solidity snapshot. + _snapStart(); + vm.snapshotValue("ComparisonGroup", "testGasComparisonEmptyB", _snapEnd()); + } + + function testGasComparisonInternalCold() public { + // Start a cheatcode snapshot. + vm.startSnapshotGas("ComparisonGroup", "testGasComparisonInternalColdA"); + slotA = 1; + vm.stopSnapshotGas(); + + // Start a comparitive Solidity snapshot. + _snapStart(); + slotB = 1; + vm.snapshotValue("ComparisonGroup", "testGasComparisonInternalColdB", _snapEnd()); + } + + function testGasComparisonInternalWarm() public { + // Warm up the cache. + slotA = 1; + slotB = 1; + + // Start a cheatcode snapshot. + vm.startSnapshotGas("ComparisonGroup", "testGasComparisonInternalWarmA"); + slotA = 2; + vm.stopSnapshotGas(); + + // Start a comparitive Solidity snapshot. + _snapStart(); + slotB = 2; + vm.snapshotValue("ComparisonGroup", "testGasComparisonInternalWarmB", _snapEnd()); + } + + function testGasComparisonExternal() public { + // Warm up the cache. + TargetB targetA = new TargetB(); + targetA.update(1); + TargetB targetB = new TargetB(); + targetB.update(1); + + // Start a cheatcode snapshot. + vm.startSnapshotGas("ComparisonGroup", "testGasComparisonExternalA"); + targetA.update(2); + vm.stopSnapshotGas(); + + // Start a comparitive Solidity snapshot. + _snapStart(); + targetB.update(2); + vm.snapshotValue("ComparisonGroup", "testGasComparisonExternalB", _snapEnd()); + } + + function testGasComparisonCreateA() public { + // Start a cheatcode snapshot. + vm.startSnapshotGas("ComparisonGroup", "testGasComparisonCreateA"); + new TargetEmpty(); + vm.stopSnapshotGas(); + } + + function testGasComparisonCreateB() public { + // Start a comparitive Solidity snapshot. + _snapStart(); + new TargetEmpty(); + vm.snapshotValue("ComparisonGroup", "testGasComparisonCreateB", _snapEnd()); + } + + // Internal function to start a Solidity snapshot. + function _snapStart() internal { + cachedGas = 1; + cachedGas = gasleft(); + } + + // Internal function to end a Solidity snapshot. + function _snapEnd() internal returns (uint256 gasUsed) { + gasUsed = cachedGas - gasleft() - 174; + cachedGas = 2; + } +} + +contract Flare { + TargetA public target; + bytes32[] public data; + + constructor() { + target = new TargetA(); + } + + function run(uint256 n_) public { + for (uint256 i = 0; i < n_; i++) { + data.push(keccak256(abi.encodePacked(i))); + } + } + + function update(uint256 x_) public { + target.update(x_); + } +} + +contract TargetA { + TargetB public target; + + constructor() { + target = new TargetB(); + } + + function update(uint256 x_) public { + target.update(x_); + } +} + +contract TargetB { + uint256 public x; + + function update(uint256 x_) public { + x = x_; + } +} + +contract TargetEmpty {} diff --git a/testdata/foundry.toml b/testdata/foundry.toml index e9189bb008a32..30621914fa353 100644 --- a/testdata/foundry.toml +++ b/testdata/foundry.toml @@ -1,5 +1,5 @@ [profile.default] -solc = "0.8.18" +# solc = "0.8.18" block_base_fee_per_gas = 0 block_coinbase = "0x0000000000000000000000000000000000000000" block_difficulty = 0