refactor callstack pop result into struct

src/callframe.rs

@@ -104,15 +104,16 @@ impl Callframe {
         self.gas = gas_to_call;
     }
 
-    pub(crate) fn pop_near_call(&mut self) -> Option<(u16, u16, Snapshot)> {
+    pub(crate) fn pop_near_call(&mut self) -> Option<FrameRemnant> {
         self.near_calls.pop().map(|f| {
             self.sp = f.previous_frame_sp;
             self.gas = f.previous_frame_gas;
-            (
-                f.call_instruction,
-                f.exception_handler,
-                f.world_before_this_frame,
-            )
+
+            FrameRemnant {
+                program_counter: f.call_instruction,
+                exception_handler: f.exception_handler,
+                snapshot: f.world_before_this_frame,
+            }
         })
     }
 
@@ -137,3 +138,9 @@ impl Callframe {
                 .sum::<u32>()
     }
 }
+
+pub(crate) struct FrameRemnant {
+    pub(crate) program_counter: u16,
+    pub(crate) exception_handler: u16,
+    pub(crate) snapshot: Snapshot,
+}

src/instruction_handlers/far_call.rs

@@ -69,7 +69,7 @@ fn far_call<const CALLING_MODE: u8, const IS_STATIC: bool>(
         .checked_add(stipend)
         .expect("stipend must not cause overflow");
 
-    vm.state.push_frame::<CALLING_MODE>(
+    vm.push_frame::<CALLING_MODE>(
         instruction,
         u256_into_address(destination_address),
         program,
@@ -79,7 +79,6 @@ fn far_call<const CALLING_MODE: u8, const IS_STATIC: bool>(
         IS_STATIC && !is_evm_interpreter,
         calldata.memory_page,
         vm.world.snapshot(),
-        vm.stack_pool.get(),
     );
 
     vm.state.flags = Flags::new(false, false, false);

src/instruction_handlers/ret.rs

@@ -1,6 +1,7 @@
 use super::far_call::get_far_call_calldata;
 use crate::{
     addressing_modes::{Arguments, Immediate1, Register1, Source, INVALID_INSTRUCTION_COST},
+    callframe::FrameRemnant,
     instruction::{ExecutionEnd, InstructionResult},
     predication::Flags,
     Instruction, Predicate, VirtualMachine,
@@ -39,16 +40,19 @@ fn ret<const RETURN_TYPE: u8, const TO_LABEL: bool>(
     let mut return_type = ReturnType::from_u8(RETURN_TYPE);
     let near_call_leftover_gas = vm.state.current_frame.gas;
 
-    let (pc, snapshot, leftover_gas) = if let Some((pc, eh, snapshot)) =
-        vm.state.current_frame.pop_near_call()
+    let (pc, snapshot, leftover_gas) = if let Some(FrameRemnant {
+        program_counter,
+        exception_handler,
+        snapshot,
+    }) = vm.state.current_frame.pop_near_call()
     {
         (
             if TO_LABEL {
                 Immediate1::get(args, &mut vm.state).low_u32() as u16
             } else if return_type.is_failure() {
-                eh
+                exception_handler
             } else {
-                pc.wrapping_add(1)
+                program_counter.wrapping_add(1)
             },
             snapshot,
             near_call_leftover_gas,
@@ -76,11 +80,16 @@ fn ret<const RETURN_TYPE: u8, const TO_LABEL: bool>(
             .gas
             .saturating_sub(vm.state.current_frame.stipend);
 
-        let Some((pc, eh, snapshot, stack)) = vm.state.pop_frame(
+        let Some(FrameRemnant {
+            program_counter,
+            exception_handler,
+            snapshot,
+        }) = vm.pop_frame(
             return_value_or_panic
                 .as_ref()
                 .map(|pointer| pointer.memory_page),
-        ) else {
+        )
+        else {
             if return_type.is_failure() {
                 vm.world
                     .rollback(vm.state.current_frame.world_before_this_frame);
@@ -99,7 +108,6 @@ fn ret<const RETURN_TYPE: u8, const TO_LABEL: bool>(
                 Err(ExecutionEnd::Panicked)
             };
         };
-        vm.stack_pool.recycle(stack);
 
         vm.state.set_context_u128(0);
         vm.state.registers = [U256::zero(); 16];
@@ -111,9 +119,9 @@ fn ret<const RETURN_TYPE: u8, const TO_LABEL: bool>(
 
         (
             if return_type.is_failure() {
-                eh
+                exception_handler
             } else {
-                pc.wrapping_add(1)
+                program_counter.wrapping_add(1)
             },
             snapshot,
             leftover_gas,

src/state.rs

@@ -1,11 +1,9 @@
 use crate::{
-    addressing_modes::Addressable, bitset::Bitset, callframe::Callframe,
-    decommit::u256_into_address, fat_pointer::FatPointer, instruction_handlers::CallingMode,
-    modified_world::Snapshot, predication::Flags, program::Program, stack::Stack, Instruction,
+    addressing_modes::Addressable, bitset::Bitset, callframe::Callframe, fat_pointer::FatPointer,
+    modified_world::Snapshot, predication::Flags, program::Program, stack::Stack,
 };
 use std::ops::{Index, IndexMut};
 use u256::{H160, U256};
-use zkevm_opcode_defs::system_params::NEW_FRAME_MEMORY_STIPEND;
 
 #[derive(Clone, PartialEq, Debug)]
 pub struct State {
@@ -24,7 +22,7 @@ pub struct State {
 
     pub transaction_number: u16,
 
-    context_u128: u128,
+    pub(crate) context_u128: u128,
 }
 
 pub const FIRST_HEAP: u32 = 2;
@@ -100,95 +98,6 @@ impl State {
                 .sum::<u32>()
     }
 
-    #[allow(clippy::too_many_arguments)]
-    pub(crate) fn push_frame<const CALLING_MODE: u8>(
-        &mut self,
-        instruction_pointer: *const Instruction,
-        code_address: H160,
-        program: Program,
-        gas: u32,
-        stipend: u32,
-        exception_handler: u16,
-        is_static: bool,
-        calldata_heap: u32,
-        world_before_this_frame: Snapshot,
-        stack: Box<Stack>,
-    ) {
-        let new_heap = self.heaps.0.len() as u32;
-
-        self.heaps.0.extend([
-            vec![0; NEW_FRAME_MEMORY_STIPEND as usize],
-            vec![0; NEW_FRAME_MEMORY_STIPEND as usize],
-        ]);
-
-        let mut new_frame = Callframe::new(
-            if CALLING_MODE == CallingMode::Delegate as u8 {
-                self.current_frame.address
-            } else {
-                code_address
-            },
-            code_address,
-            if CALLING_MODE == CallingMode::Normal as u8 {
-                self.current_frame.address
-            } else if CALLING_MODE == CallingMode::Delegate as u8 {
-                self.current_frame.caller
-            } else {
-                // Mimic call
-                u256_into_address(self.registers[15])
-            },
-            program,
-            stack,
-            new_heap,
-            new_heap + 1,
-            calldata_heap,
-            gas,
-            stipend,
-            exception_handler,
-            if CALLING_MODE == CallingMode::Delegate as u8 {
-                self.current_frame.context_u128
-            } else {
-                self.context_u128
-            },
-            is_static || self.current_frame.is_static,
-            world_before_this_frame,
-        );
-        self.context_u128 = 0;
-
-        let old_pc = self.current_frame.pc_to_u16(instruction_pointer);
-        std::mem::swap(&mut new_frame, &mut self.current_frame);
-        self.previous_frames.push((old_pc, new_frame));
-    }
-
-    pub(crate) fn pop_frame(
-        &mut self,
-        heap_to_keep: Option<u32>,
-    ) -> Option<(u16, u16, Snapshot, Box<Stack>)> {
-        self.previous_frames.pop().map(|(pc, mut frame)| {
-            for &heap in [self.current_frame.heap, self.current_frame.aux_heap]
-                .iter()
-                .chain(&self.current_frame.heaps_i_am_keeping_alive)
-            {
-                if Some(heap) != heap_to_keep {
-                    self.heaps.deallocate(heap);
-                }
-            }
-
-            std::mem::swap(&mut self.current_frame, &mut frame);
-            let Callframe {
-                exception_handler,
-                world_before_this_frame,
-                stack,
-                ..
-            } = frame;
-
-            self.current_frame
-                .heaps_i_am_keeping_alive
-                .extend(heap_to_keep);
-
-            (pc, exception_handler, world_before_this_frame, stack)
-        })
-    }
-
     pub(crate) fn set_context_u128(&mut self, value: u128) {
         self.context_u128 = value;
     }
@@ -224,7 +133,7 @@ impl Addressable for State {
 }
 
 #[derive(Debug, Clone)]
-pub struct Heaps(Vec<Vec<u8>>);
+pub struct Heaps(pub(crate) Vec<Vec<u8>>);
 
 impl Heaps {
     pub(crate) fn deallocate(&mut self, heap: u32) {

src/vm.rs

@@ -1,8 +1,14 @@
 use crate::{
-    instruction_handlers::free_panic, modified_world::ModifiedWorld, stack::StackPool,
-    state::State, ExecutionEnd, Instruction, Program, World,
+    callframe::{Callframe, FrameRemnant},
+    decommit::u256_into_address,
+    instruction_handlers::{free_panic, CallingMode},
+    modified_world::{ModifiedWorld, Snapshot},
+    stack::StackPool,
+    state::State,
+    ExecutionEnd, Instruction, Program, World,
 };
 use u256::H160;
+use zkevm_opcode_defs::system_params::NEW_FRAME_MEMORY_STIPEND;
 
 pub struct Settings {
     pub default_aa_code_hash: [u8; 32],
@@ -139,6 +145,104 @@ impl VirtualMachine {
             .map(|left| (left, end))
     }
 
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) fn push_frame<const CALLING_MODE: u8>(
+        &mut self,
+        instruction_pointer: *const Instruction,
+        code_address: H160,
+        program: Program,
+        gas: u32,
+        stipend: u32,
+        exception_handler: u16,
+        is_static: bool,
+        calldata_heap: u32,
+        world_before_this_frame: Snapshot,
+    ) {
+        let new_heap = self.state.heaps.0.len() as u32;
+
+        self.state.heaps.0.extend([
+            vec![0; NEW_FRAME_MEMORY_STIPEND as usize],
+            vec![0; NEW_FRAME_MEMORY_STIPEND as usize],
+        ]);
+
+        let mut new_frame = Callframe::new(
+            if CALLING_MODE == CallingMode::Delegate as u8 {
+                self.state.current_frame.address
+            } else {
+                code_address
+            },
+            code_address,
+            if CALLING_MODE == CallingMode::Normal as u8 {
+                self.state.current_frame.address
+            } else if CALLING_MODE == CallingMode::Delegate as u8 {
+                self.state.current_frame.caller
+            } else {
+                // Mimic call
+                u256_into_address(self.state.registers[15])
+            },
+            program,
+            self.stack_pool.get(),
+            new_heap,
+            new_heap + 1,
+            calldata_heap,
+            gas,
+            stipend,
+            exception_handler,
+            if CALLING_MODE == CallingMode::Delegate as u8 {
+                self.state.current_frame.context_u128
+            } else {
+                self.state.context_u128
+            },
+            is_static || self.state.current_frame.is_static,
+            world_before_this_frame,
+        );
+        self.state.context_u128 = 0;
+
+        let old_pc = self.state.current_frame.pc_to_u16(instruction_pointer);
+        std::mem::swap(&mut new_frame, &mut self.state.current_frame);
+        self.state.previous_frames.push((old_pc, new_frame));
+    }
+
+    pub(crate) fn pop_frame(&mut self, heap_to_keep: Option<u32>) -> Option<FrameRemnant> {
+        self.state
+            .previous_frames
+            .pop()
+            .map(|(program_counter, mut frame)| {
+                for &heap in [
+                    self.state.current_frame.heap,
+                    self.state.current_frame.aux_heap,
+                ]
+                .iter()
+                .chain(&self.state.current_frame.heaps_i_am_keeping_alive)
+                {
+                    if Some(heap) != heap_to_keep {
+                        self.state.heaps.deallocate(heap);
+                    }
+                }
+
+                std::mem::swap(&mut self.state.current_frame, &mut frame);
+                let Callframe {
+                    exception_handler,
+                    world_before_this_frame,
+                    stack,
+                    ..
+                } = frame;
+
+                self.stack_pool.recycle(stack);
+
+                self.state
+                    .current_frame
+                    .heaps_i_am_keeping_alive
+                    .extend(heap_to_keep);
+
+                FrameRemnant {
+                    program_counter,
+                    exception_handler,
+                    snapshot: world_before_this_frame,
+                }
+            })
+    }
+
     #[cfg(trace)]
     fn print_instruction(&self, instruction: *const Instruction) {
         print!("{:?}: ", unsafe {