perf: Make far call faster by making it allocate less

src/callframe.rs

@@ -1,5 +1,5 @@
-use crate::{bitset::Bitset, modified_world::Snapshot, program::Program, Instruction};
-use u256::{H160, U256};
+use crate::{modified_world::Snapshot, program::Program, stack::Stack, Instruction};
+use u256::H160;
 
 #[derive(Clone, PartialEq, Debug)]
 pub struct Callframe {
@@ -12,8 +12,7 @@ pub struct Callframe {
     pub is_static: bool,
 
     // TODO: joint allocate these.
-    pub stack: Box<[U256; 1 << 16]>,
-    pub stack_pointer_flags: Box<Bitset>,
+    pub stack: Box<Stack>,
 
     pub heap: u32,
     pub aux_heap: u32,
@@ -55,6 +54,7 @@ impl Callframe {
         code_address: H160,
         caller: H160,
         program: Program,
+        stack: Box<Stack>,
         heap: u32,
         aux_heap: u32,
         calldata_heap: u32,
@@ -72,11 +72,7 @@ impl Callframe {
             program,
             context_u128,
             is_static,
-            stack: vec![U256::zero(); 1 << 16]
-                .into_boxed_slice()
-                .try_into()
-                .unwrap(),
-            stack_pointer_flags: Default::default(),
+            stack,
             heap,
             aux_heap,
             calldata_heap,

src/instruction_handlers/far_call.rs

@@ -56,9 +56,11 @@ fn far_call<const CALLING_MODE: u8, const IS_STATIC: bool>(
     let new_frame_gas = abi.gas_to_pass.min(maximum_gas);
     vm.state.current_frame.gas -= new_frame_gas;
 
+    let stack = vm.stack_pool.get();
+
     let (Some(calldata), Some((program, is_evm_interpreter))) = (calldata, decommit_result) else {
         vm.state
-            .push_dummy_frame(instruction, exception_handler, vm.world.snapshot());
+            .push_dummy_frame(instruction, exception_handler, vm.world.snapshot(), stack);
         return Ok(&INVALID_INSTRUCTION);
     };
 
@@ -81,6 +83,7 @@ fn far_call<const CALLING_MODE: u8, const IS_STATIC: bool>(
         IS_STATIC && !is_evm_interpreter,
         calldata.memory_page,
         vm.world.snapshot(),
+        stack,
     );
 
     vm.state.flags = Flags::new(false, false, false);

src/instruction_handlers/heap_access.rs

@@ -1,15 +1,18 @@
 use super::{common::instruction_boilerplate_with_panic, PANIC};
 use crate::{
+    address_into_u256,
     addressing_modes::{
         Arguments, Destination, DestinationWriter, Immediate1, Register1, Register2,
         RegisterOrImmediate, Source,
     },
+    decommit::is_kernel,
     fat_pointer::FatPointer,
     instruction::InstructionResult,
     state::State,
     ExecutionEnd, Instruction, Predicate, VirtualMachine,
 };
 use u256::U256;
+use zkevm_opcode_defs::system_params::NEW_KERNEL_FRAME_MEMORY_STIPEND;
 
 pub trait HeapFromState {
     fn get_heap(state: &mut State) -> &mut Vec<u8>;
@@ -114,10 +117,15 @@ fn store<H: HeapFromState, In: Source, const INCREMENT: bool, const HOOKING_ENAB
 }
 
 pub fn grow_heap<H: HeapFromState>(state: &mut State, new_bound: u32) -> Result<(), ()> {
-    if let Some(growth) = new_bound.checked_sub(H::get_heap(state).len() as u32) {
-        state.use_gas(growth)?;
-
-        // This will not cause frequent reallocations; it allocates in a geometric series like push.
+    let heap_length = H::get_heap(state).len() as u32;
+    let already_paid = if is_kernel(address_into_u256(state.current_frame.code_address)) {
+        heap_length.max(NEW_KERNEL_FRAME_MEMORY_STIPEND)
+    } else {
+        heap_length
+    };
+
+    state.use_gas(new_bound.saturating_sub(already_paid))?;
+    if heap_length < new_bound {
         H::get_heap(state).resize(new_bound as usize, 0);
     }
     Ok(())

src/instruction_handlers/ret.rs

@@ -76,7 +76,7 @@ fn ret<const RETURN_TYPE: u8, const TO_LABEL: bool>(
             .gas
             .saturating_sub(vm.state.current_frame.stipend);
 
-        let Some((pc, eh, snapshot)) = vm.state.pop_frame(
+        let Some((pc, eh, snapshot, stack)) = vm.state.pop_frame(
             return_value_or_panic
                 .as_ref()
                 .map(|pointer| pointer.memory_page),
@@ -99,6 +99,7 @@ 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];

src/lib.rs

@@ -12,6 +12,7 @@ mod modified_world;
 mod predication;
 mod program;
 mod rollback;
+mod stack;
 mod state;
 pub mod testworld;
 mod vm;

src/stack.rs

@@ -0,0 +1,35 @@
+use crate::bitset::Bitset;
+use u256::U256;
+
+#[derive(Clone, PartialEq, Debug)]
+pub struct Stack {
+    pub pointer_flags: Bitset,
+    pub slots: [U256; 1 << 16],
+}
+
+#[derive(Default)]
+pub struct StackPool {
+    stacks: Vec<Box<Stack>>,
+}
+
+impl StackPool {
+    pub fn get(&mut self) -> Box<Stack> {
+        self.stacks
+            .pop()
+            .map(|mut s| {
+                s.slots = [U256::zero(); 1 << 16];
+                s.pointer_flags = Default::default();
+                s
+            })
+            .unwrap_or_else(|| {
+                Box::new(Stack {
+                    pointer_flags: Default::default(),
+                    slots: [U256::zero(); 1 << 16],
+                })
+            })
+    }
+
+    pub fn recycle(&mut self, stack: Box<Stack>) {
+        self.stacks.push(stack);
+    }
+}

src/state.rs

@@ -1,18 +1,11 @@
 use crate::{
-    address_into_u256,
-    addressing_modes::Addressable,
-    bitset::Bitset,
-    callframe::Callframe,
-    decommit::{is_kernel, u256_into_address},
-    fat_pointer::FatPointer,
-    instruction_handlers::CallingMode,
-    modified_world::Snapshot,
-    predication::Flags,
-    program::Program,
-    Instruction,
+    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,
 };
 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 {
@@ -44,6 +37,7 @@ impl State {
         gas: u32,
         program: Program,
         world_before_this_frame: Snapshot,
+        stack: Box<Stack>,
     ) -> Self {
         let mut registers: [U256; 16] = Default::default();
         registers[1] = FatPointer {
@@ -63,6 +57,7 @@ impl State {
                 address,
                 caller,
                 program,
+                stack,
                 FIRST_HEAP,
                 3,
                 1,
@@ -116,16 +111,14 @@ impl State {
         is_static: bool,
         calldata_heap: u32,
         world_before_this_frame: Snapshot,
+        stack: Box<Stack>,
     ) {
         let new_heap = self.heaps.0.len() as u32;
-        let new_heap_len = if is_kernel(address_into_u256(code_address)) {
-            zkevm_opcode_defs::system_params::NEW_KERNEL_FRAME_MEMORY_STIPEND
-        } else {
-            zkevm_opcode_defs::system_params::NEW_FRAME_MEMORY_STIPEND
-        } as usize;
-        self.heaps
-            .0
-            .extend([vec![0; new_heap_len], vec![0; new_heap_len]]);
+
+        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 {
@@ -143,6 +136,7 @@ impl State {
                 u256_into_address(self.registers[15])
             },
             program,
+            stack,
             new_heap,
             new_heap + 1,
             calldata_heap,
@@ -164,8 +158,11 @@ impl State {
         self.previous_frames.push((old_pc, new_frame));
     }
 
-    pub(crate) fn pop_frame(&mut self, heap_to_keep: Option<u32>) -> Option<(u16, u16, Snapshot)> {
-        self.previous_frames.pop().map(|(pc, 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)
@@ -175,15 +172,19 @@ impl State {
                 }
             }
 
-            let eh = self.current_frame.exception_handler;
-            let snapshot = self.current_frame.world_before_this_frame;
+            std::mem::swap(&mut self.current_frame, &mut frame);
+            let Callframe {
+                exception_handler,
+                world_before_this_frame,
+                stack,
+                ..
+            } = frame;
 
-            self.current_frame = frame;
             self.current_frame
                 .heaps_i_am_keeping_alive
                 .extend(heap_to_keep);
 
-            (pc, eh, snapshot)
+            (pc, exception_handler, world_before_this_frame, stack)
         })
     }
 
@@ -194,12 +195,14 @@ impl State {
         instruction_pointer: *const Instruction,
         exception_handler: u16,
         world_before_this_frame: Snapshot,
+        stack: Box<Stack>,
     ) {
         let mut new_frame = Callframe::new(
             H160::zero(),
             H160::zero(),
             H160::zero(),
             Program::new(vec![], vec![]),
+            stack,
             0,
             0,
             0,
@@ -233,10 +236,10 @@ impl Addressable for State {
         &mut self.register_pointer_flags
     }
     fn stack(&mut self) -> &mut [U256; 1 << 16] {
-        &mut self.current_frame.stack
+        &mut self.current_frame.stack.slots
     }
     fn stack_pointer_flags(&mut self) -> &mut Bitset {
-        &mut self.current_frame.stack_pointer_flags
+        &mut self.current_frame.stack.pointer_flags
     }
     fn stack_pointer(&mut self) -> &mut u16 {
         &mut self.current_frame.sp

src/vm.rs

@@ -1,6 +1,6 @@
 use crate::{
-    instruction_handlers::free_panic, modified_world::ModifiedWorld, state::State, ExecutionEnd,
-    Instruction, Program, World,
+    instruction_handlers::free_panic, modified_world::ModifiedWorld, stack::StackPool,
+    state::State, ExecutionEnd, Instruction, Program, World,
 };
 use u256::H160;
 
@@ -20,6 +20,8 @@ pub struct VirtualMachine {
     pub state: State,
 
     pub(crate) settings: Settings,
+
+    pub(crate) stack_pool: StackPool,
 }
 
 impl VirtualMachine {
@@ -34,6 +36,7 @@ impl VirtualMachine {
     ) -> Self {
         let world = ModifiedWorld::new(world);
         let world_before_this_frame = world.snapshot();
+        let mut stack_pool = StackPool::default();
 
         Self {
             world,
@@ -44,8 +47,10 @@ impl VirtualMachine {
                 gas,
                 program,
                 world_before_this_frame,
+                stack_pool.get(),
             ),
             settings,
+            stack_pool,
         }
     }