diff --git a/src/callframe.rs b/src/callframe.rs index f8b14357..0abd8302 100644 --- a/src/callframe.rs +++ b/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,6 +159,46 @@ impl Callframe { .map(|f| f.previous_frame_gas) .sum::() } + + 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 { @@ -162,3 +206,19 @@ pub(crate) struct FrameRemnant { 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, + + heap_size: u32, + aux_heap_size: u32, + + heaps_i_am_keeping_alive: Vec, + world_before_this_frame: Snapshot, +} diff --git a/src/single_instruction_test/stack.rs b/src/single_instruction_test/stack.rs index 988bc4b5..802fc2eb 100644 --- a/src/single_instruction_test/stack.rs +++ b/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) {} } + +pub(crate) struct StackSnapshot; diff --git a/src/stack.rs b/src/stack.rs index dc8552b4..6b628278 100644 --- a/src/stack.rs +++ b/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 { diff --git a/src/state.rs b/src/state.rs index 3d8681d0..798c1248 100644 --- a/src/state.rs +++ b/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, +} diff --git a/src/vm.rs b/src/vm.rs index 8243d016..ad52026c 100644 --- a/src/vm.rs +++ b/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, }