ethrex_levm/

memory.rs

use std::{cell::RefCell, rc::Rc};

use crate::{
    constants::{MEMORY_EXPANSION_QUOTIENT, WORD_SIZE_IN_BYTES_U64, WORD_SIZE_IN_BYTES_USIZE},
    errors::{ExceptionalHalt, InternalError, VMError},
};
use ExceptionalHalt::OutOfBounds;
use bytes::Bytes;
use ethrex_common::{
    U256,
    utils::{u256_from_big_endian_const, u256_to_big_endian},
};

/// A cheaply clonable callframe-shared memory buffer.
///
/// When a new callframe is created a RC clone of this memory is made, with the current base offset at the length of the buffer at that time.
#[derive(Debug, Clone)]
pub struct Memory {
    pub buffer: Rc<RefCell<Vec<u8>>>,
    pub len: usize,
    current_base: usize,
}

impl Memory {
    #[inline]
    pub fn new() -> Self {
        Self {
            buffer: Rc::new(RefCell::new(Vec::new())),
            len: 0,
            current_base: 0,
        }
    }

    /// Resets this memory so its buffer can be reused from a pool by the next transaction:
    /// drops all contents (length → 0, capacity retained) and rebases to 0.
    ///
    /// Truncating the buffer to length 0 is REQUIRED for correctness, not just hygiene:
    /// [`Memory::resize`] only zero-fills bytes grown *past* `buffer.len()`, so handing a
    /// non-empty buffer to the next tx would expose stale data from the previous one (a
    /// consensus bug). Capacity is kept so the grown allocation is reused.
    #[inline]
    pub fn reset_for_reuse(&mut self) {
        self.buf_mut().clear();
        self.len = 0;
        self.current_base = 0;
    }

    /// Gets the Memory for the next children callframe.
    #[inline]
    pub fn next_memory(&self) -> Memory {
        let mut mem = self.clone();
        mem.current_base = mem.buffer.borrow().len();
        mem.len = 0;
        mem
    }

    /// Cleans the memory from base onwards, this must be used in callframes when handling returns.
    ///
    /// On the callframe that is about to be dropped.
    #[inline]
    pub fn clean_from_base(&self) {
        #[expect(unsafe_code)]
        unsafe {
            self.buffer
                .borrow_mut()
                .get_unchecked_mut(self.current_base..(self.current_base.wrapping_add(self.len)))
                .fill(0);
        }
    }

    /// Truncates the memory back to base. This is crucial for constrained
    /// memory in zkVMs. The memory is not freed, but rather shrunk in `len`,
    /// so that the already allocated `capacity` is reused.
    #[cfg(target_arch = "riscv64")]
    #[inline]
    pub fn truncate_to_base(&self) {
        self.buf_mut().truncate(self.current_base);
    }

    /// Unchecked shared access to the backing buffer, skipping `RefCell`'s runtime
    /// borrow-flag read/branch/inc/dec. Sound because the guest is single-threaded
    /// and only the current call frame ever mutates its own memory region, so
    /// these borrows never overlap or alias — the same invariant `clean_from_base`
    /// and the `get_unchecked` reads already rely on. Used only on the hot
    /// MLOAD/MSTORE paths where the borrow is taken, used locally, and dropped.
    #[expect(unsafe_code, reason = "single active frame; no overlapping borrows")]
    #[inline]
    fn buf(&self) -> &Vec<u8> {
        unsafe { &*self.buffer.as_ptr() }
    }

    /// Unchecked mutable counterpart of [`Memory::buf`]; same invariant.
    #[expect(unsafe_code, reason = "single active frame; no overlapping borrows")]
    #[expect(clippy::mut_from_ref, reason = "single active frame; see buf()")]
    #[inline]
    fn buf_mut(&self) -> &mut Vec<u8> {
        unsafe { &mut *self.buffer.as_ptr() }
    }

