evm_revert keeps same snapshot around; cleans up all later ones

crates/anvil/core/src/types.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use ethers_core::types::{TxHash, H256, U256, U64};
 use revm::primitives::SpecId;
 
@@ -217,6 +219,7 @@ pub struct AnvilMetadata {
     pub latest_block_number: u64,
     pub latest_block_hash: H256,
     pub forked_network: Option<ForkedNetwork>,
+    pub snapshots: HashMap<U256, (u64, H256)>,
 }
 
 /// Information about the forked network.

crates/anvil/src/eth/api.rs

@@ -1709,6 +1709,8 @@ impl EthApi {
         node_info!("anvil_metadata");
         let fork_config = self.backend.get_fork();
 
+        let snapshots = self.backend.list_snapshots().await;
+
         Ok(AnvilMetadata {
             client_version: CLIENT_VERSION,
             chain_id: self.backend.chain_id().try_into().unwrap_or(u64::MAX),
@@ -1720,6 +1722,7 @@ impl EthApi {
                 fork_block_number: cfg.block_number().into(),
                 fork_block_hash: cfg.block_hash(),
             }),
+            snapshots,
         })
     }
 

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

@@ -72,7 +72,7 @@ impl Db for MemDb {
     }
 
     fn revert(&mut self, id: U256) -> bool {
-        if let Some(snapshot) = self.snapshots.remove(id.to_alloy()) {
+        if let Some(snapshot) = self.snapshots.revert_to(id.to_alloy()) {
             self.inner = snapshot;
             trace!(target: "backend::memdb", "Reverted snapshot {}", id);
             true

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

@@ -622,6 +622,13 @@ impl Backend {
     pub async fn create_snapshot(&self) -> U256 {
         let num = self.best_number().as_u64();
         let hash = self.best_hash();
+
+        if let Some((id, _)) =
+            self.active_snapshots.lock().iter().find(|(_, val)| &(num, hash) == *val)
+        {
+            return *id
+        }
+
         let id = self.db.write().await.snapshot();
         trace!(target: "backend", "creating snapshot {} at {}", id, num);
         self.active_snapshots.lock().insert(id, (num, hash));
@@ -630,7 +637,7 @@ impl Backend {
 
     /// Reverts the state to the snapshot identified by the given `id`.
     pub async fn revert_snapshot(&self, id: U256) -> Result<bool, BlockchainError> {
-        let block = { self.active_snapshots.lock().remove(&id) };
+        let block = { self.active_snapshots.lock().get(&id).cloned() };
         if let Some((num, hash)) = block {
             let best_block_hash = {
                 // revert the storage that's newer than the snapshot
@@ -673,9 +680,15 @@ impl Backend {
                 ..Default::default()
             };
         }
+
+        self.active_snapshots.lock().retain(|key, _| *key <= id);
         Ok(self.db.write().await.revert(id))
     }
 
+    pub async fn list_snapshots(&self) -> HashMap<U256, (u64, H256)> {
+        self.active_snapshots.lock().clone()
+    }
+
     /// Get the current state.
     pub async fn serialized_state(&self) -> Result<SerializableState, BlockchainError> {
         let state = self.db.read().await.dump_state()?;

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

@@ -111,9 +111,8 @@ impl ForkedDatabase {
     }
 
     pub fn revert_snapshot(&mut self, id: U256) -> bool {
-        let snapshot = { self.snapshots().lock().remove_at(id) };
+        let snapshot = { self.snapshots().lock().revert_to(id) };
         if let Some(snapshot) = snapshot {
-            self.snapshots().lock().insert_at(snapshot.clone(), id);
             let ForkDbSnapshot {
                 local,
                 snapshot: StateSnapshot { accounts, storage, block_hashes },

crates/evm/core/src/snapshot.rs

@@ -10,7 +10,7 @@ pub struct Snapshots<T> {
     snapshots: HashMap<U256, T>,
 }
 
-impl<T> Snapshots<T> {
+impl<T: Clone> Snapshots<T> {
     fn next_id(&mut self) -> U256 {
         let id = self.id;
         self.id = id.saturating_add(U256::from(1));
@@ -39,6 +39,20 @@ impl<T> Snapshots<T> {
         snapshot
     }
 
+    /// Removes all snapshots after `id`.
+    ///
+    /// This will not remove the `id` itself, only clone its data, since the snapshot will still
+    /// point to a valid state. e.g.: reverting
+    /// to id 1 will delete snapshots with ids 2, 3, etc.)
+    pub fn revert_to(&mut self, id: U256) -> Option<T> {
+        let snapshot = self.snapshots.get(&id).cloned();
+
+        // revert all snapshots taken after the snapshot
+        self.snapshots.retain(|k, _| *k <= id);
+
+        snapshot
+    }
+
     /// Removes the snapshot with the given `id`.
     ///
     /// Does not remove snapshots after it.