A Buffer Stack Allocator for Zig

stack_alloc.zig

const std = @import("std");

const testing = std.testing;
const Allocator = std.mem.Allocator;
const Error = Allocator.Error;
const assert = std.debug.assert;
const mem = std.mem;
const math = std.math;

pub fn StackAllocator(comptime buffer_size: comptime_int, comptime max_alignment: comptime_int) type {
    const Offset = if (buffer_size < 2 * 1024 * 1024 * 1024) u32 else u64;
    const no_mark = ~@as(Offset, 0);
    return struct {
        const Self = @This();

        buffer: [buffer_size]u8 align(std.math.max(max_alignment, @alignOf(Offset))) = undefined,
        allocator: std.mem.Allocator = .{ .allocFn = alloc, .resizeFn = resize },
        current_offset: Offset = 0,
        mark_offset: Offset = no_mark,

        pub fn alloc(allocator: *Allocator, len: usize, ptr_align: u29, len_align: u29, ret_addr: usize) Error![]u8 {
            const self = @fieldParentPtr(Self, "allocator", allocator);
            std.debug.assert(ptr_align <= max_alignment);
            std.debug.assert(len <= buffer_size);
            const offset = std.mem.alignForwardGeneric(Offset, self.current_offset, ptr_align);
            const end = offset + len;
            if (end > buffer_size) return error.OutOfMemory;
            self.current_offset = @intCast(Offset, end);
            return self.buffer[offset..end];
        }

        pub fn resize(allocator: *Allocator, buf: []u8, buf_align: u29, new_len: usize, len_align: u29, ret_addr: usize) Error!usize {
            const self = @fieldParentPtr(Self, "allocator", allocator);
            const oldOffset = @ptrToInt(buf.ptr) - @ptrToInt(&self.buffer);
            const oldEnd = oldOffset + buf.len;
            if (oldEnd == self.current_offset) {
                const newEnd = oldOffset + new_len;
                if (newEnd > buffer_size) return error.OutOfMemory;
                self.current_offset = @intCast(Offset, newEnd);
                return new_len;
            } else if (new_len <= buf.len) {
                return new_len;
            } else {
                return error.OutOfMemory;
            }
        }

        pub fn pushMark(self: *Self) Error!void {
            const markOffset = std.mem.alignForwardGeneric(Offset, self.current_offset, @alignOf(Offset));
            const markEnd = markOffset + @sizeOf(Offset);
            if (markEnd >= buffer_size) return error.OutOfMemory;
            const markStorage = @intToPtr(*Offset, @ptrToInt(&self.buffer) + markOffset);
            markStorage.* = self.mark_offset;
            self.mark_offset = markOffset;
            self.current_offset = markEnd;
        }

        pub fn popMark(self: *Self) void {
            assert(self.mark_offset != no_mark);
            const markStorage = @intToPtr(*Offset, @ptrToInt(&self.buffer) + self.mark_offset);
            self.current_offset = self.mark_offset;
            self.mark_offset = markStorage.*;
        }

        pub fn reset(self: *Self) void {
            self.current_offset = 0;
            self.mark_offset = no_mark;
        }
    };
}

test "StackAllocator" {
    var stack = StackAllocator(256, 256){};
    const allocator = &stack.allocator;

    testing.expectEqual(@as(u32, 0), stack.current_offset);
    testing.expectEqual(@as(u32, 0xFFFFFFFF), stack.mark_offset);
    try stack.pushMark();
    testing.expectEqual(@as(u32, 4), stack.current_offset);
    testing.expectEqual(@as(u32, 0), stack.mark_offset);
    stack.popMark();
    testing.expectEqual(@as(u32, 0), stack.current_offset);
    testing.expectEqual(@as(u32, 0xFFFFFFFF), stack.mark_offset);

    var a = try allocator.alloc(u16, 8);
    testing.expectEqual(@as(u32, 16), stack.current_offset);
    a = try allocator.realloc(a, 12);
    testing.expectEqual(@as(u32, 24), stack.current_offset);
    a = try allocator.realloc(a, 5);
    testing.expectEqual(@as(u32, 10), stack.current_offset);
    
    try stack.pushMark();
    testing.expectEqual(@as(u32, 16), stack.current_offset);
    testing.expectEqual(@as(u32, 12), stack.mark_offset);
    stack.popMark();
    testing.expectEqual(@as(u32, 12), stack.current_offset);
    testing.expectEqual(@as(u32, 0xFFFFFFFF), stack.mark_offset);

    try stack.pushMark();
    var b = try allocator.alloc(u128, 2);
    testing.expectEqual(@as(u32, 16+32), stack.current_offset);
    // a = try allocator.resize(a, 1); // pending #6881
    testing.expectEqual(@as(u32, 16+32), stack.current_offset);
    stack.popMark();
    testing.expectEqual(@as(u32, 12), stack.current_offset);
    testing.expectEqual(@as(u32, 0xFFFFFFFF), stack.mark_offset);
    // testing.expectError(error.OutOfMemory, allocator.resize(a, 2)); // pending #6881

    stack.reset();
    testing.expectEqual(@as(u32, 0), stack.current_offset);
    testing.expectEqual(@as(u32, 0xFFFFFFFF), stack.mark_offset);
}