chore: remove fork db snapshot (#6486)

* chore: remove fork db snapshot

* feat: add revert action

* clippy happy

crates/anvil/src/eth/backend/db.rs

@@ -10,7 +10,7 @@ use ethers::{
 };
 use foundry_common::{errors::FsPathError, types::ToAlloy};
 use foundry_evm::{
-    backend::{DatabaseError, DatabaseResult, MemDb, StateSnapshot},
+    backend::{DatabaseError, DatabaseResult, MemDb, RevertSnapshotAction, StateSnapshot},
     fork::BlockchainDb,
     hashbrown::HashMap,
     revm::{
@@ -172,7 +172,7 @@ pub trait Db:
     /// Reverts a snapshot
     ///
     /// Returns `true` if the snapshot was reverted
-    fn revert(&mut self, snapshot: U256) -> bool;
+    fn revert(&mut self, snapshot: U256, action: RevertSnapshotAction) -> bool;
 
     /// Returns the state root if possible to compute
     fn maybe_state_root(&self) -> Option<H256> {
@@ -208,7 +208,7 @@ impl<T: DatabaseRef<Error = DatabaseError> + Send + Sync + Clone + fmt::Debug> D
         U256::zero()
     }
 
-    fn revert(&mut self, _snapshot: U256) -> bool {
+    fn revert(&mut self, _snapshot: U256, _action: RevertSnapshotAction) -> bool {
         false
     }
 

crates/anvil/src/eth/backend/mem/fork_db.rs

@@ -9,7 +9,7 @@ use crate::{
 use ethers::{prelude::H256, types::BlockId};
 use foundry_common::types::{ToAlloy, ToEthers};
 use foundry_evm::{
-    backend::{DatabaseResult, StateSnapshot},
+    backend::{DatabaseResult, RevertSnapshotAction, StateSnapshot},
     fork::{database::ForkDbSnapshot, BlockchainDb},
     revm::Database,
 };
@@ -68,8 +68,8 @@ impl Db for ForkedDatabase {
         self.insert_snapshot().to_ethers()
     }
 
-    fn revert(&mut self, id: U256) -> bool {
-        self.revert_snapshot(id.to_alloy())
+    fn revert(&mut self, id: U256, action: RevertSnapshotAction) -> bool {
+        self.revert_snapshot(id.to_alloy(), action)
     }
 
     fn current_state(&self) -> StateDb {

crates/anvil/src/eth/backend/mem/in_memory_db.rs

@@ -17,6 +17,7 @@ use foundry_evm::{
 };
 
 // reexport for convenience
+use foundry_evm::backend::RevertSnapshotAction;
 pub use foundry_evm::{backend::MemDb, revm::db::DatabaseRef};
 
 impl Db for MemDb {
@@ -71,8 +72,11 @@ impl Db for MemDb {
         id.to_ethers()
     }
 
-    fn revert(&mut self, id: U256) -> bool {
+    fn revert(&mut self, id: U256, action: RevertSnapshotAction) -> bool {
         if let Some(snapshot) = self.snapshots.remove(id.to_alloy()) {
+            if action.is_keep() {
+                self.snapshots.insert_at(snapshot.clone(), id.to_alloy());
+            }
             self.inner = snapshot;
             trace!(target: "backend::memdb", "Reverted snapshot {}", id);
             true

crates/anvil/src/eth/backend/mem/mod.rs

@@ -58,7 +58,7 @@ use ethers::{
 use flate2::{read::GzDecoder, write::GzEncoder, Compression};
 use foundry_common::types::{ToAlloy, ToEthers};
 use foundry_evm::{
-    backend::{DatabaseError, DatabaseResult},
+    backend::{DatabaseError, DatabaseResult, RevertSnapshotAction},
     constants::DEFAULT_CREATE2_DEPLOYER_RUNTIME_CODE,
     decode::decode_revert,
     inspectors::AccessListTracer,
@@ -686,7 +686,7 @@ impl Backend {
                 ..Default::default()
             };
         }
-        Ok(self.db.write().await.revert(id))
+        Ok(self.db.write().await.revert(id, RevertSnapshotAction::RevertRemove))
     }
 
     pub fn list_snapshots(&self) -> BTreeMap<U256, (u64, H256)> {

crates/anvil/tests/it/fork.rs

@@ -289,6 +289,51 @@ async fn test_fork_snapshotting() {
     assert_eq!(block_number, provider.get_block_number().await.unwrap());
 }
 
+#[tokio::test(flavor = "multi_thread")]
+async fn test_fork_snapshotting_repeated() {
+    let (api, handle) = spawn(fork_config()).await;
+    let provider = handle.http_provider();
+
+    let snapshot = api.evm_snapshot().await.unwrap();
+
+    let accounts: Vec<_> = handle.dev_wallets().collect();
+    let from = accounts[0].address();
+    let to = accounts[1].address();
+    let block_number = provider.get_block_number().await.unwrap();
+
+    let initial_nonce = provider.get_transaction_count(from, None).await.unwrap();
+    let balance_before = provider.get_balance(to, None).await.unwrap();
+    let amount = handle.genesis_balance().checked_div(2u64.into()).unwrap();
+
+    let tx = TransactionRequest::new().to(to).value(amount).from(from);
+
+    let _ = provider.send_transaction(tx, None).await.unwrap().await.unwrap().unwrap();
+
+    let nonce = provider.get_transaction_count(from, None).await.unwrap();
+    assert_eq!(nonce, initial_nonce + 1);
+    let to_balance = provider.get_balance(to, None).await.unwrap();
+    assert_eq!(balance_before.saturating_add(amount), to_balance);
+
+    let _second_snapshot = api.evm_snapshot().await.unwrap();
+
+    assert!(api.evm_revert(snapshot).await.unwrap());
+
+    let nonce = provider.get_transaction_count(from, None).await.unwrap();
+    assert_eq!(nonce, initial_nonce);
+    let balance = provider.get_balance(from, None).await.unwrap();
+    assert_eq!(balance, handle.genesis_balance());
+    let balance = provider.get_balance(to, None).await.unwrap();
+    assert_eq!(balance, handle.genesis_balance());
+    assert_eq!(block_number, provider.get_block_number().await.unwrap());
+
+    // invalidated
+    // TODO enable after <https://github.com/foundry-rs/foundry/pull/6366>
+    // assert!(!api.evm_revert(second_snapshot).await.unwrap());
+
+    // nothing is reverted, snapshot gone
+    assert!(!api.evm_revert(snapshot).await.unwrap());
+}
+
 /// tests that the remote state and local state are kept separate.
 /// changes don't make into the read only Database that holds the remote state, which is flushed to
 /// a cache file.

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

@@ -41,7 +41,7 @@ mod in_memory_db;
 pub use in_memory_db::{EmptyDBWrapper, FoundryEvmInMemoryDB, MemDb};
 
 mod snapshot;
-pub use snapshot::{BackendSnapshot, StateSnapshot};
+pub use snapshot::{BackendSnapshot, RevertSnapshotAction, StateSnapshot};
 
 // A `revm::Database` that is used in forking mode
 type ForkDB = CacheDB<SharedBackend>;

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

@@ -42,3 +42,22 @@ impl<T> BackendSnapshot<T> {
         self.journaled_state.logs = current.logs.clone();
     }
 }
+
+/// What to do when reverting a snapshot
+///
+/// Whether to remove the snapshot or keep it
+#[derive(Debug, Clone, Copy, Eq, PartialEq, Default)]
+pub enum RevertSnapshotAction {
+    /// Remove the snapshot after reverting
+    #[default]
+    RevertRemove,
+    /// Keep the snapshot after reverting
+    RevertKeep,
+}
+
+impl RevertSnapshotAction {
+    /// Returns `true` if the action is to keep the snapshot
+    pub fn is_keep(&self) -> bool {
+        matches!(self, RevertSnapshotAction::RevertKeep)
+    }
+}

crates/evm/core/src/fork/database.rs

@@ -1,7 +1,7 @@
 //! A revm database that forks off a remote client
 
 use crate::{
-    backend::{DatabaseError, StateSnapshot},
+    backend::{DatabaseError, RevertSnapshotAction, StateSnapshot},
     fork::{BlockchainDb, SharedBackend},
     snapshot::Snapshots,
 };
@@ -110,10 +110,13 @@ impl ForkedDatabase {
         id
     }
 
-    pub fn revert_snapshot(&mut self, id: U256) -> bool {
+    /// Removes the snapshot from the tracked snapshot and sets it as the current state
+    pub fn revert_snapshot(&mut self, id: U256, action: RevertSnapshotAction) -> bool {
         let snapshot = { self.snapshots().lock().remove_at(id) };
         if let Some(snapshot) = snapshot {
-            self.snapshots().lock().insert_at(snapshot.clone(), id);
+            if action.is_keep() {
+                self.snapshots().lock().insert_at(snapshot.clone(), id);
+            }
             let ForkDbSnapshot {
                 local,
                 snapshot: StateSnapshot { accounts, storage, block_hashes },