make and restore snapshots

core/lib/multivm/src/lib.rs

@@ -12,8 +12,8 @@ pub use crate::{
         tracers::{MultiVMTracer, MultiVmTracerPointer},
     },
     versions::{
-        vm_1_3_2, vm_1_4_1, vm_1_4_2, vm_boojum_integration, vm_fast, vm_latest, vm_m5, vm_m6,
-        vm_refunds_enhancement, vm_virtual_blocks, era_vm,
+        era_vm, vm_1_3_2, vm_1_4_1, vm_1_4_2, vm_boojum_integration, vm_fast, vm_latest, vm_m5,
+        vm_m6, vm_refunds_enhancement, vm_virtual_blocks,
     },
     vm_instance::VmInstance,
 };

core/lib/multivm/src/versions/era_vm/snapshot.rs

@@ -1,20 +1,7 @@
+use era_vm::store::SnapShot;
 use zksync_types::H256;
 
-#[derive(Debug, Clone)]
-pub(crate) struct BootloaderStateSnapshot {
-    /// ID of the next transaction to be executed.
-    pub(crate) tx_to_execute: usize,
-    /// Stored L2 blocks in bootloader memory
-    pub(crate) l2_blocks_len: usize,
-    /// Snapshot of the last L2 block. Only this block could be changed during the rollback
-    pub(crate) last_l2_block: L2BlockSnapshot,
-    /// The number of 32-byte words spent on the already included compressed bytecodes.
-    pub(crate) compressed_bytecodes_encoding: usize,
-    /// Current offset of the free space in the bootloader memory.
-    pub(crate) free_tx_offset: usize,
-    /// Whether the pubdata information has been provided already
-    pub(crate) is_pubdata_information_provided: bool,
-}
+use super::bootloader_state::BootloaderStateSnapshot;
 
 #[derive(Debug, Clone)]
 pub(crate) struct L2BlockSnapshot {
@@ -25,11 +12,9 @@ pub(crate) struct L2BlockSnapshot {
 }
 
 pub struct VmSnapshot {
-    state: era_vm::VMState,
-
-    // TODO: Implement snapshots in era vm
-    // world_snapshot: vm2::ExternalSnapshot,
-    bootloader_snapshot: BootloaderStateSnapshot,
-    suspended_at: u16,
-    gas_for_account_validation: u32,
+    pub(crate) storage_snapshot: SnapShot,
+    pub(crate) transient_storage_snapshot: SnapShot,
+    pub(crate) bootloader_snapshot: BootloaderStateSnapshot,
+    pub(crate) suspended_at: u16,
+    pub(crate) gas_for_account_validation: u32,
 }

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

@@ -8,7 +8,7 @@ use std::{
 };
 
 use era_vm::{
-    store::{L2ToL1Log, StorageError, StorageKey as EraStorageKey},
+    store::{L2ToL1Log, SnapShot, StorageError, StorageKey as EraStorageKey},
     vm::ExecutionOutput,
     EraVM, VMState,
 };
@@ -60,7 +60,7 @@ pub struct Vm<S: ReadStorage> {
     pub(crate) batch_env: L1BatchEnv,
     pub(crate) system_env: SystemEnv,
 
-    snapshots: Vec<VmSnapshot>, // TODO: Implement snapshots logic
+    snapshot: Option<VmSnapshot>, // TODO: Implement snapshots logic
 }
 
 /// Encapsulates creating VM instance based on the provided environment.
@@ -149,7 +149,7 @@ impl<S: ReadStorage + 'static> VmFactory<S> for Vm<S> {
             storage,
             batch_env,
             system_env,
-            snapshots: Vec::new(),
+            snapshot: None,
         };
 
         mv.write_to_bootloader_heap(bootloader_memory);
@@ -471,15 +471,45 @@ impl<S: ReadStorage + 'static> VmInterface for Vm<S> {
 
 impl<S: ReadStorage + 'static> VmInterfaceHistoryEnabled for Vm<S> {
     fn make_snapshot(&mut self) {
-        todo!()
+        assert!(
+            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 {
+            storage_snapshot: self.inner.state_storage.create_snapshot(),
+            transient_storage_snapshot: self.inner.transient_storage.create_snapshot(),
+            bootloader_snapshot: self.bootloader_state.get_snapshot(),
+            suspended_at: self.suspended_at,
+            gas_for_account_validation: self.gas_for_account_validation,
+        });
     }
 
     fn rollback_to_the_latest_snapshot(&mut self) {
-        todo!()
+        let VmSnapshot {
+            storage_snapshot,
+            transient_storage_snapshot,
+            bootloader_snapshot,
+            suspended_at,
+            gas_for_account_validation,
+        } = self.snapshot.take().expect("no snapshots to rollback to");
+
+        self.inner.state_storage.rollback(&storage_snapshot);
+        self.inner
+            .transient_storage
+            .rollback(&transient_storage_snapshot);
+
+        self.bootloader_state.apply_snapshot(bootloader_snapshot);
+        self.suspended_at = suspended_at;
+        self.gas_for_account_validation = gas_for_account_validation;
+
+        // self.delete_history_if_appropriate();
     }
 
     fn pop_snapshot_no_rollback(&mut self) {
-        todo!()
+        self.snapshot = None;
+        // self.delete_history_if_appropriate();
     }
 }
 

core/lib/multivm/src/versions/mod.rs

@@ -1,5 +1,5 @@
-pub mod shadow;
 pub mod era_vm;
+pub mod shadow;
 pub mod vm_1_3_2;
 pub mod vm_1_4_1;
 pub mod vm_1_4_2;