use core::{
alloc::{GlobalAlloc, Layout},
sync::atomic::{AtomicUsize, Ordering::Relaxed},
};
pub struct TalcAllocator(
talc::cell::TalcSyncCell<talc::wasm::WasmGrowAndExtend, talc::wasm::WasmBinning>,
);
impl TalcAllocator {
pub const fn new() -> Self {
Self(talc::cell::TalcSyncCell::new_wasm(
talc::wasm::WasmGrowAndExtend::new(),
))
}
}
impl Default for TalcAllocator {
fn default() -> Self {
Self::new()
}
}
unsafe impl GlobalAlloc for TalcAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
unsafe { self.0.alloc(layout) }
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
unsafe {
self.0.dealloc(ptr, layout);
}
}
unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
unsafe { self.0.alloc_zeroed(layout) }
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
unsafe { self.0.realloc(ptr, layout, new_size) }
}
}
const CHUNK_SIZE: usize = 64;
const PAGE_SIZE: usize = 64 * 1024;
unsafe extern "C" {
safe static __heap_base: usize;
}
pub struct BumpAllocator {
offset: AtomicUsize,
capacity: AtomicUsize,
free_8: AtomicUsize,
free_16: AtomicUsize,
free_32: AtomicUsize,
free_64: AtomicUsize,
free_128: AtomicUsize,
free_256: AtomicUsize,
free_512: AtomicUsize,
free_1024: AtomicUsize,
free_2048: AtomicUsize,
free_4096: AtomicUsize,
free_8192: AtomicUsize,
free_16384: AtomicUsize,
free_32768: AtomicUsize,
}
fn space_for(layout: Layout) -> usize {
usize::div_ceil(layout.size(), CHUNK_SIZE) * CHUNK_SIZE
}
fn pop_free_list(list: &AtomicUsize) -> Option<*mut u8> {
let offset = list.swap(0, Relaxed);
if offset != 0 {
// safety: offset is a valid pointer for these operations when non-zero
unsafe {
let ptr = offset as *mut u8;
let next = *ptr.cast::<usize>();
list.swap(next, Relaxed);
Some(ptr)
}
} else {
None
}
}
unsafe fn push_free_list(list: &AtomicUsize, ptr: *mut u8) {
// safety: all of the pointers here are valid
unsafe {
let new_offset = ptr.addr();
let current_head = list.swap(new_offset, Relaxed);
*ptr.cast::<usize>() = current_head;
}
}
impl BumpAllocator {
pub const fn new() -> Self {
Self {
offset: AtomicUsize::new(8),
capacity: AtomicUsize::new(PAGE_SIZE - 8),
free_8: AtomicUsize::new(0),
free_16: AtomicUsize::new(0),
free_32: AtomicUsize::new(0),
free_64: AtomicUsize::new(0),
free_128: AtomicUsize::new(0),
free_256: AtomicUsize::new(0),
free_512: AtomicUsize::new(0),
free_1024: AtomicUsize::new(0),
free_2048: AtomicUsize::new(0),
free_4096: AtomicUsize::new(0),
free_8192: AtomicUsize::new(0),
free_16384: AtomicUsize::new(0),
free_32768: AtomicUsize::new(0),
}
}
}
impl Default for BumpAllocator {
fn default() -> Self {
Self::new()
}
}
unsafe impl GlobalAlloc for BumpAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
self.offset.fetch_max(__heap_base, Relaxed);
let space_needed = space_for(layout);
let ptr = match space_needed {
8 => pop_free_list(&self.free_8),
16 => pop_free_list(&self.free_16),
32 => pop_free_list(&self.free_32),
64 => pop_free_list(&self.free_64),
128 => pop_free_list(&self.free_128),
256 => pop_free_list(&self.free_256),
512 => pop_free_list(&self.free_512),
1024 => pop_free_list(&self.free_1024),
2048 => pop_free_list(&self.free_2048),
4096 => pop_free_list(&self.free_4096),
8192 => pop_free_list(&self.free_8192),
16384 => pop_free_list(&self.free_16384),
32768 => pop_free_list(&self.free_32768),
_ => None,
};
if let Some(ptr) = ptr {
return ptr;
}
let offset = self.offset.fetch_add(space_needed, Relaxed);
let pointer = offset as *mut u8;
let available_space = self.capacity.load(Relaxed) - offset;
if available_space < layout.size() {
let pages = layout.size().div_ceil(PAGE_SIZE);
core::arch::wasm32::memory_grow(0, pages);
self.capacity.fetch_add(pages * PAGE_SIZE, Relaxed);
}
pointer
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
let space_used = space_for(layout);
let ptr_offset = ptr.addr();
if self
.offset
.compare_exchange(ptr_offset + space_used, ptr_offset, Relaxed, Relaxed)
.is_ok()
{
return;
}
// safety: ptr must be valid and on the heap
unsafe {
match space_used {
8 => push_free_list(&self.free_8, ptr),
16 => push_free_list(&self.free_16, ptr),
32 => push_free_list(&self.free_32, ptr),
64 => push_free_list(&self.free_64, ptr),
128 => push_free_list(&self.free_128, ptr),
256 => push_free_list(&self.free_256, ptr),
512 => push_free_list(&self.free_512, ptr),
1024 => push_free_list(&self.free_1024, ptr),
2048 => push_free_list(&self.free_2048, ptr),
4096 => push_free_list(&self.free_4096, ptr),
8192 => push_free_list(&self.free_8192, ptr),
16384 => push_free_list(&self.free_16384, ptr),
32768 => push_free_list(&self.free_32768, ptr),
_ => {}
};
}
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
let space_used = space_for(layout);
if new_size <= space_used {
return ptr;
}
unsafe {
self.dealloc(ptr, layout);
self.alloc(Layout::from_size_align(new_size, layout.align()).unwrap_unchecked())
}
}
}
|
