diff --git a/crates/forge/tests/it/revive/cheat_snapshot.rs b/crates/forge/tests/it/revive/cheat_snapshot.rs index 36b04cabc91f3..174ae495198d5 100644 --- a/crates/forge/tests/it/revive/cheat_snapshot.rs +++ b/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; +} diff --git a/crates/revive-strategy/src/cheatcodes/mod.rs b/crates/revive-strategy/src/cheatcodes/mod.rs index 5ad87b4459d4d..dd66dfa4ccb02 100644 --- a/crates/revive-strategy/src/cheatcodes/mod.rs +++ b/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::(t) => { - ctx.externalities.start_snapshotting(); - cheatcode.dyn_apply(ccx, executor) + t if is::(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::(t) => { + t if is::(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::(t) => { + t if is::(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::(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::(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::(t) => { diff --git a/crates/revive-strategy/src/state.rs b/crates/revive-strategy/src/state.rs index aa870401f2d55..36e03a1d32e40 100644 --- a/crates/revive-strategy/src/state.rs +++ b/crates/revive-strategy/src/state.rs @@ -17,6 +17,7 @@ use polkadot_sdk::{ }; use revive_env::{Balances, BlockAuthor, ExtBuilder, NativeToEthRatio, Runtime, System, Timestamp}; use std::{ + collections::HashMap, fmt::Debug, sync::{Arc, Mutex}, }; @@ -24,6 +25,8 @@ use std::{ 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, } #[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>(&mut self, f: F) -> R { diff --git a/testdata/default/revive/Snapshot.t.sol b/testdata/default/revive/Snapshot.t.sol index 4e3928967a9eb..26bd95cc68e6e 100644 --- a/testdata/default/revive/Snapshot.t.sol +++ b/testdata/default/revive/Snapshot.t.sol @@ -151,4 +151,285 @@ contract StateSnapshotTest is DSTest { assertEq(store.blockNumber(), num, "snapshot revert for block.number unsuccessful"); assertEq(store.blockTimestamp(), time, "snapshot revert for block.timestamp unsuccessful"); } + + function testOutOfOrderRevert() public { + uint256 snap1 = vm.snapshotState(); + store.setSlots(300, 400); + + vm.snapshotState(); + store.setSlots(500, 600); + + vm.snapshotState(); + store.setSlots(700, 800); + + assertEq(store.slot0(), 700, "should be at 700"); + + vm.revertToState(snap1); + assertEq(store.slot0(), 10, "out-of-order revert to snap1 failed"); + + store.setSlots(999, 888); + assertEq(store.slot0(), 999, "post-revert write failed"); + assertEq(store.slot1(), 888, "post-revert write failed for slot1"); + } + + function testSnapshotAfterOutOfOrderRevert() public { + uint256 snap1 = vm.snapshotState(); + store.setSlots(300, 400); + + vm.snapshotState(); + store.setSlots(500, 600); + + vm.snapshotState(); + store.setSlots(700, 800); + + vm.revertToState(snap1); + assertEq(store.slot0(), 10, "initial revert failed"); + + uint256 newSnap = vm.snapshotState(); + store.setSlots(999, 888); + assertEq(store.slot0(), 999, "write after new snapshot failed"); + + vm.revertToState(newSnap); + assertEq(store.slot0(), 10, "revert to new snapshot failed - pallet-revive depth may be wrong"); + assertEq(store.slot1(), 20, "revert to new snapshot failed for slot1"); + } + + function testRevertChangeRevertAgain() public { + uint256 snapshotId = vm.snapshotState(); + + store.setSlots(300, 400); + assertEq(store.slot0(), 300, "should be 300 after first change"); + + vm.revertToState(snapshotId); + assertEq(store.slot0(), 10, "should restore to 10 after first revert"); + + store.setSlots(500, 600); + assertEq(store.slot0(), 500, "should be 500 after second change"); + + vm.revertToState(snapshotId); + assertEq(store.slot0(), 10, "should restore to 10 after second revert"); + assertEq(store.slot1(), 20, "should restore to 20 after second revert"); + + store.setSlots(700, 800); + assertEq(store.slot0(), 700, "should be 700 after third change"); + + vm.revertToState(snapshotId); + assertEq(store.slot0(), 10, "should restore to 10 after third revert"); + assertEq(store.slot1(), 20, "should restore to 20 after third revert"); + } + + function testNewContractRollback() public { + uint256 snapshotId = vm.snapshotState(); + + Storage newStore = new Storage(); + address newAddr = address(newStore); + newStore.setSlots(777, 888); + assertEq(newStore.slot0(), 777, "new contract should have value"); + + uint256 codeSizeBefore = newAddr.code.length; + assertTrue(codeSizeBefore > 0, "contract should have code before revert"); + + vm.revertToState(snapshotId); + + uint256 codeSizeAfter = newAddr.code.length; + assertEq(codeSizeAfter, 0, "contract code should be gone after revert"); + } + + function testExistingContractStorageRollback() public { + assertEq(store.slot0(), 10, "initial state"); + + uint256 snapshotId = vm.snapshotState(); + + store.setSlots(999, 888); + assertEq(store.slot0(), 999, "modified state"); + + Storage newStore = new Storage(); + newStore.setSlots(111, 222); + + vm.revertToState(snapshotId); + + assertEq(store.slot0(), 10, "existing contract should be restored"); + assertEq(store.slot1(), 20, "existing contract slot1 should be restored"); + } +} + +contract SnapshotAcrossModeSwitchTest is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + + Storage store; + + function setUp() public { + store = new Storage(); + store.setSlots(100, 200); + } + + function testSnapshotInRevmThenSwitchToPolkadot() public { + vm.polkadot(false); + + uint256 snapshotId = vm.snapshotState(); + + store.setSlots(300, 400); + assertEq(store.slot0(), 300, "REVM should have 300"); + + vm.polkadot(true); + + store.setSlots(500, 600); + assertEq(store.slot0(), 500, "pallet-revive should have 500"); + + vm.revertToState(snapshotId); + + assertEq(store.slot0(), 100, "should restore to 100 after revert"); + assertEq(store.slot1(), 200, "should restore to 200 after revert"); + } + + function testSnapshotInPolkadotThenSwitchToRevm() public { + uint256 snapshotId = vm.snapshotState(); + + store.setSlots(300, 400); + assertEq(store.slot0(), 300, "pallet-revive should have 300"); + + vm.polkadot(false); + + store.setSlots(500, 600); + assertEq(store.slot0(), 500, "REVM should have 500"); + + vm.revertToState(snapshotId); + + assertEq(store.slot0(), 100, "should restore to 100 after revert"); + assertEq(store.slot1(), 200, "should restore to 200 after revert"); + } + + function testNewSnapshotAfterModeSwitchBack() public { + vm.polkadot(false); + + uint256 revmSnap = vm.snapshotState(); + + store.setSlots(300, 400); + assertEq(store.slot0(), 300, "REVM should have 300"); + + vm.polkadot(true); + + uint256 polkadotSnap = vm.snapshotState(); + + store.setSlots(500, 600); + assertEq(store.slot0(), 500, "pallet-revive should have 500"); + + vm.revertToState(polkadotSnap); + assertEq(store.slot0(), 300, "polkadotSnap should restore to 300"); + + vm.revertToState(revmSnap); + assertEq(store.slot0(), 100, "revmSnap should restore to 100"); + assertEq(store.slot1(), 200, "revmSnap should restore to 200"); + } + + function testMultipleModeSwitches() public { + assertEq(store.slot0(), 100, "initial slot0"); + + vm.polkadot(false); + store.setSlots(200, 201); + assertEq(store.slot0(), 200, "REVM cycle 1"); + + vm.polkadot(true); + store.setSlots(300, 301); + assertEq(store.slot0(), 300, "pallet-revive cycle 1"); + + vm.polkadot(false); + store.setSlots(400, 401); + assertEq(store.slot0(), 400, "REVM cycle 2"); + + vm.polkadot(true); + store.setSlots(500, 501); + assertEq(store.slot0(), 500, "pallet-revive cycle 2"); + + vm.polkadot(false); + store.setSlots(600, 601); + assertEq(store.slot0(), 600, "REVM cycle 3"); + + vm.polkadot(true); + store.setSlots(700, 701); + assertEq(store.slot0(), 700, "pallet-revive cycle 3"); + + assertEq(store.slot0(), 700, "final value should be 700"); + } + + function testMultipleModeSwitchesWithSnapshot() public { + uint256 initialSnap = vm.snapshotState(); + + vm.polkadot(false); + store.setSlots(200, 201); + + vm.polkadot(true); + store.setSlots(300, 301); + + vm.polkadot(false); + store.setSlots(400, 401); + + vm.polkadot(true); + store.setSlots(500, 501); + + assertEq(store.slot0(), 500, "should be 500 after switches"); + + vm.revertToState(initialSnap); + + assertEq(store.slot0(), 100, "should restore to initial 100"); + assertEq(store.slot1(), 200, "should restore to initial 200"); + } + + function testStatePersistsAcrossMigrations() public { + store.setSlots(111, 222); + + vm.polkadot(false); + assertEq(store.slot0(), 111, "REVM should have migrated value"); + + vm.polkadot(true); + assertEq(store.slot0(), 111, "pallet-revive should have value back"); + + vm.polkadot(false); + store.setSlots(333, 444); + + vm.polkadot(true); + assertEq(store.slot0(), 333, "pallet-revive should have REVM changes"); + + vm.polkadot(false); + vm.polkadot(true); + assertEq(store.slot0(), 333, "value should persist through empty migration"); + } +} + +contract SnapshotConstructorContractTest is DSTest { + Vm constant vm = Vm(HEVM_ADDRESS); + + Storage constructorStore; + Storage setupStore; + + constructor() { + constructorStore = new Storage(); + constructorStore.setSlots(1, 2); + } + + function setUp() public { + setupStore = new Storage(); + setupStore.setSlots(10, 20); + } + + function testConstructorContractPreserved() public { + assertEq(constructorStore.slot0(), 1, "constructor contract initial slot0"); + assertEq(constructorStore.slot1(), 2, "constructor contract initial slot1"); + assertEq(setupStore.slot0(), 10, "setUp contract initial slot0"); + + uint256 snapshotId = vm.snapshotState(); + + constructorStore.setSlots(100, 200); + setupStore.setSlots(300, 400); + + assertEq(constructorStore.slot0(), 100, "constructor contract modified"); + assertEq(setupStore.slot0(), 300, "setUp contract modified"); + + vm.revertToState(snapshotId); + + assertEq(constructorStore.slot0(), 1, "constructor contract should be restored"); + assertEq(constructorStore.slot1(), 2, "constructor contract slot1 should be restored"); + assertEq(setupStore.slot0(), 10, "setUp contract should be restored"); + assertEq(setupStore.slot1(), 20, "setUp contract slot1 should be restored"); + } }