optimize snapshotting of stack

src/callframe.rs

@@ -1,5 +1,9 @@
 use crate::{
-    decommit::is_kernel, heap::HeapId, program::Program, stack::Stack, world_diff::Snapshot,
+    decommit::is_kernel,
+    heap::HeapId,
+    program::Program,
+    stack::{Stack, StackSnapshot},
+    world_diff::Snapshot,
     Instruction,
 };
 use u256::H160;
@@ -155,10 +159,66 @@ impl Callframe {
                 .map(|f| f.previous_frame_gas)
                 .sum::<u32>()
     }
+
+    pub(crate) fn snapshot(&self) -> CallframeSnapshot {
+        CallframeSnapshot {
+            stack: self.stack.snapshot(),
+
+            context_u128: self.context_u128,
+            sp: self.sp,
+            gas: self.gas,
+            near_calls: self.near_calls.clone(),
+            heap_size: self.heap_size,
+            aux_heap_size: self.aux_heap_size,
+            heaps_i_am_keeping_alive: self.heaps_i_am_keeping_alive.clone(),
+            world_before_this_frame: self.world_before_this_frame.clone(),
+        }
+    }
+
+    pub fn rollback(&mut self, snapshot: CallframeSnapshot) {
+        let CallframeSnapshot {
+            stack,
+            context_u128,
+            sp,
+            gas,
+            near_calls,
+            heap_size,
+            aux_heap_size,
+            heaps_i_am_keeping_alive,
+            world_before_this_frame,
+        } = snapshot;
+
+        self.stack.rollback(stack);
+
+        self.context_u128 = context_u128;
+        self.sp = sp;
+        self.gas = gas;
+        self.near_calls = near_calls;
+        self.heap_size = heap_size;
+        self.aux_heap_size = aux_heap_size;
+        self.heaps_i_am_keeping_alive = heaps_i_am_keeping_alive;
+        self.world_before_this_frame = world_before_this_frame;
+    }
 }
 
 pub(crate) struct FrameRemnant {
     pub(crate) program_counter: u16,
     pub(crate) exception_handler: u16,
     pub(crate) snapshot: Snapshot,
 }
+
+/// Only contains the fields that can change (other than via tracer).
+pub struct CallframeSnapshot {
+    stack: StackSnapshot,
+
+    context_u128: u128,
+    sp: u16,
+    gas: u32,
+    near_calls: Vec<NearCallFrame>,
+
+    heap_size: u32,
+    aux_heap_size: u32,
+
+    heaps_i_am_keeping_alive: Vec<HeapId>,
+    world_before_this_frame: Snapshot,
+}

src/single_instruction_test/stack.rs

@@ -72,6 +72,14 @@ impl Stack {
             && (self.slot_written.is_none()
                 || is_valid_tagged_value((self.value_written, self.pointer_tag_written)))
     }
+
+    pub(crate) fn snapshot(&self) -> StackSnapshot {
+        unimplemented!()
+    }
+
+    pub(crate) fn rollback(&mut self, _: StackSnapshot) {
+        unimplemented!()
+    }
 }
 
 #[derive(Default, Debug)]
@@ -91,3 +99,5 @@ impl StackPool {
 
     pub fn recycle(&mut self, _: Box<Stack>) {}
 }
+
+pub(crate) struct StackSnapshot;

src/stack.rs

