feat(zkvm): implement freeing allocator and reserved memory section f…

…or read_vec

Co-Authored-By: John Guibas <[email protected]>

crates/zkvm/entrypoint/src/heap.rs

@@ -11,46 +11,44 @@ struct FreeBlock {
     next: Option<NonNull<FreeBlock>>,
 }
 
-/// A simple heap allocator with free list.
-///
-/// Uses a first-fit strategy with coalescing of adjacent free blocks.
-/// Designed for single-threaded embedded systems with a memory limit of 0x78000000.
-pub struct SimpleAlloc {
-    head: AtomicUsize,  // Stores the raw pointer value of the head FreeBlock
-}
+// Global free list head stored as raw pointer value
+static FREE_LIST_HEAD: AtomicUsize = AtomicUsize::new(0);
 
-// SAFETY: The allocator is thread-safe due to atomic operations
-unsafe impl Sync for SimpleAlloc {}
+/// A simple heap allocator that supports freeing memory.
+///
+/// Uses a first-fit strategy for allocation and maintains a free list
+/// for memory reuse. Designed for single-threaded embedded systems
+/// with a memory limit of 0x78000000.
+#[derive(Copy, Clone)]
+pub struct SimpleAlloc;
 
+// Implementation detail functions
 impl SimpleAlloc {
-    const fn new() -> Self {
-        Self {
-            head: AtomicUsize::new(0),
-        }
-    }
-
-    unsafe fn add_free_block(&self, ptr: *mut u8, size: usize) {
+    unsafe fn add_free_block(ptr: *mut u8, size: usize) {
         let block = ptr as *mut FreeBlock;
         (*block).size = size;
 
         loop {
-            let current_head = self.head.load(Ordering::Relaxed);
+            let current_head = FREE_LIST_HEAD.load(Ordering::Relaxed);
             (*block).next = NonNull::new(current_head as *mut FreeBlock);
 
-            if self.head.compare_exchange(
-                current_head,
-                block as usize,
-                Ordering::Release,
-                Ordering::Relaxed,
-            ).is_ok() {
+            if FREE_LIST_HEAD
+                .compare_exchange(
+                    current_head,
+                    block as usize,
+                    Ordering::Release,
+                    Ordering::Relaxed,
+                )
+                .is_ok()
+            {
                 break;
             }
         }
     }
 
-    unsafe fn find_block(&self, size: usize, align: usize) -> Option<(*mut u8, usize)> {
+    unsafe fn find_block(size: usize, align: usize) -> Option<(*mut u8, usize)> {
         let mut prev: Option<*mut FreeBlock> = None;
-        let mut current_ptr = self.head.load(Ordering::Acquire) as *mut FreeBlock;
+        let mut current_ptr = FREE_LIST_HEAD.load(Ordering::Acquire) as *mut FreeBlock;
 
         while !current_ptr.is_null() {
             let addr = current_ptr as *mut u8;
@@ -64,7 +62,7 @@ impl SimpleAlloc {
                 match prev {
                     Some(p) => (*p).next = next,
                     None => {
-                        self.head.store(next_raw, Ordering::Release);
+                        FREE_LIST_HEAD.store(next_raw, Ordering::Release);
                     }
                 }
                 return Some((aligned_addr, (*current_ptr).size));
@@ -82,7 +80,7 @@ unsafe impl GlobalAlloc for SimpleAlloc {
         let align = layout.align();
 
         // Try to find a block in free list
-        if let Some((ptr, _)) = self.find_block(size, align) {
+        if let Some((ptr, _)) = Self::find_block(size, align) {
             return ptr;
         }
 
@@ -92,9 +90,10 @@ unsafe impl GlobalAlloc for SimpleAlloc {
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
         let size = layout.size().max(core::mem::size_of::<FreeBlock>());
-        self.add_free_block(ptr, size);
+        Self::add_free_block(ptr, size);
     }
 }
 
-#[global_allocator]
-static ALLOCATOR: SimpleAlloc = SimpleAlloc::new();
+#[used]
+#[no_mangle]
+pub static HEAP_ALLOCATOR: SimpleAlloc = SimpleAlloc;

crates/zkvm/entrypoint/src/syscalls/memory.rs

@@ -15,6 +15,10 @@
 // Memory addresses must be lower than BabyBear prime.
 const MAX_MEMORY: usize = 0x78000000;
 
+const RESERVED_SIZE: usize = 1024 * 1024; // 1MB reserved section
+static mut RESERVED_POS: usize = 0;
+static mut RESERVED_START: usize = 0;
+
 #[allow(clippy::missing_safety_doc)]
 #[no_mangle]
 pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u8 {
@@ -32,6 +36,11 @@ pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u
 
     if heap_pos == 0 {
         heap_pos = unsafe { (&_end) as *const u8 as usize };
+        unsafe {
+            RESERVED_START = heap_pos;
+            RESERVED_POS = heap_pos;
+            heap_pos += RESERVED_SIZE;
+        }
     }
 
     let offset = heap_pos & (align - 1);
@@ -49,3 +58,22 @@ pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u
     unsafe { HEAP_POS = heap_pos };
     ptr
 }
+#[allow(clippy::missing_safety_doc)]
+#[no_mangle]
+pub unsafe extern "C" fn sys_alloc_reserved(bytes: usize, align: usize) -> *mut u8 {
+    let mut pos = RESERVED_POS;
+
+    // Align the position
+    let offset = pos & (align - 1);
+    if offset != 0 {
+        pos += align - offset;
+    }
+
+    let new_pos = pos + bytes;
+    if new_pos > RESERVED_START + RESERVED_SIZE {
+        panic!("Reserved memory section full");
+    }
+
+    RESERVED_POS = new_pos;
+    pos as *mut u8
+}

crates/zkvm/lib/src/io.rs

@@ -1,10 +1,7 @@
 #![allow(unused_unsafe)]
-use crate::{syscall_hint_len, syscall_hint_read, syscall_write};
+use crate::{sys_alloc_reserved, syscall_hint_len, syscall_hint_read, syscall_write};
 use serde::{de::DeserializeOwned, Serialize};
-use std::{
-    alloc::Layout,
-    io::{Result, Write},
-};
+use std::io::{Result, Write};
 
 /// The file descriptor for public values.
 pub const FD_PUBLIC_VALUES: u32 = 3;
@@ -43,24 +40,12 @@ impl Write for SyscallWriter {
 /// let data: Vec<u8> = sp1_zkvm::io::read_vec();
 /// ```
 pub fn read_vec() -> Vec<u8> {
-    // Round up to the nearest multiple of 4 so that the memory allocated is in whole words
     let len = unsafe { syscall_hint_len() };
     let capacity = (len + 3) / 4 * 4;
+    let ptr = unsafe { sys_alloc_reserved(capacity, 4) };
 
-    // Allocate a buffer of the required length that is 4 byte aligned
-    let layout = Layout::from_size_align(capacity, 4).expect("vec is too large");
-    let ptr = unsafe { std::alloc::alloc(layout) };
-
-    // SAFETY:
-    // 1. `ptr` was allocated using alloc
-    // 2. We assuume that the VM global allocator doesn't dealloc
-    // 3/6. Size is correct from above
-    // 4/5. Length is 0
-    // 7. Layout::from_size_align already checks this
     let mut vec = unsafe { Vec::from_raw_parts(ptr, 0, capacity) };
 
-    // Read the vec into uninitialized memory. The syscall assumes the memory is uninitialized,
-    // which should be true because the allocator does not dealloc, so a new alloc should be fresh.
     unsafe {
         syscall_hint_read(ptr, len);
         vec.set_len(len);

crates/zkvm/lib/src/lib.rs

@@ -97,6 +97,9 @@ extern "C" {
     /// Allocates a buffer aligned to the given alignment.
     pub fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u8;
 
+    /// Allocates from the reserved memory section
+    pub fn sys_alloc_reserved(bytes: usize, align: usize) -> *mut u8;
+
     /// Decompresses a BLS12-381 point.
     pub fn syscall_bls12381_decompress(point: &mut [u8; 96], is_odd: bool);