perf: optimize external snapshots (#47)

The bootloader only uses a small part of the stack and only changes
small areas of its heap, yet both are copied entirely when snapshotting.
This fixes that.

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,69 @@ 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_was_keeping_alive: self.heaps_i_am_keeping_alive.len(),
+        }
+    }
+
+    /// Returns heaps that were created during the period that is rolled back
+    /// and thus can't be referenced anymore and should be deallocated.
+    pub(crate) fn rollback(
+        &mut self,
+        snapshot: CallframeSnapshot,
+    ) -> impl Iterator<Item = HeapId> + '_ {
+        let CallframeSnapshot {
+            stack,
+            context_u128,
+            sp,
+            gas,
+            near_calls,
+            heap_size,
+            aux_heap_size,
+            heaps_i_was_keeping_alive,
+        } = 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
+            .drain(heaps_i_was_keeping_alive..)
+    }
 }
 
 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(crate) struct CallframeSnapshot {
+    stack: StackSnapshot,
+
+    context_u128: u128,
+    sp: u16,
+    gas: u32,
+    near_calls: Vec<NearCallFrame>,
+
+    heap_size: u32,
+    aux_heap_size: u32,
+
+    heaps_i_was_keeping_alive: usize,
+}

src/heap.rs

@@ -1,5 +1,5 @@
 use crate::instruction_handlers::HeapInterface;
-use std::ops::{Index, IndexMut, Range};
+use std::ops::{Index, Range};
 use u256::U256;
 use zkevm_opcode_defs::system_params::NEW_FRAME_MEMORY_STIPEND;
 
@@ -21,8 +21,17 @@ impl HeapId {
 pub struct Heap(Vec<u8>);
 
 impl Heap {
-    pub fn reserve(&mut self, additional: usize) {
-        self.0.reserve_exact(additional);
+    fn write_u256(&mut self, start_address: u32, value: U256) {
+        let end = (start_address + 32) as usize;
+        if end > self.0.len() {
+            self.0.resize(end, 0);
+        }
+
+        value.to_big_endian(&mut self.0[start_address as usize..end]);
+    }
+
+    pub(crate) fn is_empty(&self) -> bool {
+        self.0.is_empty()
     }
 }
 
@@ -39,16 +48,6 @@ impl HeapInterface for Heap {
         }
         U256::from_big_endian(&bytes)
     }
-    fn write_u256(&mut self, start_address: u32, value: U256) {
-        let end = (start_address + 32) as usize;
-        if end > self.0.len() {
-            self.0.resize(end, 0);
-        }
-
-        let mut bytes = [0; 32];
-        value.to_big_endian(&mut bytes);
-        self.0[start_address as usize..end].copy_from_slice(&bytes);
-    }
     fn read_range_big_endian(&self, range: Range<u32>) -> Vec<u8> {
         let end = (range.end as usize).min(self.0.len());
         let mut result = vec![0; range.len()];
@@ -57,13 +56,13 @@ impl HeapInterface for Heap {
         }
         result
     }
-    fn memset(&mut self, src: &[u8]) {
-        self.0 = src.to_vec();
-    }
 }
 
 #[derive(Debug, Clone)]
-pub struct Heaps(Vec<Heap>);
+pub struct Heaps {
+    heaps: Vec<Heap>,
+    bootloader_heap_rollback_info: Vec<(u32, U256)>,
+}
 
 pub(crate) const CALLDATA_HEAP: HeapId = HeapId(1);
 pub const FIRST_HEAP: HeapId = HeapId(2);
