feat: add additional snapshot cheatcodes

crates/cheatcodes/spec/src/vm.rs

@@ -411,10 +411,34 @@ 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.
+    ///
+    /// 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);
+
+    /// Removes _all_ snapshots previously created by `snapshot`.
+    #[cheatcode(group = Evm, safety = Unsafe)]
+    function deleteSnapshots() external;
+
     // -------- Forking --------
     // --- Creation and Selection ---
 

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<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> 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<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> 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<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { snapshotId } = self;
+        let result = ccx.data.db.delete_snapshot(*snapshotId);
+        Ok(result.abi_encode())
+    }
+}
+impl Cheatcode for deleteSnapshotsCall {
+    fn apply_full<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> 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;

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<JournaledState> {
         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<LocalForkId> {

crates/evm/core/src/backend/mod.rs

@@ -82,14 +82,26 @@ pub trait DatabaseExt: Database<Error = DatabaseError> {
     /// 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<JournaledState>;
 
+    /// 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<JournaledState> {
         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<LocalForkId> {
         trace!("create fork");
         let (fork_id, fork, _) = self.forks.create_fork(create_fork.clone())?;

crates/evm/core/src/snapshot.rs

@@ -39,6 +39,11 @@ impl<T> Snapshots<T> {
         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.

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));
+    }
+
+    // <https://github.com/foundry-rs/foundry/issues/6411>
+    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;