Add leak detection to Ona heap allocator

source/coral/arena.zig

@@ -49,7 +49,7 @@ pub const Stacking = struct {
 	}
 
 	fn allocate_page(self: *Stacking, page_size: usize) io.AllocationError!*Page {
-		var buffer = try io.allocate_many(u8, page_size, self.base_allocator);
+		var buffer = try io.allocate_many(self.base_allocator, page_size, u8);
 
 		errdefer io.deallocate(self.base_allocator, buffer);
 

source/coral/io.zig

@@ -176,7 +176,7 @@ pub const GrowingBuffer = struct {
 
 pub const Writer = Generator(?usize, []const u8);
 
-pub fn allocate_many(comptime Type: type, amount: usize, allocator: Allocator) AllocationError![]Type {
+pub fn allocate_many(allocator: Allocator, amount: usize, comptime Type: type) AllocationError![]Type {
 	if (@sizeOf(Type) == 0) {
 		@compileError("Cannot allocate memory for 0-byte type " ++ @typeName(Type));
 	}
@@ -232,14 +232,12 @@ pub fn copy(target: []u8, source: []const u8) void {
 }
 
 pub fn deallocate(allocator: Allocator, allocation: anytype) void {
-	const Allocation = @TypeOf(allocation);
-
-	switch (@typeInfo(Allocation)) {
-		.Pointer => |allocation_pointer| {
+	switch (@typeInfo(@TypeOf(allocation))) {
+		.Pointer => |pointer| {
 			_ = allocator.invoke(.{
-				.allocation = switch (allocation_pointer.size) {
-					.One => @ptrCast([*]u8, allocation)[0 .. @sizeOf(Allocation)],
-					.Slice => @ptrCast([*]u8, allocation.ptr)[0 .. (@sizeOf(Allocation) * allocation.len)],
+				.allocation = switch (pointer.size) {
+					.One => @ptrCast([*]u8, allocation)[0 .. @sizeOf(pointer.child)],
+					.Slice => @ptrCast([*]u8, allocation.ptr)[0 .. (@sizeOf(pointer.child) * allocation.len)],
 					.Many, .C => @compileError("length of allocation must be known to deallocate"),
 				},
 

source/coral/list.zig

@@ -87,7 +87,7 @@ pub fn Stack(comptime Value: type) type {
 		///
 		pub fn grow(self: *Self, allocator: io.Allocator, growth_amount: usize) io.AllocationError!void {
 			const grown_capacity = self.capacity + growth_amount;
-			const values = (try io.allocate_many(Value, grown_capacity, allocator))[0 .. self.values.len];
+			const values = (try io.allocate_many(allocator, grown_capacity, Value))[0 .. self.values.len];
 
 			errdefer io.deallocate(allocator, values);
 
@@ -96,7 +96,7 @@ pub fn Stack(comptime Value: type) type {
 					values[index] = self.values[index];
 				}
 
-				io.deallocate(allocator, self.values);
+				io.deallocate(allocator, self.values.ptr[0 .. self.capacity]);
 			}
 
 			self.values = values;

source/coral/slab.zig

@@ -95,7 +95,7 @@ pub fn Map(comptime index_int: std.builtin.Type.Int, comptime Value: type) type
 		///
 		pub fn grow(self: *Self, allocator: io.Allocator, growth_amount: usize) io.AllocationError!void {
 			const grown_capacity = self.table.len + growth_amount;
-			const entries = try io.allocate_many(Entry, grown_capacity, allocator);
+			const entries = try io.allocate_many(allocator, grown_capacity, Entry);
 
 			errdefer io.deallocate(allocator, entries);
 

source/coral/table.zig

@@ -228,7 +228,7 @@ pub fn Hashed(comptime Key: type, comptime Value: type, comptime keyer: Keyer(Ke
 		pub fn rehash(self: *Self, allocator: io.Allocator, requested_range: usize) io.AllocationError!void {
 			const old_table = self.table;
 
-			self.table = try io.allocate_many(?Entry, math.max(requested_range, self.count), allocator);
+			self.table = try io.allocate_many(allocator, math.max(requested_range, self.count), ?Entry);
 
 			errdefer {
 				io.deallocate(allocator, self.table);

source/ona/heap.zig

@@ -2,53 +2,171 @@ const coral = @import("coral");
 
 const ext = @import("./ext.zig");
 
+const std = @import("std");
+
+///
+///
+///
+const stack_trace_frame_size = if (std.debug.sys_can_stack_trace) 4 else 0;
+
+///
+///
+///
 const Context = struct {
-	live_allocations: usize,
+	allocation_info_head: ?*AllocationInfo = null,
 
-	const Self = @This();
+	///
+	///
+	///
+	const AllocationInfo = struct {
+		stack_frames: [stack_trace_frame_size]usize,
+		stack_trace: std.builtin.StackTrace,
+		next_info: ?*AllocationInfo,
+		size: usize,
+	};
 
-	const empty_allocation = [0]u8{};
+	///
+	///
+	///
+	fn allocate(self: *Context, size: usize) ?[]u8 {
+		const allocation_info_size = @sizeOf(AllocationInfo);
+		const total_allocation_size = allocation_info_size + size;
+		const allocation = ext.SDL_malloc(total_allocation_size) orelse return null;
+		const allocation_info = @ptrCast(*AllocationInfo, @alignCast(@alignOf(AllocationInfo), allocation));
 
-	fn reallocate(self: *Self, options: coral.io.AllocationOptions) ?[]u8 {
-		if (options.size == 0) {
-			if (options.allocation) |allocation| {
-				if (allocation.ptr != &empty_allocation) {
-					ext.SDL_free(allocation.ptr);
+		allocation_info.* = .{
+			.size = size,
+			.next_info = self.allocation_info_head,
+			.stack_frames = [_]usize{0} ** stack_trace_frame_size,
+
+			.stack_trace = .{
+				.instruction_addresses = &allocation_info.stack_frames,
+				.index = 0,
+			},
+		};
+
+		std.debug.captureStackTrace(null, &allocation_info.stack_trace);
+
+		self.allocation_info_head = allocation_info;
+
+		return @ptrCast([*]u8, allocation)[allocation_info_size .. total_allocation_size];
+	}
+
+	///
+	///
+	///
+	fn deallocate(self: *Context, allocation: []u8) void {
+		const target_allocation_info = @intToPtr(*AllocationInfo, @ptrToInt(allocation.ptr) - @sizeOf(AllocationInfo));
+
+		if (target_allocation_info.size != allocation.len) {
+			@panic("incorrect allocation length");
+		}
+
+		if (self.allocation_info_head) |allocation_info_head| {
+			if (target_allocation_info == allocation_info_head) {
+				self.allocation_info_head = allocation_info_head.next_info;
+
+				ext.SDL_free(target_allocation_info);
+
+				return;
+			}
+
+			var previous_allocation_info = allocation_info_head;
+			var current_allocation_info = allocation_info_head.next_info;
+
+			while (current_allocation_info) |allocation_info| {
+				if (allocation_info == target_allocation_info) {
+					previous_allocation_info.next_info = allocation_info.next_info;
+
+					ext.SDL_free(target_allocation_info);
+
+					return;
 				}
 
-				self.live_allocations -= 1;
-
-				return null;
-			}
-
-			self.live_allocations += 1;
-
-			return &empty_allocation;
-		}
-
-		if (options.allocation) |allocation| {
-			if (ext.SDL_realloc(allocation.ptr, options.size)) |reallocation| {
-				self.live_allocations += 1;
-
-				return @ptrCast([*]u8, reallocation)[0 .. options.size];
+				previous_allocation_info = allocation_info;
+				current_allocation_info = allocation_info.next_info;
 			}
 		}
 
-		if (ext.SDL_malloc(options.size)) |allocation| {
-			self.live_allocations += 1;
+		@panic("double-free detected");
+	}
 
-			return @ptrCast([*]u8, allocation)[0 .. options.size];
+	///
+	///
+	///
+	fn reallocate(self: *Context, allocation: []u8, size: usize) ?[]u8 {
+		const allocation_info_size = @sizeOf(AllocationInfo);
+		const target_allocation_info = @intToPtr(*AllocationInfo, @ptrToInt(allocation.ptr) - allocation_info_size);
+
+		if (target_allocation_info.size != allocation.len) {
+			@panic("incorrect allocation length");
 		}
 
-		return null;
+		if (self.allocation_info_head) |allocation_info_head| {
+			if (target_allocation_info == allocation_info_head) {
+				self.allocation_info_head = allocation_info_head.next_info;
+
+				return @ptrCast([*]u8, ext.SDL_realloc(target_allocation_info, size) orelse {
+					return null;
+				})[allocation_info_size .. allocation_info_size + size];
+			}
+
+			var previous_allocation_info = allocation_info_head;
+			var current_allocation_info = allocation_info_head.next_info;
+
+			while (current_allocation_info) |allocation_info| {
+				if (allocation_info == target_allocation_info) {
+					previous_allocation_info.next_info = allocation_info.next_info;
+
+					return @ptrCast([*]u8, ext.SDL_realloc(target_allocation_info, size) orelse {
+						return null;
+					})[allocation_info_size .. allocation_info_size + size];
+				}
+
+				previous_allocation_info = allocation_info;
+				current_allocation_info = allocation_info.next_info;
+			}
+		}
+
+		@panic("use-after-free detected");
 	}
 };
 
-var context = Context{
-	.live_allocations = 0,
-};
+///
+///
+///
+var context = Context{};
 
 ///
 /// Heap allocator.
 ///
-pub const allocator = coral.io.Allocator.bind(Context, &context, Context.reallocate);
+pub const allocator = coral.io.Allocator.bind(Context, &context, struct {
+	fn reallocate(self: *Context, options: coral.io.AllocationOptions) ?[]u8 {
+		if (options.size == 0) {
+			if (options.allocation) |allocation| {
+				self.deallocate(allocation);
+
+				return null;
+			}
+
+			return self.allocate(0);
+		}
+
+		if (options.allocation) |allocation| {
+			return self.reallocate(allocation, options.size);
+		}
+
+		return self.allocate(options.size);
+	}
+}.reallocate);
+
+///
+///
+///
+pub fn trace_allocations() void {
+	var current_allocation_info = context.allocation_info_head;
+
+	while (current_allocation_info) |allocation_info| : (current_allocation_info = allocation_info) {
+		std.debug.dumpStackTrace(allocation_info.stack_trace);
+	}
+}

source/ona/kym/Environment.zig

@@ -361,7 +361,7 @@ pub fn new_array(self: *Self) coral.io.AllocationError!types.Val {
 }
 
 pub fn new_object(self: *Self, userdata: []const u8, info: ObjectInfo) coral.io.AllocationError!types.Val {
-	const allocation = try coral.io.allocate_many(u8, userdata.len, self.allocator);
+	const allocation = try coral.io.allocate_many(self.allocator, userdata.len, u8);
 
 	errdefer coral.io.deallocate(self.allocator, allocation);
 

source/ona/ona.zig

@@ -135,4 +135,6 @@ pub fn run_app(base_file_system: file.System) void {
 			ext.SDL_RenderPresent(renderer);
 		}
 	}
+
+	heap.trace_allocations();
 }