Skip to content

Commit

Permalink
use externally defined allocator
Browse files Browse the repository at this point in the history
  • Loading branch information
ncitron committed Apr 19, 2024
1 parent b391f3c commit 3c5f0ec
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 44 deletions.
2 changes: 1 addition & 1 deletion library/std/src/sys/pal/zkvm/abi.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
extern "C" {
pub static _HEAP_PTR: u8;
pub fn sys_alloc(size: usize, align: usize) -> *mut u8;
}
46 changes: 3 additions & 43 deletions library/std/src/sys/pal/zkvm/alloc.rs
Original file line number Diff line number Diff line change
@@ -1,51 +1,11 @@
use crate::{alloc::{GlobalAlloc, Layout, System}, cell::UnsafeCell};
use super::abi::_HEAP_PTR;

static mut BUMP_ALLOC: BumpAllocator = BumpAllocator::new();
use crate::alloc::{GlobalAlloc, Layout, System};
use super::abi::sys_alloc;

#[stable(feature = "alloc_system_type", since = "1.28.0")]
unsafe impl GlobalAlloc for System {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
BUMP_ALLOC.alloc(layout)
}

unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
}

pub struct BumpAllocator {
offset: UnsafeCell<usize>,
}

impl BumpAllocator {
pub const fn new() -> Self {
Self {
offset: UnsafeCell::new(0),
}
}

pub fn free_memory(&self) -> usize {
heap_start() + (self.offset.get() as usize)
}
}

unsafe impl GlobalAlloc for BumpAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
let alloc_start = align_up(self.free_memory(), layout.align());
let alloc_end = alloc_start + layout.size();
*self.offset.get() = alloc_end - self.free_memory();

alloc_start as *mut u8
sys_alloc(layout.size(), layout.align())
}

unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
}

unsafe impl Sync for BumpAllocator {}

fn align_up(addr: usize, align: usize) -> usize {
(addr + align - 1) & !(align - 1)
}

fn heap_start() -> usize {
unsafe { _HEAP_PTR as *const u8 as usize }
}

0 comments on commit 3c5f0ec

Please sign in to comment.