test(vm): Fix fast VM unit tests

core/lib/multivm/Cargo.toml

@@ -16,7 +16,7 @@ zk_evm_1_4_1.workspace = true
 zk_evm_1_4_0.workspace = true
 zk_evm_1_3_3.workspace = true
 zk_evm_1_3_1.workspace = true
-vm2 = { git = "https://github.com/matter-labs/vm2.git", rev = "28770597a3f150dbe4373cb57929bd8db82e884f" }
+vm2 = { git = "https://github.com/matter-labs/vm2.git", rev = "a0cf04b03ac1c486a48e5f2e32422a00c27a1b9d" }
 
 circuit_sequencer_api_1_3_3.workspace = true
 circuit_sequencer_api_1_4_0.workspace = true

core/lib/multivm/src/versions/vm_fast/tests/rollbacks.rs

@@ -10,7 +10,6 @@ use crate::{
     },
 };
 
-#[ignore] // FIXME(PLA-1002): VM state assertion in `execute_tx_and_verify` fails
 #[test]
 fn test_vm_rollbacks() {
     let mut vm = VmTesterBuilder::new()
@@ -60,7 +59,6 @@ fn test_vm_rollbacks() {
     assert_eq!(result_without_rollbacks, result_with_rollbacks);
 }
 
-#[ignore] // FIXME(PLA-1002): VM state assertion in `execute_tx_and_verify` fails
 #[test]
 fn test_vm_loadnext_rollbacks() {
     let mut vm = VmTesterBuilder::new()

core/lib/multivm/src/versions/vm_fast/tests/tester/transaction_test_info.rs

@@ -1,9 +1,13 @@
-use zksync_types::{ExecuteTransactionCommon, Transaction};
+use zksync_state::ReadStorage;
+use zksync_types::{ExecuteTransactionCommon, Transaction, H160, U256};
 
 use super::VmTester;
-use crate::interface::{
-    CurrentExecutionState, ExecutionResult, Halt, TxRevertReason, VmExecutionMode,
-    VmExecutionResultAndLogs, VmInterface, VmInterfaceHistoryEnabled, VmRevertReason,
+use crate::{
+    interface::{
+        CurrentExecutionState, ExecutionResult, Halt, TxRevertReason, VmExecutionMode,
+        VmExecutionResultAndLogs, VmInterface, VmInterfaceHistoryEnabled, VmRevertReason,
+    },
+    vm_fast::Vm,
 };
 
 #[derive(Debug, Clone)]
@@ -179,6 +183,30 @@ impl TransactionTestInfo {
     }
 }
 
+// TODO this doesn't include all the state of ModifiedWorld
+#[derive(Debug, PartialEq)]
+struct VmStateDump {
+    state: vm2::State,
+    storage_writes: Vec<((H160, U256), U256)>,
+    events: Box<[vm2::Event]>,
+}
+
+impl<S: ReadStorage> Vm<S> {
+    fn dump_state(&self) -> VmStateDump {
+        VmStateDump {
+            state: self.inner.state.clone(),
+            storage_writes: self
+                .inner
+                .world_diff
+                .get_storage_state()
+                .iter()
+                .map(|(k, v)| (*k, *v))
+                .collect(),
+            events: self.inner.world_diff.events().into(),
+        }
+    }
+}
+
 impl VmTester {
     pub(crate) fn execute_and_verify_txs(
         &mut self,
@@ -197,16 +225,17 @@ impl VmTester {
         &mut self,
         tx_test_info: TransactionTestInfo,
     ) -> VmExecutionResultAndLogs {
-        let inner_state_before = self.vm.dump_state();
         self.vm.make_snapshot();
+        let inner_state_before = self.vm.dump_state();
         self.vm.push_transaction(tx_test_info.tx.clone());
         let result = self.vm.execute(VmExecutionMode::OneTx);
         tx_test_info.verify_result(&result);
         if tx_test_info.should_rollback() {
             self.vm.rollback_to_the_latest_snapshot();
             let inner_state_after = self.vm.dump_state();
-            assert_eq!(
-                inner_state_before, inner_state_after,
+            pretty_assertions::assert_eq!(
+                inner_state_before,
+                inner_state_after,
                 "Inner state before and after rollback should be equal"
             );
         } else {

core/lib/multivm/src/versions/vm_fast/tests/tracing_execution_error.rs

@@ -9,7 +9,6 @@ use crate::{
 };
 
 #[test]
-#[ignore] // FIXME(PLA-1002): VM state assertion in `execute_tx_and_verify` fails
 fn test_tracing_of_execution_errors() {
     let contract_address = H160::random();
     let mut vm = VmTesterBuilder::new()

core/lib/multivm/src/versions/vm_fast/vm.rs

@@ -337,23 +337,6 @@ impl<S: ReadStorage> Vm<S> {
         }
     }
 
-    /// Returns the current state of the VM in a format that can be compared for equality.
-    #[cfg(test)]
-    #[allow(clippy::type_complexity)] // OK for tests
-    pub(crate) fn dump_state(&self) -> (vm2::State, Vec<((H160, U256), U256)>, Box<[vm2::Event]>) {
-        // TODO this doesn't include all the state of ModifiedWorld
-        (
-            self.inner.state.clone(),
-            self.inner
-                .world_diff
-                .get_storage_state()
-                .iter()
-                .map(|(k, v)| (*k, *v))
-                .collect(),
-            self.inner.world_diff.events().into(),
-        )
-    }
-
     pub(crate) fn push_transaction_inner(
         &mut self,
         tx: zksync_types::Transaction,
@@ -480,6 +463,12 @@ impl<S: ReadStorage> Vm<S> {
 
         me
     }
+
+    fn delete_history_if_appropriate(&mut self) {
+        if self.snapshot.is_none() && self.inner.state.previous_frames.is_empty() {
+            self.inner.delete_history();
+        }
+    }
 }
 
 impl<S: ReadStorage> VmInterface for Vm<S> {
@@ -672,6 +661,7 @@ impl<S: ReadStorage> VmInterface for Vm<S> {
     }
 }
 
+#[derive(Debug)]
 struct VmSnapshot {
     vm_snapshot: vm2::Snapshot,
     bootloader_snapshot: BootloaderStateSnapshot,
@@ -685,6 +675,8 @@ impl<S: ReadStorage> VmInterfaceHistoryEnabled for Vm<S> {
             self.snapshot.is_none(),
             "cannot create a VM snapshot until a previous snapshot is rolled back to or popped"
         );
+
+        self.delete_history_if_appropriate();
         self.snapshot = Some(VmSnapshot {
             vm_snapshot: self.inner.snapshot(),
             bootloader_snapshot: self.bootloader_state.get_snapshot(),
@@ -715,14 +707,6 @@ impl<S: ReadStorage> VmInterfaceHistoryEnabled for Vm<S> {
     }
 }
 
-impl<S: ReadStorage> Vm<S> {
-    fn delete_history_if_appropriate(&mut self) {
-        if self.snapshot.is_none() && self.inner.state.previous_frames.is_empty() {
-            self.inner.delete_history();
-        }
-    }
-}
-
 impl<S: fmt::Debug> fmt::Debug for Vm<S> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("Vm")