    /// Returns the len of the current memory, from the current base.
    #[inline]
    pub fn len(&self) -> usize {
        self.len
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns a copy of the live byte slice for this frame (from `current_base` to
    /// `current_base + len`).  Used by the struct-log tracer for memory capture.
    pub fn live_bytes(&self) -> Vec<u8> {
        if self.len == 0 {
            return Vec::new();
        }
        let buf = self.buf();
        let end = self.current_base.saturating_add(self.len);
        buf.get(self.current_base..end)
            .map(<[u8]>::to_vec)
            .unwrap_or_default()
    }

    /// Resizes the from the current base to fit the memory specified at new_memory_size.
    ///
    /// Note: new_memory_size is increased to the next 32 byte multiple.
    #[inline(always)]
    pub fn resize(&mut self, new_memory_size: usize) -> Result<(), VMError> {
        if new_memory_size == 0 {
            return Ok(());
        }

        let new_memory_size = new_memory_size
            .checked_next_multiple_of(WORD_SIZE_IN_BYTES_USIZE)
            .ok_or(OutOfBounds)?;

        let current_len = self.len();

        if new_memory_size <= current_len {
            return Ok(());
        }

        self.len = new_memory_size;

        let buffer = self.buf_mut();

        #[allow(clippy::arithmetic_side_effects)]
        let real_new_memory_size = new_memory_size + self.current_base;

        if real_new_memory_size > buffer.len() {
            // when resizing, avoid really small resizes.
            let new_size = real_new_memory_size.next_multiple_of(64);
            buffer.resize(new_size, 0);
        }

        Ok(())
    }

    /// Load `size` bytes from the given offset. Returning a Bytes.
    #[inline]
    pub fn load_range(&mut self, offset: usize, size: usize) -> Result<Bytes, VMError> {
        if size == 0 {
            return Ok(Bytes::new());
        }

        let new_size = offset.checked_add(size).ok_or(OutOfBounds)?;
        self.resize(new_size)?;

        let true_offset = offset.wrapping_add(self.current_base);

        let buf = self.buf();

        // SAFETY: resize already makes sure bounds are correct.
        #[allow(unsafe_code)]
        unsafe {
            Ok(Bytes::copy_from_slice(buf.get_unchecked(
                true_offset..(true_offset.wrapping_add(size)),
            )))
        }
    }

    /// Borrow `size` bytes from the given offset and pass them to `f`, without
    /// allocating a `Bytes` copy of the range.
    ///
    /// `load_range` reads through `self.buf()`, whose `Ref` guard cannot
    /// outlive this call — so a `-> &[u8]` accessor can't be written. Callers that
    /// only need to read the range (e.g. hashing) take the borrow via this closure
    /// instead. Semantics match `load_range` exactly, including zero-padding reads
    /// past the current length (handled by `resize`).
    #[inline]
    pub fn with_range<R>(
        &mut self,
        offset: usize,
        size: usize,
        f: impl FnOnce(&[u8]) -> R,
    ) -> Result<R, VMError> {
        if size == 0 {
            return Ok(f(&[]));
        }

        let new_size = offset.checked_add(size).ok_or(OutOfBounds)?;
        self.resize(new_size)?;

        let true_offset = offset.wrapping_add(self.current_base);

        let buf = self.buf();

        // SAFETY: resize already makes sure bounds are correct.
        #[allow(unsafe_code)]
        let range = unsafe { buf.get_unchecked(true_offset..(true_offset.wrapping_add(size))) };
        Ok(f(range))
    }

    /// Load N bytes from the given offset.
    #[inline(always)]
    pub fn load_range_const<const N: usize>(&mut self, offset: usize) -> Result<[u8; N], VMError> {
        let new_size = offset.checked_add(N).ok_or(OutOfBounds)?;
        self.resize(new_size)?;

        let true_offset = offset.checked_add(self.current_base).ok_or(OutOfBounds)?;

        let buf = self.buf();
        // SAFETY: resize already makes sure bounds are correct.
        #[allow(unsafe_code)]
        unsafe {
            Ok(*buf
                .get_unchecked(true_offset..(true_offset.wrapping_add(N)))
                .as_ptr()
                .cast::<[u8; N]>())
        }
    }

    /// Load a word from at the given offset.
    #[inline(always)]
    pub fn load_word(&mut self, offset: usize) -> Result<U256, VMError> {
        let value: [u8; 32] = self.load_range_const(offset)?;
        Ok(u256_from_big_endian_const(value))
    }

    /// Stores the given data and data size at the given offset.
    ///
    /// Internal use.
    #[inline(always)]
    fn store(&self, data: &[u8], at_offset: usize, data_size: usize) -> Result<(), VMError> {
        if data_size == 0 {
            return Ok(());
        }

        let real_offset = self.current_base.wrapping_add(at_offset);

        let buffer = self.buf_mut();

        let real_data_size = data_size.min(data.len());

        // SAFETY: Used internally, resize always called before this function.
        #[allow(clippy::indexing_slicing, clippy::arithmetic_side_effects)]
        #[allow(unsafe_code)]
        unsafe {
            std::ptr::copy_nonoverlapping(
                data.get_unchecked(..real_data_size).as_ptr(),
                buffer
                    .get_unchecked_mut(real_offset..(real_offset + real_data_size))
                    .as_mut_ptr(),
                real_data_size,
            );
        }

        Ok(())
    }

    /// Stores the given data at the given offset.
    #[inline(always)]
    pub fn store_data(&mut self, offset: usize, data: &[u8]) -> Result<(), VMError> {
        if data.is_empty() {
            return Ok(());
        }
        let new_size = offset.checked_add(data.len()).ok_or(OutOfBounds)?;
        self.resize(new_size)?;
        self.store(data, offset, data.len())
    }

    /// Stores data and zero-pads up to total_size at the given offset.
    #[inline(always)]
    pub fn store_data_zero_padded(
        &mut self,
        offset: usize,
        data: &[u8],
        total_size: usize,
    ) -> Result<(), VMError> {
        if total_size == 0 {
            return Ok(());
        }

        let new_size = offset.checked_add(total_size).ok_or(OutOfBounds)?;
        self.resize(new_size)?;

        let copy_size = data.len().min(total_size);
        if copy_size > 0 {
            self.store(data, offset, copy_size)?;
        }

        #[allow(clippy::arithmetic_side_effects)]
        if copy_size < total_size {
            // SAFETY: copy_size < total_size and offset + total_size didn't overflow (checked above),
            // so offset + copy_size cannot overflow.
            let zero_offset = offset.wrapping_add(copy_size);
            let zero_size = total_size - copy_size;
            let real_offset = self.current_base.wrapping_add(zero_offset);
            let buffer = self.buf_mut();

            // resize ensures bounds are correct
            #[expect(unsafe_code)]
            unsafe {
                buffer
                    .get_unchecked_mut(real_offset..real_offset.wrapping_add(zero_size))
                    .fill(0);
            }
        }

        Ok(())
    }

    /// Stores a word at the given offset, resizing memory if needed.
    #[inline(always)]
    pub fn store_word(&mut self, offset: usize, word: U256) -> Result<(), VMError> {
        let new_size: usize = offset
            .checked_add(WORD_SIZE_IN_BYTES_USIZE)
            .ok_or(OutOfBounds)?;

        self.resize(new_size)?;
        self.store(&u256_to_big_endian(word), offset, WORD_SIZE_IN_BYTES_USIZE)?;
        Ok(())
    }

    /// Like [`resize`], but `new_memory_size` MUST already be rounded up to a
    /// multiple of `WORD_SIZE_IN_BYTES_USIZE` (e.g. the value returned by
    /// `calculate_memory_size`). Skips the internal re-rounding the caller has
    /// already performed for the gas charge.
    #[inline]
    pub fn resize_prerounded(&mut self, new_memory_size: usize) -> Result<(), VMError> {
        debug_assert_eq!(new_memory_size % WORD_SIZE_IN_BYTES_USIZE, 0);
        if new_memory_size == 0 {
            return Ok(());
        }
        let current_len = self.len();
        if new_memory_size <= current_len {
            return Ok(());
        }
        self.len = new_memory_size;
        let buffer = self.buf_mut();
        #[allow(clippy::arithmetic_side_effects)]
        let real_new_memory_size = new_memory_size + self.current_base;
        if real_new_memory_size > buffer.len() {
            let new_size = real_new_memory_size.next_multiple_of(64);
            buffer.resize(new_size, 0);
        }
        Ok(())
    }

    /// [`store_word`] variant taking the pre-rounded memory size (see
    /// [`resize_prerounded`]) so the word-rounding is computed once by the
    /// caller instead of twice per MSTORE.
    pub fn store_word_prerounded(
        &mut self,
        offset: usize,
        word: U256,
        new_memory_size: usize,
    ) -> Result<(), VMError> {
        self.resize_prerounded(new_memory_size)?;
        self.store(&u256_to_big_endian(word), offset, WORD_SIZE_IN_BYTES_USIZE)?;
        Ok(())
    }

    /// [`load_word`] variant taking the pre-rounded memory size (see
    /// [`resize_prerounded`]) so MLOAD computes the word-rounding once (for its
    /// gas charge) instead of twice.
    pub fn load_word_prerounded(
        &mut self,
        offset: usize,
        new_memory_size: usize,
    ) -> Result<U256, VMError> {
        self.resize_prerounded(new_memory_size)?;
        let true_offset = offset.checked_add(self.current_base).ok_or(OutOfBounds)?;
        let buf = self.buf();
        // SAFETY: resize_prerounded above guarantees `true_offset + 32` is in bounds.
        #[allow(unsafe_code)]
        let value: [u8; WORD_SIZE_IN_BYTES_USIZE] = unsafe {
            *buf.get_unchecked(true_offset..(true_offset.wrapping_add(WORD_SIZE_IN_BYTES_USIZE)))
                .as_ptr()
                .cast::<[u8; WORD_SIZE_IN_BYTES_USIZE]>()
        };
        Ok(u256_from_big_endian_const(value))
    }

    /// Copies memory within 2 offsets. Like a memmove.
    ///
    /// Resizes if needed, because one can copy from "expanded memory", which is initialized with zeroes.
    pub fn copy_within(
        &mut self,
        from_offset: usize,
        to_offset: usize,
        size: usize,
    ) -> Result<(), VMError> {
        if size == 0 {
            return Ok(());
        }

        self.resize(
            to_offset
                .max(from_offset)
                .checked_add(size)
                .ok_or(InternalError::Overflow)?,
        )?;

        let true_from_offset = from_offset
            .checked_add(self.current_base)
            .ok_or(OutOfBounds)?;

        let true_to_offset = to_offset
            .checked_add(self.current_base)
            .ok_or(OutOfBounds)?;
        let buffer = self.buf_mut();

        buffer.copy_within(
            true_from_offset
                ..(true_from_offset
                    .checked_add(size)
                    .ok_or(InternalError::Overflow)?),
            true_to_offset,
        );

        Ok(())
    }

    #[inline(always)]
    pub fn store_zeros(&mut self, offset: usize, size: usize) -> Result<(), VMError> {
        if size == 0 {
            return Ok(());
        }

        let new_size = offset.checked_add(size).ok_or(OutOfBounds)?;
        self.resize(new_size)?;

        let real_offset = self.current_base.wrapping_add(offset);
        let buffer = self.buf_mut();

        // resize ensures bounds are correct
        #[expect(unsafe_code)]
        unsafe {
            buffer
                .get_unchecked_mut(real_offset..(real_offset.wrapping_add(size)))
                .fill(0);
        }

        Ok(())
    }
}

impl Default for Memory {
    fn default() -> Self {
        Self::new()
    }
}

/// When a memory expansion is triggered, only the additional bytes of memory
/// must be paid for.
#[inline]
pub fn expansion_cost(new_memory_size: usize, current_memory_size: usize) -> Result<u64, VMError> {
    let cost = if new_memory_size <= current_memory_size {
        0
    } else {
        // We already know new_memory_size > current_memory_size,
        // and cost(x) > cost(y) where x > y, so cost should not underflow.
        cost(new_memory_size)?.wrapping_sub(cost(current_memory_size)?)
    };
    Ok(cost)
}

/// The total cost for a given memory size.
/// Gas cost should always be computed in u64
#[inline]
fn cost(memory_size: usize) -> Result<u64, VMError> {
    let memory_size = u64::try_from(memory_size).map_err(|_| InternalError::TypeConversion)?;

    // memory size measured in 32 byte words
    let words = memory_size.div_ceil(WORD_SIZE_IN_BYTES_U64);

    // Cost(words) ≈ floor(words^2 / q) + 3 * words
    // For this to overflow memory size in words should be 2^32, which is impossible.
    #[expect(clippy::arithmetic_side_effects)]
    let gas_cost = words * words / MEMORY_EXPANSION_QUOTIENT + 3 * words;

    Ok(gas_cost)
}

#[inline]
pub fn calculate_memory_size(offset: usize, size: usize) -> Result<usize, VMError> {
    if size == 0 {
        return Ok(0);
    }

    offset
        .checked_add(size)
        .and_then(|sum| sum.checked_next_multiple_of(WORD_SIZE_IN_BYTES_USIZE))
        .ok_or(OutOfBounds.into())
}