Fix snapshot issues

crates/forge/tests/it/revive/cheat_snapshot.rs

@@ -8,9 +8,30 @@ use rstest::rstest;
 #[case::pvm(ReviveRuntimeMode::Pvm)]
 #[case::evm(ReviveRuntimeMode::Evm)]
 #[tokio::test(flavor = "multi_thread")]
-async fn test_snapshot_cheats(#[case] runtime_mode: ReviveRuntimeMode) {
+async fn test_snapshot_state(#[case] runtime_mode: ReviveRuntimeMode) {
     let runner: forge::MultiContractRunner = TEST_DATA_REVIVE.runner_revive(runtime_mode);
     let filter = Filter::new(".*", "StateSnapshotTest", ".*/revive/.*");
 
     TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
 }
+
+#[rstest]
+#[case::pvm(ReviveRuntimeMode::Pvm)]
+#[case::evm(ReviveRuntimeMode::Evm)]
+#[tokio::test(flavor = "multi_thread")]
+async fn test_snapshot_constructor_contract(#[case] runtime_mode: ReviveRuntimeMode) {
+    let runner: forge::MultiContractRunner = TEST_DATA_REVIVE.runner_revive(runtime_mode);
+    let filter = Filter::new(".*", "SnapshotConstructorContractTest", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}
+
+#[rstest]
+#[case::evm(ReviveRuntimeMode::Evm)]
+#[tokio::test(flavor = "multi_thread")]
+async fn test_snapshot_across_mode_switch(#[case] runtime_mode: ReviveRuntimeMode) {
+    let runner: forge::MultiContractRunner = TEST_DATA_REVIVE.runner_revive(runtime_mode);
+    let filter = Filter::new(".*", "SnapshotAcrossModeSwitchTest", ".*/revive/.*");
+
+    TestConfig::with_filter(runner, filter).spec_id(SpecId::PRAGUE).run().await;
+}

crates/revive-strategy/src/cheatcodes/mod.rs

@@ -8,10 +8,11 @@ use foundry_cheatcodes::{
     CheatcodeInspectorStrategyContext, CheatcodeInspectorStrategyRunner, CheatsConfig, CheatsCtxt,
     CommonCreateInput, DynCheatcode, Ecx, EvmCheatcodeInspectorStrategyRunner, Result,
     Vm::{
-        AccountAccessKind, chainIdCall, coinbaseCall, dealCall, etchCall, getNonce_0Call, loadCall,
-        polkadot_0Call, polkadot_1Call, polkadotSkipCall, resetNonceCall,
-        revertToStateAndDeleteCall, revertToStateCall, rollCall, setBlockhashCall, setNonceCall,
-        setNonceUnsafeCall, snapshotStateCall, storeCall, warpCall,
+        AccountAccessKind, chainIdCall, coinbaseCall, dealCall, deleteStateSnapshotCall,
+        deleteStateSnapshotsCall, etchCall, getNonce_0Call, loadCall, polkadot_0Call,
+        polkadot_1Call, polkadotSkipCall, resetNonceCall, revertToStateAndDeleteCall,
+        revertToStateCall, rollCall, setBlockhashCall, setNonceCall, setNonceUnsafeCall,
+        snapshotStateCall, storeCall, warpCall,
     },
     journaled_account, precompile_error,
 };
@@ -368,21 +369,42 @@ impl CheatcodeInspectorStrategyRunner for PvmCheatcodeInspectorStrategyRunner {
 
                 rollCall { newHeight: clamped_height }.dyn_apply(ccx, executor)
             }
-            t if using_revive && is::<snapshotStateCall>(t) => {
-                ctx.externalities.start_snapshotting();
-                cheatcode.dyn_apply(ccx, executor)
+            t if is::<snapshotStateCall>(t) => {
+                let result = cheatcode.dyn_apply(ccx, executor);
+                if let Ok(ref encoded) = result
+                    && let Ok(snapshot_id) = U256::abi_decode(encoded)
+                {
+                    let ctx = get_context_ref_mut(ccx.state.strategy.context.as_mut());
+                    ctx.externalities.start_snapshotting(snapshot_id);
+                }
+                result
             }
-            t if using_revive && is::<revertToStateAndDeleteCall>(t) => {
+            t if is::<revertToStateAndDeleteCall>(t) => {
                 let &revertToStateAndDeleteCall { snapshotId } =
                     cheatcode.as_any().downcast_ref().unwrap();
-
-                ctx.externalities.revert(snapshotId.try_into().unwrap());
-                cheatcode.dyn_apply(ccx, executor)
+                let ctx = get_context_ref_mut(ccx.state.strategy.context.as_mut());
+                ctx.externalities.revert(snapshotId);
+                let result = cheatcode.dyn_apply(ccx, executor);
+                let ctx = get_context_ref_mut(ccx.state.strategy.context.as_mut());
+                ctx.externalities.delete_snapshot(snapshotId);
+                result
             }
-            t if using_revive && is::<revertToStateCall>(t) => {
+            t if is::<revertToStateCall>(t) => {
                 let &revertToStateCall { snapshotId } = cheatcode.as_any().downcast_ref().unwrap();
-
-                ctx.externalities.revert(snapshotId.try_into().unwrap());
+                let ctx = get_context_ref_mut(ccx.state.strategy.context.as_mut());
+                ctx.externalities.revert(snapshotId);
+                cheatcode.dyn_apply(ccx, executor)
+            }
+            t if is::<deleteStateSnapshotCall>(t) => {
+                let &deleteStateSnapshotCall { snapshotId } =
+                    cheatcode.as_any().downcast_ref().unwrap();
+                let ctx = get_context_ref_mut(ccx.state.strategy.context.as_mut());
+                ctx.externalities.delete_snapshot(snapshotId);
+                cheatcode.dyn_apply(ccx, executor)
+            }
+            t if is::<deleteStateSnapshotsCall>(t) => {
+                let ctx = get_context_ref_mut(ccx.state.strategy.context.as_mut());
+                ctx.externalities.delete_all_snapshots();
                 cheatcode.dyn_apply(ccx, executor)
             }
             t if using_revive && is::<warpCall>(t) => {

crates/revive-strategy/src/state.rs

@@ -17,13 +17,16 @@ use polkadot_sdk::{
 };
 use revive_env::{Balances, BlockAuthor, ExtBuilder, NativeToEthRatio, Runtime, System, Timestamp};
 use std::{
+    collections::HashMap,
     fmt::Debug,
     sync::{Arc, Mutex},
 };
 
 pub(crate) struct Inner {
     pub externalities: TestExternalities,
     pub depth: usize,
+    /// Maps REVM snapshot_id to the pallet-revive transaction depth at snapshot time.
+    pub snapshot_depths: HashMap<U256, usize>,
 }
 
 #[derive(Default)]
@@ -39,6 +42,7 @@ impl Default for Inner {
                 )])
                 .build(),
             depth: 0,
+            snapshot_depths: HashMap::new(),
         }
     }
 }
@@ -51,9 +55,9 @@ impl Debug for TestEnv {
 
 impl Clone for TestEnv {
     fn clone(&self) -> Self {
+        let mut state = self.0.lock().unwrap();
         let mut inner: Inner = Default::default();
-        inner.externalities.backend = self.0.lock().unwrap().externalities.as_backend();
-        inner.depth = self.0.lock().unwrap().depth;
+        inner.externalities.backend = state.externalities.as_backend();
         Self(Arc::new(Mutex::new(inner)))
     }
 }
@@ -63,20 +67,60 @@ impl TestEnv {
         Self(self.0.clone())
     }
 
-    pub fn start_snapshotting(&mut self) {
+    pub fn start_snapshotting(&mut self, snapshot_id: U256) {
         let mut state = self.0.lock().unwrap();
-        state.depth += 1;
+        let current_depth = state.depth;
+        state.snapshot_depths.insert(snapshot_id, current_depth);
         state.externalities.ext().storage_start_transaction();
+        state.depth += 1;
     }
 
-    pub fn revert(&mut self, depth: usize) {
+    pub fn revert(&mut self, snapshot_id: U256) {
         let mut state = self.0.lock().unwrap();
-        while state.depth > depth + 1 {
-            state.externalities.ext().storage_rollback_transaction().unwrap();
-            state.depth -= 1;
+
+        let target_depth = match state.snapshot_depths.get(&snapshot_id) {
+            Some(&depth) => depth,
+            None => {
+                // Unknown snapshot - reset pallet-revive completely.
+                // This can happen with cross-function snapshots (Clone committed transactions)
+                // in setUp or test contract constructor call
+                tracing::warn!(
+                    snapshot_id = ?snapshot_id,
+                    current_depth = state.depth,
+                    "snapshot not found, resetting pallet-revive to sync with REVM"
+                );
+                while state.depth > 0 {
+                    let _ = state.externalities.ext().storage_rollback_transaction();
+                    state.depth -= 1;
+                }
+                state.snapshot_depths.clear();
+                return;
+            }
+        };
+
+        let rollbacks_needed = state.depth.saturating_sub(target_depth);
+        for _ in 0..rollbacks_needed {
+            if state.depth > 0 {
+                let _ = state.externalities.ext().storage_rollback_transaction();
+                state.depth -= 1;
+            }
         }
-        state.externalities.ext().storage_rollback_transaction().unwrap();
+
+        // Remove snapshots that are now invalid (taken after the target snapshot)
+        state.snapshot_depths.retain(|_, &mut depth| depth <= target_depth);
+
         state.externalities.ext().storage_start_transaction();
+        state.depth = target_depth + 1;
+    }
+
+    pub fn delete_snapshot(&mut self, snapshot_id: U256) {
+        let mut state = self.0.lock().unwrap();
+        state.snapshot_depths.remove(&snapshot_id);
+    }
+
+    pub fn delete_all_snapshots(&mut self) {
+        let mut state = self.0.lock().unwrap();
+        state.snapshot_depths.clear();
     }
 
     pub fn execute_with<R, F: FnOnce() -> R>(&mut self, f: F) -> R {