🎓 A1Learner
🔍
0 of 46 lessons completed

Zig Stack Allocators

A stack allocator is an allocator that hands out memory from one buffer in order. Freed blocks only count when they are freed in reverse order.

Allocating in order

Two blocks take 8 and 4 bytes. The used total is 12.

const std = @import("std");
const StackAllocator = struct {
    buffer: []u8,
    end: usize = 0,
    frames: [32]usize = undefined,
    frames_len: usize = 0,
    fn allocFn(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len >= self.frames.len) return null;
        const start = std.mem.alignForward(usize, self.end, alignment.toByteUnits());
        if (start + len > self.buffer.len) return null;
        self.frames[self.frames_len] = start;
        self.frames_len += 1;
        self.end = start + len;
        return self.buffer.ptr + start;
    }
    fn resizeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
        _ = alignment;
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len == 0) return false;
        const start = self.frames[self.frames_len - 1];
        if (@intFromPtr(memory.ptr) != @intFromPtr(self.buffer.ptr) + start) return false;
        if (start + new_len > self.buffer.len) return false;
        self.end = start + new_len;
        return true;
    }
    fn remapFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
        _ = ctx;
        _ = memory;
        _ = alignment;
        _ = new_len;
        _ = ret_addr;
        return null;
    }
    fn freeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
        _ = alignment;
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len == 0) return;
        const start = self.frames[self.frames_len - 1];
        if (@intFromPtr(memory.ptr) != @intFromPtr(self.buffer.ptr) + start) return;
        self.frames_len -= 1;
        self.end = start;
    }
    fn allocator(self: *StackAllocator) std.mem.Allocator {
        return .{ .ptr = self, .vtable = &.{ .alloc = allocFn, .resize = resizeFn, .remap = remapFn, .free = freeFn } };
    }
    fn used(self: *StackAllocator) usize {
        return self.end;
    }
};
pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var backing: [64]u8 = undefined;
    var stack = StackAllocator{ .buffer = &backing };
    const alloc = stack.allocator();
    const a = try alloc.alloc(u8, 8);
    const b = try alloc.alloc(u8, 4);
    a[0] = 1;
    b[0] = 2;
    try out.print("{d} {d} {d}\n", .{ a[0], b[0], stack.used() });
    alloc.free(b);
    alloc.free(a);
}

Freeing out of order

Freeing the older block first is ignored. Freeing the newest block rewinds the buffer.

const std = @import("std");
const StackAllocator = struct {
    buffer: []u8,
    end: usize = 0,
    frames: [32]usize = undefined,
    frames_len: usize = 0,
    fn allocFn(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len >= self.frames.len) return null;
        const start = std.mem.alignForward(usize, self.end, alignment.toByteUnits());
        if (start + len > self.buffer.len) return null;
        self.frames[self.frames_len] = start;
        self.frames_len += 1;
        self.end = start + len;
        return self.buffer.ptr + start;
    }
    fn resizeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
        _ = alignment;
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len == 0) return false;
        const start = self.frames[self.frames_len - 1];
        if (@intFromPtr(memory.ptr) != @intFromPtr(self.buffer.ptr) + start) return false;
        if (start + new_len > self.buffer.len) return false;
        self.end = start + new_len;
        return true;
    }
    fn remapFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
        _ = ctx;
        _ = memory;
        _ = alignment;
        _ = new_len;
        _ = ret_addr;
        return null;
    }
    fn freeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
        _ = alignment;
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len == 0) return;
        const start = self.frames[self.frames_len - 1];
        if (@intFromPtr(memory.ptr) != @intFromPtr(self.buffer.ptr) + start) return;
        self.frames_len -= 1;
        self.end = start;
    }
    fn allocator(self: *StackAllocator) std.mem.Allocator {
        return .{ .ptr = self, .vtable = &.{ .alloc = allocFn, .resize = resizeFn, .remap = remapFn, .free = freeFn } };
    }
    fn used(self: *StackAllocator) usize {
        return self.end;
    }
};
pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var backing: [64]u8 = undefined;
    var stack = StackAllocator{ .buffer = &backing };
    const alloc = stack.allocator();
    const a = try alloc.alloc(u8, 8);
    const b = try alloc.alloc(u8, 4);
    alloc.free(a);
    const after_a = stack.used();
    alloc.free(b);
    const after_b = stack.used();
    alloc.free(a);
    const after_both = stack.used();
    try out.print("{d} {d} {d}\n", .{ after_a, after_b, after_both });
}

Resizing the newest block

The 4-byte block is the newest block. Growing it to 6 bytes succeeds in place.

