From 2f4cab3cc7b77598f01cefd9fa47e71df2a59017 Mon Sep 17 00:00:00 2001
From: kayomn <kayomn@kayomn.net>
Date: Sun, 29 Oct 2023 21:44:02 +0000
Subject: [PATCH] Rewrite stacking arena allocator to be correct

---
 source/coral/arena.zig | 137 +++++++++++++++++++++--------------------
 1 file changed, 71 insertions(+), 66 deletions(-)

diff --git a/source/coral/arena.zig b/source/coral/arena.zig
index 7ba2b1d..2408007 100644
--- a/source/coral/arena.zig
+++ b/source/coral/arena.zig
@@ -7,35 +7,47 @@ const list = @import("./list.zig");
 const math = @import("./math.zig");
 
 pub const Stacking = struct {
-	page_allocator: io.Allocator,
-	min_page_size: usize,
-	allocations: list.Stack(usize),
-	pages: list.Stack(Page),
+	allocator: io.Allocator,
+	min_region_size: usize,
+	head_region: ?*Region,
+	tail_region: ?*Region,
 
-	const AllocationsList = list.Stack(usize);
+	const Region = struct {
+		next: ?*Region,
+		count: usize,
+		capacity: usize,
 
-	const Page = struct {
-		buffer: []io.Byte,
-		used: usize,
+		fn allocate(self: *Region, region_size: usize) []usize {
+			debug.assert(self.can_allocate(region_size));
 
-		fn available(self: Page) usize {
-			return self.buffer.len - self.used;
+			const offset = (@sizeOf(Region) / @alignOf(usize)) + self.count;
+			const allocation = @as([*]usize, @ptrCast(self))[offset .. (offset + region_size)];
+
+			self.count += region_size;
+
+			return allocation;
+		}
+
+		fn as_raw(self: *Region) []usize {
+			return @as([*]usize, @ptrCast(self))[0 .. ((@sizeOf(Region) / @alignOf(usize)) + self.capacity)];
+		}
+
+		fn can_allocate(self: Region, region_size: usize) bool {
+			return (self.count + region_size) <= self.capacity;
 		}
 	};
 
-	const PageList = list.Stack(Page);
+	fn allocate_region(self: *Stacking, requested_region_size: usize) io.AllocationError!*Region {
+		const region_size = @max(requested_region_size, self.min_region_size);
+		const region = @as(*Region, @ptrCast(@alignCast(try self.allocator.reallocate(null, @alignOf(Region) + (@alignOf(usize) * region_size)))));
 
-	fn allocate_page(self: *Stacking, page_size: usize) io.AllocationError!*Page {
-		var buffer = try self.page_allocator.reallocate(null, page_size);
+		region.* = .{
+			.next = null,
+			.count = 0,
+			.capacity = region_size,
+		};
 
-		errdefer self.page_allocator.deallocate(buffer);
-
-		try self.pages.push_one(.{
-			.buffer = buffer,
-			.used = 0,
-		});
-
-		return (self.current_page() orelse unreachable);
+		return region;
 	}
 
 	pub fn as_allocator(self: *Stacking) io.Allocator {
@@ -45,67 +57,60 @@ pub const Stacking = struct {
 		});
 	}
 
-	fn current_page(self: Stacking) ?*Page {
-		if (self.pages.values.len == 0) {
-			return null;
-		}
-
-		return &self.pages.values[self.pages.values.len - 1];
-	}
-
 	pub fn free(self: *Stacking) void {
-		for (self.pages.values) |page| {
-			self.page_allocator.deallocate(page.buffer);
+		while (self.head_region) |region| {
+			const next_region = region.next;
+
+			self.allocator.deallocate(region.as_raw());
+
+			self.head_region = next_region;
 		}
 
-		self.pages.free();
-		self.allocations.free();
+		self.head_region = null;
+		self.tail_region = null;
 	}
 
-	pub fn deallocate(_: *Stacking, _: []io.Byte) void {
+	fn deallocate(_: *Stacking, _: []io.Byte) void {
 		// TODO: Decide how to implement.
+		unreachable;
 	}
 
-	pub fn reallocate(self: *Stacking, return_address: usize, existing_allocation: ?[]io.Byte, size: usize) io.AllocationError![]io.Byte {
-		// TODO: Safety-check existing allocation is from allocator or null.
-		_ = return_address;
-
-		const alignment = @as(usize, 4);
-		const aligned_size = (size + alignment - 1) & ~(alignment - 1);
-
-		if (self.pages.values.len == 0) {
-			const page = try self.allocate_page(@max(self.min_page_size, aligned_size));
-
-			page.used = size;
-
-			return page.buffer[0 .. size];
+	fn reallocate(self: *Stacking, _: usize, allocation: ?[]io.Byte, byte_size: usize) io.AllocationError![]io.Byte {
+		if (allocation) |buffer| {
+			if (byte_size < buffer.len) {
+				return buffer[0 .. byte_size];
+			}
 		}
 
-		var page = self.current_page() orelse unreachable;
+		const region_size = (byte_size + @sizeOf(usize) - 1) / @sizeOf(usize);
 
-		if (page.available() <= aligned_size) {
-			page = try self.allocate_page(@max(self.min_page_size, aligned_size));
+		const tail_region = self.tail_region orelse {
+			const region = try self.allocate_region(region_size);
+
+			self.tail_region = region;
+			self.head_region = self.tail_region;
+
+			return @as([*]io.Byte, @ptrCast(region.allocate(region_size)))[0 .. byte_size];
+		};
+
+		if (!tail_region.can_allocate(region_size)) {
+			const region = try self.allocate_region(region_size);
+
+			tail_region.next = region;
+			self.tail_region = region;
+
+			return @as([*]io.Byte, @ptrCast(region.allocate(region_size)))[0 .. byte_size];
 		}
 
-		debug.assert(page.available() >= size);
-
-		defer page.used += aligned_size;
-
-		const reallocation = page.buffer[page.used .. (page.used + size)];
-
-		if (existing_allocation) |allocation| {
-			io.copy(reallocation, allocation);
-		}
-
-		return reallocation;
+		return @as([*]io.Byte, @ptrCast(tail_region.allocate(region_size)))[0 .. byte_size];
 	}
 
-	pub fn make(allocator: io.Allocator, min_page_size: usize) Stacking {
+	pub fn make(allocator: io.Allocator, min_region_size: usize) Stacking {
 		return Stacking{
-			.allocations = AllocationsList.make(allocator),
-			.pages = PageList.make(allocator),
-			.page_allocator = allocator,
-			.min_page_size = min_page_size,
+			.allocator = allocator,
+			.min_region_size = min_region_size,
+			.head_region = null,
+			.tail_region = null,
 		};
 	}
 };