feat(heaperion): add a growable heap

Author: mathisbot

beskar-lib/src/lib.rs

@@ -7,6 +7,7 @@ extern crate alloc;
 
 pub use beskar_core::syscall::ExitCode;
 use beskar_core::{syscall::SyscallExitCode, time::Duration};
+use core::sync::atomic::{AtomicBool, Ordering};
 use hyperdrive::call_once;
 
 mod arch;
@@ -20,9 +21,20 @@ pub mod surface;
 mod sys;
 pub mod time;
 
+static PANIC_NESTED: AtomicBool = AtomicBool::new(false);
+
+/// Returns `true` if the current thread is already panicking (i.e., if we're in a nested panic).
+pub fn panicking() -> bool {
+    PANIC_NESTED.load(Ordering::SeqCst)
+}
+
 #[panic_handler]
 fn panic(info: &::core::panic::PanicInfo) -> ! {
-    println!("Panic occurred: {}", info);
+    if !panicking() {
+        PANIC_NESTED.store(true, Ordering::SeqCst);
+        println!("Panic occurred: {}", info);
+    }
+
     sys::sc_exit(ExitCode::Failure);
 }
 
@@ -70,13 +82,7 @@ macro_rules! entry_point {
 pub fn __init() {
     call_once!({
         // Heap
-        {
-            let heap_size = mem::HEAP_SIZE;
-            let res = mem::mmap(heap_size, None, mem::MemoryProtection::ReadWrite)
-                .expect("Memory mapping failed");
-            unsafe { mem::init_heap(res.as_ptr(), heap_size.try_into().unwrap()) };
-        }
-
+        mem::init_heap();
         // Time
         time::init();
     });

beskar-lib/src/mem.rs

@@ -1,17 +1,18 @@
 use crate::error::{MemoryError, MemoryErrorKind, MemoryResult};
 use beskar_core::arch::paging::{M4KiB, MemSize as _};
 use core::{num::NonZeroU64, ptr::NonNull};
+use heaperion::{DefaultGrowableHeap, HybridAllocator};
 use hyperdrive::locks::mcs::MUMcsLock;
 
-static ALLOCATOR: MUMcsLock<heaperion::Heap> = MUMcsLock::uninit();
-
-struct Heap;
+static ALLOCATOR: MUMcsLock<DefaultGrowableHeap<Mmap, 4>> = MUMcsLock::uninit();
 
 #[global_allocator]
 static HEAP: Heap = Heap;
 
-pub(crate) const HEAP_SIZE: u64 = 20 * 1024 * 1024; // 20 MiB
-beskar_core::static_assert!(HEAP_SIZE.is_multiple_of(M4KiB::SIZE));
+const HEAP_START_SIZE: u64 = 16 * 1024 * 1024; // 16 MiB
+beskar_core::static_assert!(HEAP_START_SIZE.is_multiple_of(M4KiB::SIZE));
+
+struct Heap;
 
 unsafe impl core::alloc::GlobalAlloc for Heap {
     unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 {
@@ -30,8 +31,15 @@ unsafe impl core::alloc::GlobalAlloc for Heap {
 
 #[inline]
 /// Initialize the heap allocator
-pub(crate) unsafe fn init_heap(start: *mut u8, size: usize) {
-    ALLOCATOR.init(unsafe { heaperion::Heap::new(start, size) }.unwrap());
+pub(super) fn init_heap() {
+    let size = HEAP_START_SIZE;
+    let start = mmap(size, None, MemoryProtection::ReadWrite).unwrap();
+
+    // SAFETY: `start` and `size` come from a successful call to `mmap` and are not used after this point.
+    let heap =
+        unsafe { HybridAllocator::new(start.as_ptr(), usize::try_from(size).unwrap()) }.unwrap();
+    let growable = DefaultGrowableHeap::new(heap, Mmap);
+    ALLOCATOR.init(growable);
 }
 
 /// Map memory into the address space
@@ -76,6 +84,77 @@ pub fn mprotect(ptr: *mut u8, size: u64, flags: MemoryProtection) -> bool {
     res.is_success()
 }
 
+pub struct MmapReadWrite {
+    ptr: NonNull<u8>,
+    size: u64,
+}
+
+impl MmapReadWrite {
+    #[inline]
+    /// Create a new read-write memory mapping of the given size.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if the memory cannot be mapped.
+    pub fn new(size: u64) -> MemoryResult<Self> {
+        let ptr = mmap(size, None, MemoryProtection::ReadWrite)?;
+        Ok(Self { ptr, size })
+    }
+
+    #[must_use]
+    #[inline]
+    pub const fn size(&self) -> u64 {
+        self.size
+    }
+
+    #[must_use]
+    #[inline]
+    pub const fn as_ptr(&self) -> *mut u8 {
+        self.ptr.as_ptr()
+    }
+
+    #[must_use]
+    #[inline]
+    #[expect(clippy::missing_panics_doc, reason = "Never panics")]
+    pub fn as_slice(&self) -> &[u8] {
+        let data = self.as_ptr();
+        let len = usize::try_from(self.size).unwrap();
+        unsafe { core::slice::from_raw_parts(data.cast::<u8>(), len) }
+    }
+
+    #[must_use]
+    #[inline]
+    #[expect(clippy::missing_panics_doc, reason = "Never panics")]
+    pub fn as_mut_slice(&mut self) -> &mut [u8] {
+        let data = self.as_ptr();
+        let len = usize::try_from(self.size).unwrap();
+        unsafe { core::slice::from_raw_parts_mut(data.cast::<u8>(), len) }
+    }
+}
+
+impl Drop for MmapReadWrite {
+    fn drop(&mut self) {
+        // Safety: `ptr` and `size` come from a previous, successful call to `mmap`
+        // and are not used after this point.
+        unsafe {
+            munmap(self.ptr.as_ptr(), self.size);
+        }
+    }
+}
+
+struct Mmap;
+unsafe impl heaperion::MemorySource for Mmap {
+    fn request(&mut self, min_size: usize) -> Option<(*mut u8, usize)> {
+        let ptr = mmap(
+            u64::try_from(min_size).unwrap(),
+            None,
+            MemoryProtection::ReadWrite,
+        )
+        .ok()?;
+        Some((ptr.as_ptr(), min_size))
+    }
+}
+
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 #[repr(u64)]
 pub enum MemoryProtection {

heaperion/README.md

@@ -28,3 +28,11 @@ Heaperion is a robust and efficient memory allocator designed for embedded syste
 - O(log(c)) allocation and deallocation
 - Power-of-two sized blocks
 - Automatic coalescing to reduce fragmentation
+
+### Hybrid Allocator
+
+Dispatch allocations to Slab or Buddy depending on the size.
+
+### GrowableHeap
+
+Wrapper around heaps that can grow at runtime.

heaperion/src/buddy.rs

@@ -72,15 +72,12 @@ impl BuddyAllocator {
             max_order,
         };
 
-        // Add the initial block to the appropriate free list
-        let initial_order = size_to_order(adjusted_size.next_power_of_two() / 2);
-        let initial_block_size = order_to_size(initial_order);
-
-        if initial_block_size <= adjusted_size {
-            // SAFETY: We've validated the heap region and aligned it properly
-            unsafe {
-                buddy.add_block_to_free_list(aligned_ptr, initial_order);
-            }
+        // Add the initial block to the free list.
+        // `max_order` is capped to `MAX_ORDER - 1`, so `order_to_size(max_order) <= adjusted_size`
+        // is guaranteed and we cannot go out-of-bounds on `free_lists`.
+        // SAFETY: We've validated the heap region and aligned it properly
+        unsafe {
+            buddy.add_block_to_free_list(aligned_ptr, max_order);
         }
 
         Ok(buddy)
@@ -148,6 +145,17 @@ impl BuddyAllocator {
         Ok(())
     }
 
+    /// Returns `true` if this allocator owns the given pointer.
+    ///
+    /// A pointer is owned when it falls within the heap region provided at construction.
+    #[must_use]
+    pub fn contains(&self, ptr: NonNull<u8>) -> bool {
+        let ptr_raw = ptr.as_ptr();
+        // SAFETY: heap_start + heap_size is within the original allocation
+        let heap_end = unsafe { self.heap_start.add(self.heap_size) };
+        ptr_raw >= self.heap_start && ptr_raw < heap_end
+    }
+
     /// Find and allocate a block of the given order
     fn find_block(&mut self, order: usize) -> Result<*mut u8> {
         // Try to find a block in the current order
@@ -219,11 +227,7 @@ impl BuddyAllocator {
                 }
 
                 // Merge with buddy - the merged block starts at the lower address
-                current_ptr = if current_ptr < buddy_ptr {
-                    current_ptr
-                } else {
-                    buddy_ptr
-                };
+                current_ptr = current_ptr.min(buddy_ptr);
                 current_order += 1;
             } else {
                 break;