fix(zkvm): optimize memory usage for low memory configurations

Co-Authored-By: John Guibas <[email protected]>
succinctlabs · Dec 12, 2024 · dfebebf · dfebebf
1 parent 9569a6e
commit dfebebf
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/crates/zkvm/entrypoint/src/heap.rs b/crates/zkvm/entrypoint/src/heap.rs
@@ -5,7 +5,7 @@ use core::sync::atomic::{AtomicUsize, Ordering};
 use crate::syscalls::sys_alloc_aligned;
 
 /// A block in our free list
-#[repr(C)]
+#[repr(C, align(8))]
 struct FreeBlock {
     size: usize,
     next: Option<NonNull<FreeBlock>>,

diff --git a/crates/zkvm/entrypoint/src/syscalls/memory.rs b/crates/zkvm/entrypoint/src/syscalls/memory.rs
@@ -15,7 +15,7 @@
 // Memory addresses must be lower than BabyBear prime.
 const MAX_MEMORY: usize = 0x78000000;
 
-const RESERVED_SIZE: usize = 1024 * 1024; // 1MB reserved section
+const RESERVED_SIZE: usize = MAX_MEMORY / 64;
 static mut RESERVED_POS: usize = 0;
 static mut RESERVED_START: usize = 0;
 
@@ -37,6 +37,10 @@ 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 {
+            // Check if reserved section would exceed memory limit
+            if heap_pos + RESERVED_SIZE > MAX_MEMORY {
+                panic!("Reserved section would exceed memory limit");
+            }
             RESERVED_START = heap_pos;
             RESERVED_POS = heap_pos;
             heap_pos += RESERVED_SIZE;
@@ -63,7 +67,6 @@ pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u
 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;