@@ -73,44 +72,67 @@ impl Heaps {
     pub(crate) fn new(calldata: Vec<u8>) -> Self {
         // The first heap can never be used because heap zero
         // means the current heap in precompile calls
-        Self(vec![
-            Heap(vec![]),
-            Heap(calldata),
-            Heap(vec![]),
-            Heap(vec![]),
-        ])
+        Self {
+            heaps: vec![Heap(vec![]), Heap(calldata), Heap(vec![]), Heap(vec![])],
+            bootloader_heap_rollback_info: vec![],
+        }
     }
 
     pub(crate) fn allocate(&mut self) -> HeapId {
-        let id = HeapId(self.0.len() as u32);
-        self.0
-            .push(Heap(vec![0; NEW_FRAME_MEMORY_STIPEND as usize]));
+        self.allocate_inner(vec![0; NEW_FRAME_MEMORY_STIPEND as usize])
+    }
+
+    pub(crate) fn allocate_with_content(&mut self, content: &[u8]) -> HeapId {
+        self.allocate_inner(content.to_vec())
+    }
+
+    fn allocate_inner(&mut self, memory: Vec<u8>) -> HeapId {
+        let id = HeapId(self.heaps.len() as u32);
+        self.heaps.push(Heap(memory));
         id
     }
 
     pub(crate) fn deallocate(&mut self, heap: HeapId) {
-        self.0[heap.0 as usize].0 = vec![];
+        self.heaps[heap.0 as usize].0 = vec![];
+    }
+
+    pub fn write_u256(&mut self, heap: HeapId, start_address: u32, value: U256) {
+        if heap == FIRST_HEAP {
+            self.bootloader_heap_rollback_info
+                .push((start_address, self[heap].read_u256(start_address)));
+        }
+        self.heaps[heap.0 as usize].write_u256(start_address, value);
+    }
+
+    pub(crate) fn snapshot(&self) -> usize {
+        self.bootloader_heap_rollback_info.len()
+    }
+
+    pub(crate) fn rollback(&mut self, snapshot: usize) {
+        for (address, value) in self.bootloader_heap_rollback_info.drain(snapshot..).rev() {
+            self.heaps[FIRST_HEAP.0 as usize].write_u256(address, value);
+        }
+    }
+
+    pub(crate) fn delete_history(&mut self) {
+        self.bootloader_heap_rollback_info.clear();
     }
 }
 
 impl Index<HeapId> for Heaps {
     type Output = Heap;
 
     fn index(&self, index: HeapId) -> &Self::Output {
-        &self.0[index.0 as usize]
-    }
-}
-
-impl IndexMut<HeapId> for Heaps {
-    fn index_mut(&mut self, index: HeapId) -> &mut Self::Output {
-        &mut self.0[index.0 as usize]
+        &self.heaps[index.0 as usize]
     }
 }
 
 impl PartialEq for Heaps {
     fn eq(&self, other: &Self) -> bool {
-        for i in 0..self.0.len().max(other.0.len()) {
-            if self.0.get(i).unwrap_or(&Heap(vec![])) != other.0.get(i).unwrap_or(&Heap(vec![])) {
+        for i in 0..self.heaps.len().max(other.heaps.len()) {
+            if self.heaps.get(i).unwrap_or(&Heap(vec![]))
+                != other.heaps.get(i).unwrap_or(&Heap(vec![]))
+            {
                 return false;
             }
         }

src/instruction_handlers/decommit.rs

@@ -8,7 +8,7 @@ use crate::{
     Instruction, VirtualMachine, World,
 };
 
-use super::{common::instruction_boilerplate, HeapInterface};
+use super::common::instruction_boilerplate;
 
 fn decommit(
     vm: &mut VirtualMachine,
@@ -38,9 +38,8 @@ fn decommit(
             vm.state.current_frame.gas += extra_cost;
         }
 
-        let heap = vm.state.heaps.allocate();
+        let heap = vm.state.heaps.allocate_with_content(program.as_ref());
         vm.state.current_frame.heaps_i_am_keeping_alive.push(heap);
-        vm.state.heaps[heap].memset(program.as_ref());
 
         let value = FatPointer {
             offset: 0,

src/instruction_handlers/heap_access.rs

@@ -7,28 +7,26 @@ use crate::{
     fat_pointer::FatPointer,
     instruction::InstructionResult,
     state::State,
-    ExecutionEnd, Instruction, VirtualMachine, World,
+    ExecutionEnd, HeapId, Instruction, VirtualMachine, World,
 };
 use std::ops::Range;
 use u256::U256;
 
 pub trait HeapInterface {
     fn read_u256(&self, start_address: u32) -> U256;
     fn read_u256_partially(&self, range: Range<u32>) -> U256;
-    fn write_u256(&mut self, start_address: u32, value: U256);
     fn read_range_big_endian(&self, range: Range<u32>) -> Vec<u8>;
-    fn memset(&mut self, memory: &[u8]);
 }
 
 pub trait HeapFromState {
-    fn get_heap(state: &mut State) -> &mut impl HeapInterface;
+    fn get_heap(state: &State) -> HeapId;
     fn get_heap_size(state: &mut State) -> &mut u32;
 }
 
 pub struct Heap;
 impl HeapFromState for Heap {
-    fn get_heap(state: &mut State) -> &mut impl HeapInterface {
-        &mut state.heaps[state.current_frame.heap]
+    fn get_heap(state: &State) -> HeapId {
+        state.current_frame.heap
     }
     fn get_heap_size(state: &mut State) -> &mut u32 {
         &mut state.current_frame.heap_size
@@ -37,8 +35,8 @@ impl HeapFromState for Heap {
 
 pub struct AuxHeap;
 impl HeapFromState for AuxHeap {
-    fn get_heap(state: &mut State) -> &mut impl HeapInterface {
-        &mut state.heaps[state.current_frame.aux_heap]
+    fn get_heap(state: &State) -> HeapId {
+        state.current_frame.aux_heap
     }
     fn get_heap_size(state: &mut State) -> &mut u32 {
         &mut state.current_frame.aux_heap_size
@@ -72,7 +70,8 @@ fn load<H: HeapFromState, In: Source, const INCREMENT: bool>(
             return Ok(&PANIC);
         }
 
-        let value = H::get_heap(&mut vm.state).read_u256(address);
+        let heap = H::get_heap(&vm.state);
+        let value = vm.state.heaps[heap].read_u256(address);
         Register1::set(args, &mut vm.state, value);
 
         if INCREMENT {
@@ -108,7 +107,8 @@ fn store<H: HeapFromState, In: Source, const INCREMENT: bool, const HOOKING_ENAB
             return Ok(&PANIC);
         }
 
-        H::get_heap(&mut vm.state).write_u256(address, value);
+        let heap = H::get_heap(&vm.state);
+        vm.state.heaps.write_u256(heap, address, value);
 
         if INCREMENT {
             Register1::set(args, &mut vm.state, pointer + 32)

src/instruction_handlers/precompiles.rs

@@ -96,7 +96,7 @@ impl Memory for Heaps {
 
         let start = query.location.index.0 * 32;
         if query.rw_flag {
-            self[page].write_u256(start, query.value);
+            self.write_u256(page, start, query.value);
         } else {
             query.value = self[page].read_u256(start);
             query.value_is_pointer = false;