From 36bd3c10533d3db4b2387487f8d91178b8bfbfb9 Mon Sep 17 00:00:00 2001 From: Matthias Seitz Date: Fri, 8 Dec 2023 11:02:52 +0100 Subject: [PATCH 1/3] feat: add additional snapshot cheatcodes --- crates/cheatcodes/spec/src/vm.rs | 17 ++++++++++- crates/cheatcodes/src/evm.rs | 45 ++++++++++++++++++++++++--- crates/evm/core/src/backend/fuzz.rs | 26 +++++++++++++++- crates/evm/core/src/backend/mod.rs | 27 +++++++++++++++-- crates/evm/core/src/snapshot.rs | 5 +++ testdata/cheats/Snapshots.t.sol | 47 +++++++++++++++++++++++++++++ testdata/cheats/Vm.sol | 3 ++ 7 files changed, 162 insertions(+), 8 deletions(-) diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index a9421d087daa..5e618742f470 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -411,10 +411,25 @@ interface Vm { /// Revert the state of the EVM to a previous snapshot /// Takes the snapshot ID to revert to. - /// This deletes the snapshot and all snapshots taken after the given snapshot ID. + /// + /// **Note:** This does not automatically delete the snapshot. To delete the snapshot use `deleteSnapshot`. #[cheatcode(group = Evm, safety = Unsafe)] function revertTo(uint256 snapshotId) external returns (bool success); + /// Revert the state of the EVM to a previous snapshot and automatically deletes the snapshots + /// Takes the snapshot ID to revert to. + #[cheatcode(group = Evm, safety = Unsafe)] + function revertToAndDelete(uint256 snapshotId) external returns (bool success); + + /// Removes the snapshot with the given ID created by `snapshot`. + /// Takes the snapshot ID to delete. + #[cheatcode(group = Evm, safety = Unsafe)] + function deleteSnapshot(uint256 snapshotId) external returns (bool success); + + /// Removes _all_ snapshots previously created by `snapshot`. + #[cheatcode(group = Evm, safety = Unsafe)] + function deleteSnapshots() external; + // -------- Forking -------- // --- Creation and Selection --- diff --git a/crates/cheatcodes/src/evm.rs b/crates/cheatcodes/src/evm.rs index 136cdb8f8048..b8afe62a291f 100644 --- a/crates/cheatcodes/src/evm.rs +++ b/crates/cheatcodes/src/evm.rs @@ -6,7 +6,7 @@ use alloy_sol_types::SolValue; use ethers_core::utils::{Genesis, GenesisAccount}; use ethers_signers::Signer; use foundry_common::{fs::read_json_file, types::ToAlloy}; -use foundry_evm_core::backend::DatabaseExt; +use foundry_evm_core::backend::{DatabaseExt, RevertSnapshotAction}; use revm::{ primitives::{Account, Bytecode, SpecId, KECCAK_EMPTY}, EVMData, @@ -327,9 +327,12 @@ impl Cheatcode for snapshotCall { impl Cheatcode for revertToCall { fn apply_full(&self, ccx: &mut CheatsCtxt) -> Result { let Self { snapshotId } = self; - let result = if let Some(journaled_state) = - ccx.data.db.revert(*snapshotId, &ccx.data.journaled_state, ccx.data.env) - { + let result = if let Some(journaled_state) = ccx.data.db.revert( + *snapshotId, + &ccx.data.journaled_state, + ccx.data.env, + RevertSnapshotAction::RevertKeep, + ) { // we reset the evm's journaled_state to the state of the snapshot previous state ccx.data.journaled_state = journaled_state; true @@ -340,6 +343,40 @@ impl Cheatcode for revertToCall { } } +impl Cheatcode for revertToAndDeleteCall { + fn apply_full(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { snapshotId } = self; + let result = if let Some(journaled_state) = ccx.data.db.revert( + *snapshotId, + &ccx.data.journaled_state, + ccx.data.env, + RevertSnapshotAction::RevertRemove, + ) { + // we reset the evm's journaled_state to the state of the snapshot previous state + ccx.data.journaled_state = journaled_state; + true + } else { + false + }; + Ok(result.abi_encode()) + } +} + +impl Cheatcode for deleteSnapshotCall { + fn apply_full(&self, ccx: &mut CheatsCtxt) -> Result { + let Self { snapshotId } = self; + let result = ccx.data.db.delete_snapshot(*snapshotId); + Ok(result.abi_encode()) + } +} +impl Cheatcode for deleteSnapshotsCall { + fn apply_full(&self, ccx: &mut CheatsCtxt) -> Result { + let Self {} = self; + ccx.data.db.delete_snapshots(); + Ok(Default::default()) + } +} + impl Cheatcode for startStateDiffRecordingCall { fn apply(&self, state: &mut Cheatcodes) -> Result { let Self {} = self; diff --git a/crates/evm/core/src/backend/fuzz.rs b/crates/evm/core/src/backend/fuzz.rs index bf32c6f8800f..0a963810abde 100644 --- a/crates/evm/core/src/backend/fuzz.rs +++ b/crates/evm/core/src/backend/fuzz.rs @@ -3,6 +3,7 @@ use crate::{ backend::{ diagnostic::RevertDiagnostic, error::DatabaseError, Backend, DatabaseExt, LocalForkId, + RevertSnapshotAction, }, fork::{CreateFork, ForkId}, }; @@ -83,6 +84,14 @@ impl<'a> FuzzBackendWrapper<'a> { } self.backend.to_mut() } + + /// Returns a mutable instance of the Backend if it is initialized. + fn initialized_backend_mut(&mut self) -> Option<&mut Backend> { + if self.is_initialized { + return Some(self.backend.to_mut()) + } + None + } } impl<'a> DatabaseExt for FuzzBackendWrapper<'a> { @@ -96,9 +105,24 @@ impl<'a> DatabaseExt for FuzzBackendWrapper<'a> { id: U256, journaled_state: &JournaledState, current: &mut Env, + action: RevertSnapshotAction, ) -> Option { trace!(?id, "fuzz: revert snapshot"); - self.backend_mut(current).revert(id, journaled_state, current) + self.backend_mut(current).revert(id, journaled_state, current, action) + } + + fn delete_snapshot(&mut self, id: U256) -> bool { + // delete snapshot requires a previous snapshot to be initialized + if let Some(backend) = self.initialized_backend_mut() { + return backend.delete_snapshot(id) + } + false + } + + fn delete_snapshots(&mut self) { + if let Some(backend) = self.initialized_backend_mut() { + backend.delete_snapshots() + } } fn create_fork(&mut self, fork: CreateFork) -> eyre::Result { diff --git a/crates/evm/core/src/backend/mod.rs b/crates/evm/core/src/backend/mod.rs index c0ffdfd8f613..246a38320e2b 100644 --- a/crates/evm/core/src/backend/mod.rs +++ b/crates/evm/core/src/backend/mod.rs @@ -82,14 +82,26 @@ pub trait DatabaseExt: Database { /// since the snapshots was created. This way we can show logs that were emitted between /// snapshot and its revert. /// This will also revert any changes in the `Env` and replace it with the captured `Env` of - /// `Self::snapshot` + /// `Self::snapshot`. + /// + /// Depending on [RevertSnapshotAction] it will keep the snapshot alive or delete it. fn revert( &mut self, id: U256, journaled_state: &JournaledState, env: &mut Env, + action: RevertSnapshotAction, ) -> Option; + /// Deletes the snapshot with the given `id` + /// + /// Returns `true` if the snapshot was successfully deleted, `false` if no snapshot for that id + /// exists. + fn delete_snapshot(&mut self, id: U256) -> bool; + + /// Deletes all snapshots. + fn delete_snapshots(&mut self); + /// Creates and also selects a new fork /// /// This is basically `create_fork` + `select_fork` @@ -918,11 +930,14 @@ impl DatabaseExt for Backend { id: U256, current_state: &JournaledState, current: &mut Env, + action: RevertSnapshotAction, ) -> Option { trace!(?id, "revert snapshot"); if let Some(mut snapshot) = self.inner.snapshots.remove_at(id) { // Re-insert snapshot to persist it - self.inner.snapshots.insert_at(snapshot.clone(), id); + if action.is_keep() { + self.inner.snapshots.insert_at(snapshot.clone(), id); + } // need to check whether there's a global failure which means an error occurred either // during the snapshot or even before if self.is_global_failure(current_state) { @@ -969,6 +984,14 @@ impl DatabaseExt for Backend { } } + fn delete_snapshot(&mut self, id: U256) -> bool { + self.inner.snapshots.remove_at(id).is_some() + } + + fn delete_snapshots(&mut self) { + self.inner.snapshots.clear() + } + fn create_fork(&mut self, mut create_fork: CreateFork) -> eyre::Result { trace!("create fork"); let (fork_id, fork, _) = self.forks.create_fork(create_fork.clone())?; diff --git a/crates/evm/core/src/snapshot.rs b/crates/evm/core/src/snapshot.rs index bcd1c9291319..705c5baf7f27 100644 --- a/crates/evm/core/src/snapshot.rs +++ b/crates/evm/core/src/snapshot.rs @@ -39,6 +39,11 @@ impl Snapshots { snapshot } + /// Removes all snapshots + pub fn clear(&mut self) { + self.snapshots.clear(); + } + /// Removes the snapshot with the given `id`. /// /// Does not remove snapshots after it. diff --git a/testdata/cheats/Snapshots.t.sol b/testdata/cheats/Snapshots.t.sol index 9a174be98c52..baf82e2e57c1 100644 --- a/testdata/cheats/Snapshots.t.sol +++ b/testdata/cheats/Snapshots.t.sol @@ -32,6 +32,53 @@ contract SnapshotTest is DSTest { assertEq(store.slot1, 20, "snapshot revert for slot 1 unsuccessful"); } + function testSnapshotRevertDelete() public { + uint256 snapshot = vm.snapshot(); + store.slot0 = 300; + store.slot1 = 400; + + assertEq(store.slot0, 300); + assertEq(store.slot1, 400); + + vm.revertToAndDelete(snapshot); + assertEq(store.slot0, 10, "snapshot revert for slot 0 unsuccessful"); + assertEq(store.slot1, 20, "snapshot revert for slot 1 unsuccessful"); + // nothing to revert to anymore + assert(!vm.revertTo(snapshot)); + } + + function testSnapshotDelete() public { + uint256 snapshot = vm.snapshot(); + store.slot0 = 300; + store.slot1 = 400; + + vm.deleteSnapshot(snapshot); + // nothing to revert to anymore + assert(!vm.revertTo(snapshot)); + } + + function testSnapshotDeleteAll() public { + uint256 snapshot = vm.snapshot(); + store.slot0 = 300; + store.slot1 = 400; + + vm.deleteSnapshots(); + // nothing to revert to anymore + assert(!vm.revertTo(snapshot)); + } + + // + function testSnapshotsMany() public { + uint256 preState; + for (uint256 c = 0; c < 10; c++) { + for (uint256 cc = 0; cc < 10; cc++) { + preState = vm.snapshot(); + vm.revertToAndDelete(preState); + assert(!vm.revertTo(preState)); + } + } + } + // tests that snapshots can also revert changes to `block` function testBlockValues() public { uint256 num = block.number; diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index 4fc5c3ce51b2..38f4ea5d24c4 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -48,6 +48,8 @@ interface Vm { function createWallet(uint256 privateKey) external returns (Wallet memory wallet); function createWallet(uint256 privateKey, string calldata walletLabel) external returns (Wallet memory wallet); function deal(address account, uint256 newBalance) external; + function deleteSnapshot(uint256 snapshot) external returns (bool); + function deleteSnapshots() external; function deriveKey(string calldata mnemonic, uint32 index) external pure returns (uint256 privateKey); function deriveKey(string calldata mnemonic, string calldata derivationPath, uint32 index) external pure returns (uint256 privateKey); function deriveKey(string calldata mnemonic, uint32 index, string calldata language) external pure returns (uint256 privateKey); @@ -172,6 +174,7 @@ interface Vm { function resetNonce(address account) external; function resumeGasMetering() external; function revertTo(uint256 snapshotId) external returns (bool success); + function revertToAndDelete(uint256 snapshotId) external returns (bool success); function revokePersistent(address account) external; function revokePersistent(address[] calldata accounts) external; function roll(uint256 newHeight) external; From 71f6914442dfcd60c289fcc4a04e9fd20f841dc4 Mon Sep 17 00:00:00 2001 From: Matthias Seitz Date: Sun, 10 Dec 2023 12:24:43 +0100 Subject: [PATCH 2/3] chore: docs --- crates/cheatcodes/spec/src/vm.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/crates/cheatcodes/spec/src/vm.rs b/crates/cheatcodes/spec/src/vm.rs index 5e618742f470..b8a4ca885d18 100644 --- a/crates/cheatcodes/spec/src/vm.rs +++ b/crates/cheatcodes/spec/src/vm.rs @@ -412,17 +412,26 @@ interface Vm { /// Revert the state of the EVM to a previous snapshot /// Takes the snapshot ID to revert to. /// + /// Returns `true` if the snapshot was successfully reverted. + /// Returns `false` if the snapshot does not exist. + /// /// **Note:** This does not automatically delete the snapshot. To delete the snapshot use `deleteSnapshot`. #[cheatcode(group = Evm, safety = Unsafe)] function revertTo(uint256 snapshotId) external returns (bool success); /// Revert the state of the EVM to a previous snapshot and automatically deletes the snapshots /// Takes the snapshot ID to revert to. + /// + /// Returns `true` if the snapshot was successfully reverted and deleted. + /// Returns `false` if the snapshot does not exist. #[cheatcode(group = Evm, safety = Unsafe)] function revertToAndDelete(uint256 snapshotId) external returns (bool success); /// Removes the snapshot with the given ID created by `snapshot`. /// Takes the snapshot ID to delete. + /// + /// Returns `true` if the snapshot was successfully deleted. + /// Returns `false` if the snapshot does not exist. #[cheatcode(group = Evm, safety = Unsafe)] function deleteSnapshot(uint256 snapshotId) external returns (bool success); From 1e9b8a07eb56202b2322707059a712f89108fa59 Mon Sep 17 00:00:00 2001 From: Matthias Seitz Date: Sun, 10 Dec 2023 12:32:23 +0100 Subject: [PATCH 3/3] cargo cheats --- crates/cheatcodes/assets/cheatcodes.json | 62 +++++++++++++++++++++++- testdata/cheats/Vm.sol | 2 +- 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/crates/cheatcodes/assets/cheatcodes.json b/crates/cheatcodes/assets/cheatcodes.json index c835e4633997..51e484358306 100644 --- a/crates/cheatcodes/assets/cheatcodes.json +++ b/crates/cheatcodes/assets/cheatcodes.json @@ -997,6 +997,46 @@ "status": "stable", "safety": "unsafe" }, + { + "func": { + "id": "deleteSnapshot", + "description": "Removes the snapshot with the given ID created by `snapshot`.\nTakes the snapshot ID to delete.\nReturns `true` if the snapshot was successfully deleted.\nReturns `false` if the snapshot does not exist.", + "declaration": "function deleteSnapshot(uint256 snapshotId) external returns (bool success);", + "visibility": "external", + "mutability": "", + "signature": "deleteSnapshot(uint256)", + "selector": "0xa6368557", + "selectorBytes": [ + 166, + 54, + 133, + 87 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, + { + "func": { + "id": "deleteSnapshots", + "description": "Removes _all_ snapshots previously created by `snapshot`.", + "declaration": "function deleteSnapshots() external;", + "visibility": "external", + "mutability": "", + "signature": "deleteSnapshots()", + "selector": "0x421ae469", + "selectorBytes": [ + 66, + 26, + 228, + 105 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, { "func": { "id": "deriveKey_0", @@ -3460,7 +3500,7 @@ { "func": { "id": "revertTo", - "description": "Revert the state of the EVM to a previous snapshot\nTakes the snapshot ID to revert to.\nThis deletes the snapshot and all snapshots taken after the given snapshot ID.", + "description": "Revert the state of the EVM to a previous snapshot\nTakes the snapshot ID to revert to.\nReturns `true` if the snapshot was successfully reverted.\nReturns `false` if the snapshot does not exist.\n**Note:** This does not automatically delete the snapshot. To delete the snapshot use `deleteSnapshot`.", "declaration": "function revertTo(uint256 snapshotId) external returns (bool success);", "visibility": "external", "mutability": "", @@ -3477,6 +3517,26 @@ "status": "stable", "safety": "unsafe" }, + { + "func": { + "id": "revertToAndDelete", + "description": "Revert the state of the EVM to a previous snapshot and automatically deletes the snapshots\nTakes the snapshot ID to revert to.\nReturns `true` if the snapshot was successfully reverted and deleted.\nReturns `false` if the snapshot does not exist.", + "declaration": "function revertToAndDelete(uint256 snapshotId) external returns (bool success);", + "visibility": "external", + "mutability": "", + "signature": "revertToAndDelete(uint256)", + "selector": "0x03e0aca9", + "selectorBytes": [ + 3, + 224, + 172, + 169 + ] + }, + "group": "evm", + "status": "stable", + "safety": "unsafe" + }, { "func": { "id": "revokePersistent_0", diff --git a/testdata/cheats/Vm.sol b/testdata/cheats/Vm.sol index 38f4ea5d24c4..fc9e9f06cbb4 100644 --- a/testdata/cheats/Vm.sol +++ b/testdata/cheats/Vm.sol @@ -48,7 +48,7 @@ interface Vm { function createWallet(uint256 privateKey) external returns (Wallet memory wallet); function createWallet(uint256 privateKey, string calldata walletLabel) external returns (Wallet memory wallet); function deal(address account, uint256 newBalance) external; - function deleteSnapshot(uint256 snapshot) external returns (bool); + function deleteSnapshot(uint256 snapshotId) external returns (bool success); function deleteSnapshots() external; function deriveKey(string calldata mnemonic, uint32 index) external pure returns (uint256 privateKey); function deriveKey(string calldata mnemonic, string calldata derivationPath, uint32 index) external pure returns (uint256 privateKey);