@@ -53,6 +53,43 @@ impl Stack {
     pub(crate) fn clear_pointer_flag(&mut self, slot: u16) {
         self.pointer_flags.clear(slot);
     }
+
+    pub(crate) fn snapshot(&self) -> StackSnapshot {
+        let mut dirty_prefix_end = 0;
+        for i in 0..NUMBER_OF_DIRTY_AREAS {
+            if self.dirty_areas & (1 << i) != 0 {
+                dirty_prefix_end = i + 1;
+            }
+        }
+
+        StackSnapshot {
+            pointer_flags: self.pointer_flags.clone(),
+            dirty_areas: self.dirty_areas,
+            slots: self.slots[..DIRTY_AREA_SIZE * dirty_prefix_end].into(),
+        }
+    }
+
+    pub(crate) fn rollback(&mut self, snapshot: StackSnapshot) {
+        let StackSnapshot {
+            pointer_flags,
+            dirty_areas,
+            slots,
+        } = snapshot;
+
+        self.zero();
+
+        self.pointer_flags = pointer_flags;
+        self.dirty_areas = dirty_areas;
+        for (me, snapshot) in self.slots.iter_mut().zip(slots.iter()) {
+            *me = *snapshot;
+        }
+    }
+}
+
+pub(crate) struct StackSnapshot {
+    pointer_flags: Bitset,
+    dirty_areas: u64,
+    slots: Box<[U256]>,
 }
 
 impl Clone for Box<Stack> {

src/state.rs

@@ -1,7 +1,8 @@
 use crate::{
     addressing_modes::Addressable,
-    callframe::Callframe,
+    callframe::{Callframe, CallframeSnapshot},
     fat_pointer::FatPointer,
+    heap::Heap,
     heap::{Heaps, CALLDATA_HEAP, FIRST_AUX_HEAP, FIRST_HEAP},
     predication::Flags,
     program::Program,
@@ -106,6 +107,42 @@ impl State {
     pub(crate) fn get_context_u128(&self) -> u128 {
         self.current_frame.context_u128
     }
+
+    pub(crate) fn snapshot(&self) -> StateSnapshot {
+        StateSnapshot {
+            registers: self.registers,
+            register_pointer_flags: self.register_pointer_flags,
+            flags: self.flags.clone(),
+            current_frame: self.current_frame.snapshot(),
+            bootloader_heap: self.heaps[self.current_frame.heap].clone(),
+            bootloader_aux_heap: self.heaps[self.current_frame.aux_heap].clone(),
+            transaction_number: self.transaction_number,
+            context_u128: self.context_u128,
+        }
+    }
+
+    pub(crate) fn rollback(&mut self, snapshot: StateSnapshot) {
+        let StateSnapshot {
+            registers,
+            register_pointer_flags,
+            flags,
+            current_frame,
+            bootloader_heap,
+            bootloader_aux_heap,
+            transaction_number,
+            context_u128,
+        } = snapshot;
+
+        self.current_frame.rollback(current_frame);
+        self.heaps[self.current_frame.heap] = bootloader_heap;
+        self.heaps[self.current_frame.aux_heap] = bootloader_aux_heap;
+
+        self.registers = registers;
+        self.register_pointer_flags = register_pointer_flags;
+        self.flags = flags;
+        self.transaction_number = transaction_number;
+        self.context_u128 = context_u128;
+    }
 }
 
 impl Addressable for State {
@@ -144,3 +181,18 @@ impl Addressable for State {
         self.current_frame.is_kernel
     }
 }
+
+pub(crate) struct StateSnapshot {
+    registers: [U256; 16],
+    register_pointer_flags: u16,
+
+    flags: Flags,
+
+    current_frame: CallframeSnapshot,
+
+    bootloader_heap: Heap,
+    bootloader_aux_heap: Heap,
+    transaction_number: u16,
+
+    context_u128: u128,
+}

src/vm.rs

@@ -1,4 +1,5 @@
 use crate::heap::HeapId;
+use crate::state::StateSnapshot;
 use crate::world_diff::ExternalSnapshot;
 use crate::{
     callframe::{Callframe, FrameRemnant},
@@ -173,19 +174,27 @@ impl VirtualMachine {
     /// # Panics
     /// Calling this function outside of the initial callframe is not allowed.
     pub fn snapshot(&self) -> VmSnapshot {
-        assert!(self.state.previous_frames.is_empty());
+        assert!(
+            self.state.previous_frames.is_empty(),
+            "Snapshotting is only allowed in the bootloader!"
+        );
         VmSnapshot {
             world_snapshot: self.world_diff.external_snapshot(),
-            state_snapshot: self.state.clone(),
+            state_snapshot: self.state.snapshot(),
         }
     }
 
     /// Returns the VM to the state it was in when the snapshot was created.
     /// # Panics
     /// Rolling back snapshots in anything but LIFO order may panic.
+    /// Rolling back outside the initial callframe will panic.
     pub fn rollback(&mut self, snapshot: VmSnapshot) {
+        assert!(
+            self.state.previous_frames.is_empty(),
+            "Rolling back is only allowed in the bootloader!"
+        );
         self.world_diff.external_rollback(snapshot.world_snapshot);
-        self.state = snapshot.state_snapshot;
+        self.state.rollback(snapshot.state_snapshot);
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -306,5 +315,5 @@ impl VirtualMachine {
 
 pub struct VmSnapshot {
     world_snapshot: ExternalSnapshot,
-    state_snapshot: State,
+    state_snapshot: StateSnapshot,
 }