1 files changed, 204 insertions, 0 deletions
diff --git a/src/allocator.rs b/src/allocator.rs
new file mode 100644
index 0000000..eaaba5f
--- /dev/null
+++ b/src/allocator.rs
@@ -0,0 +1,204 @@
+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())
+		}
+	}
+}