Implement an arena allocator

1 files changed, 153 insertions, 0 deletions

diff --git a/oryxc/src/arena.rs b/oryxc/src/arena.rs
new file mode 100644
index 0000000..32d02ba
--- /dev/null
+++ b/oryxc/src/arena.rs
@@ -0,0 +1,153 @@
+use std::alloc::{
+	self,
+	Layout,
+};
+use std::cell::Cell;
+use std::ptr::{
+	self,
+	NonNull,
+};
+use std::slice;
+use std::sync::Mutex;
+
+use crate::err;
+
+#[derive(Copy, Clone)]
+struct RawBlock {
+	ptr:    NonNull<u8>,
+	layout: Layout,
+}
+
+pub struct GlobalArena {
+	blksz:  usize,
+	blocks: Mutex<Vec<RawBlock>>,
+}
+
+impl GlobalArena {
+	pub fn new(blksz: usize) -> Self {
+		Self {
+			blksz,
+			blocks: Mutex::new(Vec::new()),
+		}
+	}
+
+	fn allocate_block(&self, layout: Layout) -> RawBlock {
+		let layout = Layout::from_size_align(
+			layout.size().max(self.blksz),
+			layout.align(),
+		)
+		.unwrap_or_else(|e| err!(e, "allocation error"));
+		let ptr = NonNull::new(unsafe { alloc::alloc(layout) })
+			.unwrap_or_else(|| alloc::handle_alloc_error(layout));
+
+		let block = RawBlock { ptr, layout };
+		self.blocks.lock().unwrap().push(block.clone());
+		return block;
+	}
+}
+
+impl Drop for GlobalArena {
+	fn drop(&mut self) {
+		for block in self.blocks.lock().unwrap().iter() {
+			unsafe {
+				alloc::dealloc(block.ptr.as_ptr(), block.layout);
+			}
+		}
+	}
+}
+
+pub struct LocalArena<'a> {
+	parent: &'a GlobalArena,
+	cursor: Cell<*mut u8>,
+	end:    Cell<*mut u8>,
+}
+
+impl<'a> LocalArena<'a> {
+	pub fn new(parent: &'a GlobalArena) -> Self {
+		return Self {
+			parent,
+			cursor: Cell::new(ptr::null_mut()),
+			end: Cell::new(ptr::null_mut()),
+		};
+	}
+
+	fn alloc_raw(&self, layout: Layout) -> NonNull<u8> {
+		let cursor = self.cursor.get() as usize;
+
+		let align_offset =
+			cursor.wrapping_add(layout.align() - 1) & !(layout.align() - 1);
+		let Some(next_cursor) = align_offset.checked_add(layout.size()) else {
+			err!("allocation overflow");
+		};
+
+		return if next_cursor <= self.end.get() as usize {
+			self.cursor.set(next_cursor as *mut u8);
+			unsafe { NonNull::new_unchecked(align_offset as *mut u8) }
+		} else {
+			self.alloc_slow(layout)
+		};
+	}
+
+	#[cold]
+	fn alloc_slow(&self, layout: Layout) -> NonNull<u8> {
+		let block = self.parent.allocate_block(layout);
+		let beg = block.ptr.as_ptr() as usize;
+		let end = beg + block.layout.size();
+		let next_cursor = beg + layout.size();
+		self.cursor.set(next_cursor as *mut u8);
+		self.end.set(end as *mut u8);
+		return unsafe { NonNull::new_unchecked(beg as *mut u8) };
+	}
+
+	pub fn alloc<T>(&self, value: T) -> &'a mut T {
+		let ptr = self.alloc_raw(Layout::new::<T>()).cast::<T>().as_ptr();
+		unsafe {
+			ptr::write(ptr, value);
+			return &mut *ptr;
+		}
+	}
+
+	pub fn alloc_slice<T>(&self, len: usize) -> &'a mut [T] {
+		let layout = Layout::array::<T>(len)
+			.unwrap_or_else(|e| err!(e, "allocation error"));
+		let ptr = self.alloc_raw(layout).cast::<T>().as_ptr();
+		return unsafe { slice::from_raw_parts_mut(ptr, len) };
+	}
+}
+
+#[test]
+fn test_alloc_slice() {
+	let arena_global = GlobalArena::new(8);
+	let arena_local_1 = LocalArena::new(&arena_global);
+	let arena_local_2 = LocalArena::new(&arena_global);
+
+	let s1 = arena_local_1.alloc_slice(8);
+	let s2 = arena_local_2.alloc_slice(4);
+	assert_eq!(s1.len(), 8);
+	assert_eq!(s2.len(), 4);
+
+	for i in 0..s1.len() {
+		s1[i] = i;
+	}
+	for i in 0..s2.len() {
+		s2[i] = i;
+	}
+}
+
+#[test]
+fn test_arena_grows() {
+	let arena_global = GlobalArena::new(8);
+	let arena_local = LocalArena::new(&arena_global);
+
+	let s1 = arena_local.alloc_slice(8);
+	let s2 = arena_local.alloc_slice(69);
+	assert_eq!(s1.len(), 8);
+	assert_eq!(s2.len(), 69);
+
+	for i in 0..s1.len() {
+		s1[i] = i;
+	}
+	for i in 0..s2.len() {
+		s2[i] = i;
+	}
+}