Skip to content

Commit

Permalink
Simplify guest's memory allocator
Browse files Browse the repository at this point in the history
Extracted from #631
  • Loading branch information
matthiasgoergens committed Dec 10, 2024
1 parent 33901fd commit 73b4779
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 27 deletions.
35 changes: 16 additions & 19 deletions ceno_rt/src/allocator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,26 +4,34 @@
use core::alloc::{GlobalAlloc, Layout};

struct SimpleAllocator {
next_alloc: usize,
next_alloc: *mut u8,
}

extern "C" {
// The address of this variable is the start of the heap (growing upwards).
static mut _sheap: u8;
}

#[global_allocator]
static mut HEAP: SimpleAllocator = SimpleAllocator {
next_alloc: &raw mut _sheap,
};

unsafe impl GlobalAlloc for SimpleAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
// SAFETY: Single threaded, so nothing else can touch this while we're working.
let mut heap_pos = HEAP.next_alloc;

let align = layout.align();
// `Layout` contract forbids making a `Layout` with align=0, or align not power of 2.
// So we can safely use subtraction and a mask to ensure alignment without worrying about UB.
let offset = heap_pos & (align - 1);
if offset != 0 {
heap_pos = heap_pos.strict_add(align.strict_sub(offset));
}
core::hint::assert_unchecked(align.is_power_of_two());
core::hint::assert_unchecked(align != 0);
heap_pos = heap_pos.add(heap_pos.align_offset(align));

let ptr = heap_pos as *mut u8;
let ptr = heap_pos;
// Panic on overflow. We don't want to wrap around, and overwrite stack etc.
// (We could also return a null pointer, but only malicious programs would ever hit this.)
heap_pos = heap_pos.strict_add(layout.size());
heap_pos = heap_pos.add(layout.size());

HEAP.next_alloc = heap_pos;
ptr
Expand All @@ -36,14 +44,3 @@ unsafe impl GlobalAlloc for SimpleAllocator {
/// Never deallocate.
unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
}

#[global_allocator]
// We initialize `next_alloc` to 0xFFFF_FFFF to indicate that the heap has not been initialized.
// The value is chosen to make any premature allocation fail.
static mut HEAP: SimpleAllocator = SimpleAllocator {
next_alloc: 0xFFFF_FFFF,
};

pub unsafe fn init_heap() {
HEAP.next_alloc = core::ptr::from_ref::<u8>(&crate::_sheap).cast::<u8>() as usize;
}
11 changes: 3 additions & 8 deletions ceno_rt/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -79,19 +79,14 @@ macro_rules! entry {
/// _start_rust is called by the assembly entry point and it calls the Rust main().
#[no_mangle]
unsafe extern "C" fn _start_rust() -> ! {
allocator::init_heap();
{
extern "C" {
fn bespoke_entrypoint();
}
bespoke_entrypoint();
extern "C" {
fn bespoke_entrypoint();
}
bespoke_entrypoint();
halt(0)
}

extern "C" {
// The address of this variable is the start of the stack (growing downwards).
static _stack_start: u8;
// The address of this variable is the start of the heap (growing upwards).
static _sheap: u8;
}

0 comments on commit 73b4779

Please sign in to comment.