From 71dad714b3be6de343c172417ad23bf6f242fb99 Mon Sep 17 00:00:00 2001 From: Joonatan Saarhelo Date: Fri, 26 Apr 2024 18:36:13 +0200 Subject: [PATCH 1/3] don't preallocate any heap memory --- src/instruction_handlers/heap_access.rs | 17 +++++++++++++---- src/state.rs | 24 +++++------------------- 2 files changed, 18 insertions(+), 23 deletions(-) diff --git a/src/instruction_handlers/heap_access.rs b/src/instruction_handlers/heap_access.rs index 0b53e483..111203f5 100644 --- a/src/instruction_handlers/heap_access.rs +++ b/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_FRAME_MEMORY_STIPEND, NEW_KERNEL_FRAME_MEMORY_STIPEND}; pub trait HeapFromState { fn get_heap(state: &mut State) -> &mut Vec; @@ -114,10 +117,16 @@ fn store(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 stipend = if is_kernel(address_into_u256(state.current_frame.code_address)) { + NEW_KERNEL_FRAME_MEMORY_STIPEND + } else { + NEW_FRAME_MEMORY_STIPEND + }; + let heap_length = H::get_heap(state).len() as u32; + let already_paid = heap_length.max(stipend); + + 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(()) diff --git a/src/state.rs b/src/state.rs index 2112f5a9..e4225322 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,15 +1,7 @@ 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, Instruction, }; use std::ops::{Index, IndexMut}; use u256::{H160, U256}; @@ -118,14 +110,8 @@ impl State { world_before_this_frame: Snapshot, ) { 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![], vec![]]); let mut new_frame = Callframe::new( if CALLING_MODE == CallingMode::Delegate as u8 { From 8f20db7a9849e6241af80beb3babc009da4ad3d8 Mon Sep 17 00:00:00 2001 From: Joonatan Saarhelo Date: Fri, 26 Apr 2024 20:20:03 +0200 Subject: [PATCH 2/3] preallocate a little bit --- src/instruction_handlers/heap_access.rs | 11 +++++------ src/state.rs | 6 +++++- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/instruction_handlers/heap_access.rs b/src/instruction_handlers/heap_access.rs index 111203f5..421134f4 100644 --- a/src/instruction_handlers/heap_access.rs +++ b/src/instruction_handlers/heap_access.rs @@ -12,7 +12,7 @@ use crate::{ ExecutionEnd, Instruction, Predicate, VirtualMachine, }; use u256::U256; -use zkevm_opcode_defs::system_params::{NEW_FRAME_MEMORY_STIPEND, NEW_KERNEL_FRAME_MEMORY_STIPEND}; +use zkevm_opcode_defs::system_params::NEW_KERNEL_FRAME_MEMORY_STIPEND; pub trait HeapFromState { fn get_heap(state: &mut State) -> &mut Vec; @@ -117,13 +117,12 @@ fn store(state: &mut State, new_bound: u32) -> Result<(), ()> { - let stipend = if is_kernel(address_into_u256(state.current_frame.code_address)) { - NEW_KERNEL_FRAME_MEMORY_STIPEND + 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 { - NEW_FRAME_MEMORY_STIPEND + heap_length }; - let heap_length = H::get_heap(state).len() as u32; - let already_paid = heap_length.max(stipend); state.use_gas(new_bound.saturating_sub(already_paid))?; if heap_length < new_bound { diff --git a/src/state.rs b/src/state.rs index e4225322..478cf1ef 100644 --- a/src/state.rs +++ b/src/state.rs @@ -5,6 +5,7 @@ use crate::{ }; 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 { @@ -111,7 +112,10 @@ impl State { ) { let new_heap = self.heaps.0.len() as u32; - self.heaps.0.extend([vec![], vec![]]); + 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 { From 6810f99215bfdf80ebddcb90f1b086d9e9282f90 Mon Sep 17 00:00:00 2001 From: Joonatan Saarhelo Date: Fri, 26 Apr 2024 21:43:16 +0200 Subject: [PATCH 3/3] recycle stacks via a memory pool --- src/callframe.rs | 14 ++++------- src/instruction_handlers/far_call.rs | 5 +++- src/instruction_handlers/ret.rs | 3 ++- src/lib.rs | 1 + src/stack.rs | 35 ++++++++++++++++++++++++++++ src/state.rs | 31 +++++++++++++++++------- src/vm.rs | 9 +++++-- 7 files changed, 76 insertions(+), 22 deletions(-) create mode 100644 src/stack.rs diff --git a/src/callframe.rs b/src/callframe.rs index 7317be21..4fe84a7b 100644 --- a/src/callframe.rs +++ b/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, + pub stack: Box, pub heap: u32, pub aux_heap: u32, @@ -55,6 +54,7 @@ impl Callframe { code_address: H160, caller: H160, program: Program, + stack: Box, 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, diff --git a/src/instruction_handlers/far_call.rs b/src/instruction_handlers/far_call.rs index 93edb77e..cb552e67 100644 --- a/src/instruction_handlers/far_call.rs +++ b/src/instruction_handlers/far_call.rs @@ -56,9 +56,11 @@ fn far_call( 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( IS_STATIC && !is_evm_interpreter, calldata.memory_page, vm.world.snapshot(), + stack, ); vm.state.flags = Flags::new(false, false, false); diff --git a/src/instruction_handlers/ret.rs b/src/instruction_handlers/ret.rs index 9bcf9edd..b02cd3c1 100644 --- a/src/instruction_handlers/ret.rs +++ b/src/instruction_handlers/ret.rs @@ -76,7 +76,7 @@ fn ret( .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( Err(ExecutionEnd::Panicked) }; }; + vm.stack_pool.recycle(stack); vm.state.set_context_u128(0); vm.state.registers = [U256::zero(); 16]; diff --git a/src/lib.rs b/src/lib.rs index 540fe394..a52c718a 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/stack.rs b/src/stack.rs new file mode 100644 index 00000000..8de167bc --- /dev/null +++ b/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>, +} + +impl StackPool { + pub fn get(&mut self) -> Box { + 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) { + self.stacks.push(stack); + } +} diff --git a/src/state.rs b/src/state.rs index 478cf1ef..aed83556 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,7 +1,7 @@ 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, Instruction, + modified_world::Snapshot, predication::Flags, program::Program, stack::Stack, Instruction, }; use std::ops::{Index, IndexMut}; use u256::{H160, U256}; @@ -37,6 +37,7 @@ impl State { gas: u32, program: Program, world_before_this_frame: Snapshot, + stack: Box, ) -> Self { let mut registers: [U256; 16] = Default::default(); registers[1] = FatPointer { @@ -56,6 +57,7 @@ impl State { address, caller, program, + stack, FIRST_HEAP, 3, 1, @@ -109,6 +111,7 @@ impl State { is_static: bool, calldata_heap: u32, world_before_this_frame: Snapshot, + stack: Box, ) { let new_heap = self.heaps.0.len() as u32; @@ -133,6 +136,7 @@ impl State { u256_into_address(self.registers[15]) }, program, + stack, new_heap, new_heap + 1, calldata_heap, @@ -154,8 +158,11 @@ impl State { self.previous_frames.push((old_pc, new_frame)); } - pub(crate) fn pop_frame(&mut self, heap_to_keep: Option) -> Option<(u16, u16, Snapshot)> { - self.previous_frames.pop().map(|(pc, frame)| { + pub(crate) fn pop_frame( + &mut self, + heap_to_keep: Option, + ) -> Option<(u16, u16, Snapshot, Box)> { + 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) @@ -165,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) }) } @@ -184,12 +195,14 @@ impl State { instruction_pointer: *const Instruction, exception_handler: u16, world_before_this_frame: Snapshot, + stack: Box, ) { let mut new_frame = Callframe::new( H160::zero(), H160::zero(), H160::zero(), Program::new(vec![], vec![]), + stack, 0, 0, 0, @@ -223,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 diff --git a/src/vm.rs b/src/vm.rs index cbeab569..ef1cfe1e 100644 --- a/src/vm.rs +++ b/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, } }