const std = @import("std");
const StackAllocator = struct {
    buffer: []u8,
    end: usize = 0,
    frames: [32]usize = undefined,
    frames_len: usize = 0,
    fn allocFn(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len >= self.frames.len) return null;
        const start = std.mem.alignForward(usize, self.end, alignment.toByteUnits());
        if (start + len > self.buffer.len) return null;
        self.frames[self.frames_len] = start;
        self.frames_len += 1;
        self.end = start + len;
        return self.buffer.ptr + start;
    }
    fn resizeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
        _ = alignment;
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len == 0) return false;
        const start = self.frames[self.frames_len - 1];
        if (@intFromPtr(memory.ptr) != @intFromPtr(self.buffer.ptr) + start) return false;
        if (start + new_len > self.buffer.len) return false;
        self.end = start + new_len;
        return true;
    }
    fn remapFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
        _ = ctx;
        _ = memory;
        _ = alignment;
        _ = new_len;
        _ = ret_addr;
        return null;
    }
    fn freeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
        _ = alignment;
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len == 0) return;
        const start = self.frames[self.frames_len - 1];
        if (@intFromPtr(memory.ptr) != @intFromPtr(self.buffer.ptr) + start) return;
        self.frames_len -= 1;
        self.end = start;
    }
    fn allocator(self: *StackAllocator) std.mem.Allocator {
        return .{ .ptr = self, .vtable = &.{ .alloc = allocFn, .resize = resizeFn, .remap = remapFn, .free = freeFn } };
    }
    fn used(self: *StackAllocator) usize {
        return self.end;
    }
};
pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var backing: [64]u8 = undefined;
    var stack = StackAllocator{ .buffer = &backing };
    const alloc = stack.allocator();
    const a = try alloc.alloc(u8, 4);
    a[0] = 7;
    const b = try alloc.realloc(a, 6);
    try out.print("{d} {d}\n", .{ b.len, b[0] });
    alloc.free(b);
}

Example

const std = @import("std");
const StackAllocator = struct {
    buffer: []u8,
    end: usize = 0,
    frames: [32]usize = undefined,
    frames_len: usize = 0,
    fn allocFn(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len >= self.frames.len) return null;
        const start = std.mem.alignForward(usize, self.end, alignment.toByteUnits());
        if (start + len > self.buffer.len) return null;
        self.frames[self.frames_len] = start;
        self.frames_len += 1;
        self.end = start + len;
        return self.buffer.ptr + start;
    }
    fn resizeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
        _ = alignment;
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len == 0) return false;
        const start = self.frames[self.frames_len - 1];
        if (@intFromPtr(memory.ptr) != @intFromPtr(self.buffer.ptr) + start) return false;
        if (start + new_len > self.buffer.len) return false;
        self.end = start + new_len;
        return true;
    }
    fn remapFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
        _ = ctx;
        _ = memory;
        _ = alignment;
        _ = new_len;
        _ = ret_addr;
        return null;
    }
    fn freeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
        _ = alignment;
        _ = ret_addr;
        const self: *StackAllocator = @ptrCast(@alignCast(ctx));
        if (self.frames_len == 0) return;
        const start = self.frames[self.frames_len - 1];
        if (@intFromPtr(memory.ptr) != @intFromPtr(self.buffer.ptr) + start) return;
        self.frames_len -= 1;
        self.end = start;
    }
    fn allocator(self: *StackAllocator) std.mem.Allocator {
        return .{ .ptr = self, .vtable = &.{ .alloc = allocFn, .resize = resizeFn, .remap = remapFn, .free = freeFn } };
    }
    fn used(self: *StackAllocator) usize {
        return self.end;
    }
};
pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var backing: [64]u8 = undefined;
    var stack = StackAllocator{ .buffer = &backing };
    const alloc = stack.allocator();
    const t1 = try alloc.alloc(u8, 3);
    t1[0] = 2;
    t1[1] = 4;
    t1[2] = 6;
    const t2 = try alloc.alloc(u8, 2);
    t2[0] = 10;
    t2[1] = 20;
    var total: usize = 0;
    for (t1) |v| total += v;
    for (t2) |v| total += v;
    try out.print("{d}\n", .{total});
    alloc.free(t2);
    alloc.free(t1);
}

Key points

  • Stack allocator: a stack allocator is an allocator that hands out blocks from one buffer in order and frees them last-in-first-out.
  • Each allocation records its start so freeing can rewind the buffer end.
  • Freeing a block that is not the newest is ignored.
  • Resizing works only for the newest block.
  • The buffer lives on the caller's stack, so no separate deinit step is needed.
  • A stack allocator suits short-lived scratch memory.

Two scratch blocks hold small numbers. The first block holds 2, 4, and 6. The second block holds 10 and 20. All five numbers add up to 42.

Try it Yourself

Edit the code, then press Run.

Running is simulated

Exercise

A stack allocator holds blocks A then B. The program frees A before B. What happens?