chore(SharedMemory): refactoring

Author: Lorenzo Feroleto

crates/interpreter/src/instructions/macros.rs

@@ -52,12 +52,10 @@ macro_rules! gas_or_fail {
 
 macro_rules! shared_memory_resize {
     ($interp:expr, $offset:expr, $len:expr) => {
-        let len: usize = $len;
-        let offset: usize = $offset;
         if let Some(new_size) =
-            crate::interpreter::shared_memory::next_multiple_of_32(offset.saturating_add(len))
+            crate::interpreter::shared_memory::next_multiple_of_32($offset.saturating_add($len))
         {
-            if new_size as u64 > $interp.shared_memory.limit {
+            if new_size > $interp.shared_memory.limit {
                 $interp.instruction_result = InstructionResult::MemoryLimitOOG;
                 return;
             }

crates/interpreter/src/interpreter.rs

@@ -147,9 +147,7 @@ impl<'a> Interpreter<'a> {
     /// Returns a copy of the interpreter's return value, if any.
     #[inline]
     pub fn return_value(&mut self) -> Bytes {
-        let return_data = self.return_value_slice().to_vec().into();
-        self.shared_memory.free_context_memory();
-        return_data
+        self.return_value_slice().to_vec().into()
     }
 
     /// Returns a reference to the interpreter's return value, if any.

crates/interpreter/src/interpreter/shared_memory.rs

@@ -1,7 +1,7 @@
 use revm_primitives::U256;
 
 use crate::alloc::vec;
-use crate::alloc::vec::Vec;
+use crate::alloc::{boxed::Box, slice, vec::Vec};
 use core::{
     cmp::min,
     fmt,
@@ -14,119 +14,111 @@ use core::{
 /// the `new` static method to ensure memory safety.
 pub struct SharedMemory {
     /// Shared buffer
-    data: Vec<u8>,
+    data: Box<[u8]>,
     /// Memory checkpoints for each depth
     checkpoints: Vec<usize>,
     /// Raw pointer used for the portion of memory used
     /// by the current context
-    current_slice: *mut [u8],
+    current_ptr: *mut u8,
     /// How much memory has been used in the current context
     current_len: usize,
     /// Amount of memory left for assignment
-    pub limit: u64,
+    pub limit: usize,
 }
 
 impl fmt::Debug for SharedMemory {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("SharedMemory")
             .field(
                 "current_slice",
-                &crate::primitives::hex::encode(self.get_current_slice()),
+                &crate::primitives::hex::encode(self.current_slice()),
             )
             .finish()
     }
 }
 
 impl SharedMemory {
+    /// Calculates memory allocation upper bound using
+    /// https://2π.com/22/eth-max-mem
+    #[inline]
+    pub fn calculate_upper_bound(gas_limit: u64) -> u64 {
+        4096 * sqrt(2 * gas_limit)
+    }
+
     /// Allocate memory to be shared between calls.
     /// Memory size is estimated using https://2π.com/22/eth-max-mem
     /// which depends on transaction [gas_limit].
     /// Maximum allocation size is 2^32 - 1 bytes;
     pub fn new(gas_limit: u64) -> Self {
-        let size = min(
-            SharedMemory::calculate_upper_bound(gas_limit) as usize,
-            2usize.pow(32) - 1,
-        );
-
-        let mut data = vec![0; size];
-        let current_slice: *mut [u8] = &mut data[..];
-        SharedMemory {
-            data,
-            checkpoints: Vec::with_capacity(1024),
-            current_slice,
-            current_len: 0,
-            limit: size as u64,
-        }
+        Self::new_with_memory_limit(gas_limit, (u32::MAX - 1) as u64)
     }
 
     /// Allocate memory to be shared between calls.
     /// Memory size is estimated using https://2π.com/22/eth-max-mem
     /// which depends on transaction [gas_limit].
     /// Uses [memory_limit] as maximum allocation size
     pub fn new_with_memory_limit(gas_limit: u64, memory_limit: u64) -> Self {
-        let size = min(
-            SharedMemory::calculate_upper_bound(gas_limit) as usize,
+        let limit = min(
+            Self::calculate_upper_bound(gas_limit) as usize,
             memory_limit as usize,
         );
 
-        let mut data = vec![0; size];
-        let current_slice: *mut [u8] = &mut data[..];
-        SharedMemory {
+        let mut data = vec![0; limit].into_boxed_slice();
+        let current_slice = data.as_mut_ptr();
+        Self {
             data,
-            checkpoints: Vec::with_capacity(1024),
-            current_slice,
+            checkpoints: Vec::with_capacity(32),
+            current_ptr: current_slice,
             current_len: 0,
-            limit: size as u64,
+            limit,
         }
     }
 
     /// Prepares the shared memory for a new context
     pub fn new_context_memory(&mut self) {
-        let base_offset = self.checkpoints.last().unwrap_or(&0);
+        let base_offset = self.last_checkpoint();
         let new_checkpoint = base_offset + self.current_len;
 
         self.checkpoints.push(new_checkpoint);
 
-        self.current_slice = &mut self.data[new_checkpoint..];
+        self.current_ptr = self.data[new_checkpoint..].as_mut_ptr();
         self.current_len = 0;
     }
 
     /// Prepares the shared memory for returning to the previous context
     pub fn free_context_memory(&mut self) {
         if let Some(old_checkpoint) = self.checkpoints.pop() {
-            let last = *self.checkpoints.last().unwrap_or(&0);
-            self.current_slice = &mut self.data[last..];
-            self.current_len = old_checkpoint - last;
+            let last_checkpoint = self.last_checkpoint();
+            self.current_ptr = self.data[last_checkpoint..].as_mut_ptr();
+            self.current_len = old_checkpoint - last_checkpoint;
             self.update_limit();
         }
     }
 
     /// Get the length of the current memory range.
+    #[inline(always)]
     pub fn len(&self) -> usize {
         self.current_len
     }
 
-    /// Return true if current effective memory range is zero.
+    #[inline(always)]
     pub fn is_empty(&self) -> bool {
         self.current_len == 0
     }
 
-    /// Return the memory of the current context
-    pub fn data(&self) -> &[u8] {
-        self.get_current_slice()
-    }
-
     /// Resize the memory. assume that we already checked if:
-    /// - we have enought gas to resize this vector
+    /// - we have enough gas to resize this vector
     /// - we made new_size as multiply of 32
     /// - [new_size] is greater than `self.len()`
+    #[inline(always)]
     pub fn resize(&mut self, new_size: usize) {
         // extend with zeros
         let range = self.current_len..new_size;
 
-        self.get_current_slice_mut()[range]
-            .iter_mut()
-            .for_each(|byte| *byte = 0);
+        for byte in self.current_slice_mut()[range].iter_mut() {
+            *byte = 0;
+        }
+
         self.current_len = new_size;
         self.update_limit();
     }
@@ -137,7 +129,7 @@ impl SharedMemory {
     #[inline(always)]
     #[cfg_attr(debug_assertions, track_caller)]
     pub fn slice(&self, offset: usize, size: usize) -> &[u8] {
-        match self.get_current_slice().get(offset..offset + size) {
+        match self.current_slice().get(offset..offset + size) {
             Some(slice) => slice,
             None => debug_unreachable!("slice OOB: {offset}..{size}; len: {}", self.len()),
         }
@@ -150,7 +142,7 @@ impl SharedMemory {
     #[cfg_attr(debug_assertions, track_caller)]
     pub fn slice_mut(&mut self, offset: usize, size: usize) -> &mut [u8] {
         let len = self.current_len;
-        match self.get_current_slice_mut().get_mut(offset..offset + size) {
+        match self.current_slice_mut().get_mut(offset..offset + size) {
             Some(slice) => slice,
             None => debug_unreachable!("slice OOB: {offset}..{size}; len: {}", len),
         }
@@ -162,7 +154,7 @@ impl SharedMemory {
     #[inline(always)]
     #[cfg_attr(debug_assertions, track_caller)]
     pub fn set_byte(&mut self, index: usize, byte: u8) {
-        match self.get_current_slice_mut().get_mut(index) {
+        match self.current_slice_mut().get_mut(index) {
             Some(b) => *b = byte,
             None => debug_unreachable!("set_byte OOB: {index}; len: {}", self.len()),
         }
@@ -219,36 +211,33 @@ impl SharedMemory {
     #[inline(always)]
     #[cfg_attr(debug_assertions, track_caller)]
     pub fn copy(&mut self, dst: usize, src: usize, len: usize) {
-        self.get_current_slice_mut()
-            .copy_within(src..src + len, dst);
+        self.current_slice_mut().copy_within(src..src + len, dst);
     }
 
-    /// Get a refernce to the memory of the current context
+    /// Get a reference to the memory of the current context
     #[inline(always)]
-    fn get_current_slice(&self) -> &[u8] {
+    fn current_slice(&self) -> &[u8] {
         // Safety: if it is a valid pointer to a slice of `self.data`
-        unsafe { &*self.current_slice }
+        unsafe { slice::from_raw_parts(self.current_ptr, self.limit) }
     }
 
-    /// Get a mutable refernce to the memory of the current context
+    /// Get a mutable reference to the memory of the current context
     #[inline(always)]
-    fn get_current_slice_mut(&mut self) -> &mut [u8] {
+    fn current_slice_mut(&mut self) -> &mut [u8] {
         // Safety: it is a valid pointer to a slice of `self.data`
-        unsafe { &mut *self.current_slice }
+        unsafe { slice::from_raw_parts_mut(self.current_ptr, self.limit) }
     }
 
-    /// Calculates memory allocation upper bound using
-    /// https://2π.com/22/eth-max-mem
+    /// Update the amount of memory left for usage
     #[inline(always)]
-    fn calculate_upper_bound(gas_limit: u64) -> u64 {
-        4096 * sqrt(2 * gas_limit)
+    fn update_limit(&mut self) {
+        self.limit = self.data.len() - self.last_checkpoint() - self.current_len;
     }
 
-    /// Update the amount of memory left for usage
+    /// Get the last memory checkpoint
     #[inline(always)]
-    fn update_limit(&mut self) {
-        self.limit =
-            (self.data.len() - *self.checkpoints.last().unwrap_or(&0) - self.current_len) as u64;
+    fn last_checkpoint(&self) -> usize {
+        *self.checkpoints.last().unwrap_or(&0)
     }
 }
 

crates/revm/src/evm_impl.rs

@@ -604,7 +604,11 @@ impl<'a, GSPEC: Spec, DB: Database, const INSPECT: bool> EVMImpl<'a, GSPEC, DB,
             interpreter.run::<Self, GSPEC>(self)
         };
 
-        (exit_reason, interpreter.return_value(), *interpreter.gas())
+        let (return_value, gas) = (interpreter.return_value(), *interpreter.gas());
+
+        interpreter.shared_memory.free_context_memory();
+
+        (exit_reason, return_value, gas)
     }
 
     /// Call precompile contract