ona/source/ona/heap.zig

const builtin = @import("builtin");

const coral = @import("coral");

const ext = @import("./ext.zig");

const std = @import("std");

///
/// Recorded allocation info state.
///
const AllocationInfo = struct {
	trace: AllocationTrace,
	next_info: ?*AllocationInfo,
	size: usize,
};

///
/// Recorded stack trace of allocation call site.
///
/// *Note* this structure is reduced to zero bytes in released builds optimized for speed or size.
///
const AllocationTrace = std.debug.ConfigurableTrace(2, 4, switch (builtin.mode) {
	.Debug, .ReleaseSafe => true,
	.ReleaseFast, .ReleaseSmall => false,
});

///
/// Heap allocation context.
///
const Context = struct {
	allocation_info_head: ?*AllocationInfo = null,

	///
	/// Attempts to allocate a buffer of `size` length from `self`, with `return_address` as the location of the
	/// allocation request origin.
	///
	/// A reference to the allocated buffer is returned via a slice if the allocation was successful, otherwise `null`
	/// is returned.
	///
	/// *Note* the returned buffer must be deallocated with [deallocate] before program exit or it will cause a memory
	/// leak.
	///
	/// *Note* allocation checks are disabled in release builds optimized for speed or size.
	///
	fn allocate(self: *Context, size: usize, return_address: usize) ?[]u8 {
		switch (builtin.mode) {
			.Debug, .ReleaseSafe => {
				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));

				allocation_info.* = .{
					.size = size,
					.next_info = self.allocation_info_head,
					.trace = .{},
				};

				allocation_info.trace.addAddr(return_address, "");

				self.allocation_info_head = allocation_info;

				return @ptrCast([*]u8, allocation)[allocation_info_size .. total_allocation_size];
			},

			.ReleaseFast, .ReleaseSmall => {
				return @ptrCast([*]u8, ext.SDL_malloc(size) orelse return null)[0 .. size];
			},
		}
	}

	///
	/// Returns the assumed pointer to the [AllocationInfo] address of `allocation`.
	///
	fn allocation_info_of(allocation: [*]u8) *AllocationInfo {
		return @intToPtr(*AllocationInfo, @ptrToInt(allocation) - @sizeOf(AllocationInfo));
	}

	///
	/// Deallocates a the allocation buffer referenced by `allocation`.
	///
	/// *Note* the pointer and length of `allocation` must match valid values known to `allocator` otherwise safety-
	/// checked behavior will occur.
	///
	/// *Note* allocation checks are disabled in release builds optimized for speed or size.
	///
	fn deallocate(self: *Context, allocation: []u8) void {
		switch (builtin.mode) {
			.Debug, .ReleaseSafe => {
				const target_allocation_info = allocation_info_of(allocation.ptr);

				if (target_allocation_info.size != allocation.len) {
					@panic("incorrect allocation length for deallocating");
				}

				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;
						}

						previous_allocation_info = allocation_info;
						current_allocation_info = allocation_info.next_info;
					}
				}

				@panic("incorrect allocation address for deallocating");
			},

			.ReleaseFast, .ReleaseSmall => {
				ext.SDL_free(allocation.ptr);
			},
		}
	}

	///
	/// Attempts to reallocate the buffer referenced by `allocation` to be `size` length from `self`.
	///
	/// A reference to the reallocated buffer is returned via a slice if the reallocation was successful, otherwise
	/// `null` is returned.
	///
	/// *Note* the returned buffer must be deallocated with [deallocate] before program exit or it will cause a memory
	/// leak.
	///
	/// *Note* the pointer and length of `allocation` must match valid values known to `allocator` otherwise safety-
	/// checked behavior will occur.
	///
	/// *Note* the allocation referenced by `allocation` should be considered invalid once the function returns,
	/// discarding it in favor of the return value.
	///
	/// *Note* allocation checks are disabled in release builds optimized for speed or size.
	///
	fn reallocate(self: *Context, allocation: []u8, size: usize) ?[]u8 {
		switch (builtin.mode) {
			.Debug, .ReleaseSafe => {
				const target_allocation_info = allocation_info_of(allocation.ptr);

				if (target_allocation_info.size != allocation.len) {
					@panic("incorrect allocation length for reallocating");
				}

				const allocation_info_size = @sizeOf(AllocationInfo);

				if (self.allocation_info_head) |allocation_info_head| {
					if (target_allocation_info == allocation_info_head) {
						self.allocation_info_head = allocation_info_head.next_info;

						const allocation_address = ext.SDL_realloc(target_allocation_info, size) orelse return null;

						target_allocation_info.size = size;

						return @ptrCast([*]u8, allocation_address)[
							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;

							const allocation_address = ext.SDL_realloc(target_allocation_info, size) orelse return null;

							target_allocation_info.size = size;

							return @ptrCast([*]u8, allocation_address)[
								allocation_info_size .. (allocation_info_size + size)];
						}

						previous_allocation_info = allocation_info;
						current_allocation_info = allocation_info.next_info;
					}
				}

				@panic("incorrect allocation address for reallocating");
			},

			.ReleaseFast, .ReleaseSmall => {
				return @ptrCast([*]u8, ext.SDL_realloc(allocation.ptr, size) orelse return null)[0 .. size];
			},
		}
	}
};

///
/// Heap context.
///
var context = Context{};

///
/// Heap allocator.
///
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, options.return_address);
		}

		if (options.allocation) |allocation| {
			return self.reallocate(allocation, options.size);
		}

		return self.allocate(options.size, options.return_address);
	}
}.reallocate);

///
/// Checks for any allocations belonging to the process heap allocated through the [allocator] interface that are still
/// alive and reports the stack traces of any detected allocations to stderr along with the allocation address and
/// length.
///
/// *Note* this function becomes a no-op in release builds optimized for speed or size.
///
pub fn trace_leaks() void {
	switch (builtin.mode) {
		.Debug, .ReleaseSafe => {
			var current_allocation_info = context.allocation_info_head;

			while (current_allocation_info) |allocation_info| : (current_allocation_info = allocation_info.next_info) {
				std.debug.print("{d} byte leak at 0x{x} detected:\n", .{
					allocation_info.size,
					@ptrToInt(allocation_info) + @sizeOf(AllocationInfo),
				});

				allocation_info.trace.dump();
			}
		},

		.ReleaseFast, .ReleaseSmall => {},
	}
}