mem.rs

//Memory layout for the EVM
//uses NonNull for better performance since we don't have to check for null pointers and it's technically impossible to point to a null pointer here

use std::alloc::{Layout, alloc, dealloc, realloc};
use std::ptr::NonNull;

pub struct EvmMemory {
    // We use NonNull to satisfy the "must be non-null" rule for references
    base: NonNull<u8>,
    capacity: usize,
    last_checkpoint: usize,
}

impl EvmMemory {
    //New starts with 4096 capacity
    #[inline]
    pub fn new() -> Self {
        let layout = Layout::from_size_align(4096, 64).unwrap();
        let ptr = unsafe { alloc(layout) };

        //Initialize memory to zero
        unsafe { std::ptr::write_bytes(ptr, 0, 4096) };

        Self {
            base: NonNull::new(ptr).unwrap(),
            capacity: 4096,
            last_checkpoint: 0,
        }
    }

    #[inline]
    pub fn reset(&mut self) {
        unsafe {
            // Only zero what we actually used
            std::ptr::write_bytes(self.base.as_ptr(), 0, self.last_checkpoint);
        }
        self.last_checkpoint = 0;
    }

    #[inline]
    pub fn get_word_ptr(&self, offset: usize) -> *const u8 {
        // Preferred Strict Provenance way:
        // .addr() and .with_addr() help tools like Miri trace the
        // "permission" from 'base' to the result.
        let base_ptr = self.base.as_ptr();
        let new_addr = base_ptr.addr().wrapping_add(offset);
        base_ptr.with_addr(new_addr)
    }

    //Warning: unsafe caller needs to ensure offset is valid
    #[inline]
    pub unsafe fn get_word_ref<'a>(&'a self, offset: usize) -> &'a [u8; 32] {
        // Since EVM words are always 32 bytes, a reference to an array
        // is often faster than a slice because the '32' is known at compile time.
        let ptr = self.get_word_ptr(offset) as *const [u8; 32];

        // Checklist check: Is the ptr aligned to 1 (u8) or 32 ([u8; 32])?
        // Raw arrays of u8 have alignment 1, so this is always safe.
        unsafe { &*ptr }
    }

    #[inline]
    pub fn resize(&mut self, new_size: usize) {
        if new_size <= self.capacity {
            self.last_checkpoint = new_size;
            return;
        }

        let new_capacity = Self::calculate_new_capacity(new_size);
        let alignment = 64;

        unsafe {
            let old_ptr = self.base.as_ptr();
            let old_layout = Layout::from_size_align_unchecked(self.capacity, alignment);

            //lmaooooooooo
            let new_ptr = if self.capacity == 0 {
                alloc(Layout::from_size_align_unchecked(new_capacity, alignment))
            } else {
                // Ensure we use the exact same alignment as the original alloc
                realloc(old_ptr, old_layout, new_capacity)
            };

            //better error handling
            self.base = NonNull::new(new_ptr).expect("allocation failed");

            // Zero the new portion.
            // Because capacity is 4096-based, this pointer is 64-byte aligned.
            std::ptr::write_bytes(
                self.base.as_ptr().add(self.capacity),
                0,
                new_capacity - self.capacity,
            );

            self.capacity = new_capacity;
            self.last_checkpoint = new_size;
        }
    }

    //Happy path
    #[inline]
    pub fn resize_happy_path(&mut self, new_size: usize) {
        let new_capacity = Self::calculate_new_capacity(new_size);
        let alignment = 64;

        unsafe {
            let old_ptr = self.base.as_ptr();
            let old_layout = Layout::from_size_align_unchecked(self.capacity, alignment);

            // Ensure we use the exact same alignment as the original alloc
            let new_ptr = realloc(old_ptr, old_layout, new_capacity);

            //better error handling
            self.base = NonNull::new(new_ptr).unwrap();

            // Zero the new portion.
            // Because capacity is 4096-based, this pointer is 64-byte aligned.
            std::ptr::write_bytes(
                self.base.as_ptr().add(self.capacity),
                0,
                new_capacity.unchecked_sub(self.capacity),
            );

            self.capacity = new_capacity;
            self.last_checkpoint = new_size;
        }
    }

    #[inline]
    pub unsafe fn drop_manual(&mut self) {
        unsafe {
            let layout = Layout::from_size_align_unchecked(self.capacity, 64);
            dealloc(self.base.as_ptr(), layout);
        }
    }

    #[inline]
    //calc needed
    pub fn calculate_new_capacity(needed: usize) -> usize {
        const INITIAL: usize = 4096;
        const GROWTH_FACTOR: usize = 150;
        const MAX_CAPACITY: usize = 32 * 1024 * 1024;

        if needed <= INITIAL {
            return INITIAL;
        }

        let mut cap = INITIAL;
        while cap < needed {
            cap = (cap * GROWTH_FACTOR / 100).min(MAX_CAPACITY);
            if cap >= MAX_CAPACITY {
                break;
            }
        }
        cap.max(needed).min(MAX_CAPACITY)
    }
}

impl Drop for EvmMemory {
    fn drop(&mut self) {
        unsafe {
            // SAFETY:
            // 1. We know capacity >= 4096 (invariant from new()).
            // 2. We know 'base' was allocated with this capacity and ALIGNMENT.
            let layout = Layout::from_size_align_unchecked(self.capacity, 64);
            dealloc(self.base.as_ptr(), layout);
        }
    }
}