🎓 A1Learner
🔍
0 of 46 lessons completed

Zig Custom Allocators

A custom allocator is an allocator you build yourself from the four vtable functions. It wraps another allocator to add behavior like counting or limits.

Counting allocations

This allocator wraps the general purpose allocator and counts bytes. It delegates every vtable call to the inner allocator.

const std = @import("std");
const Counter = struct {
    inner: std.mem.Allocator,
    bytes: usize = 0,
    allocs: usize = 0,
    fn allocFn(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        const ptr = self.inner.rawAlloc(len, alignment, ret_addr) orelse return null;
        self.bytes += len;
        self.allocs += 1;
        return ptr;
    }
    fn resizeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        return self.inner.rawResize(memory, alignment, new_len, ret_addr);
    }
    fn remapFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        return self.inner.rawRemap(memory, alignment, new_len, ret_addr);
    }
    fn freeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        self.bytes -= memory.len;
        self.inner.rawFree(memory, alignment, ret_addr);
    }
    fn allocator(self: *Counter) std.mem.Allocator {
        return .{ .ptr = self, .vtable = &.{ .alloc = allocFn, .resize = resizeFn, .remap = remapFn, .free = freeFn } };
    }
};
pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    var counter = Counter{ .inner = gpa.allocator() };
    const alloc = counter.allocator();
    const a = try alloc.alloc(u8, 10);
    defer alloc.free(a);
    const b = try alloc.alloc(u8, 5);
    alloc.free(b);
    try out.print("{d} {d} {d}\n", .{ counter.allocs, counter.bytes, a.len });
}

Growing a block

Realloc passes through resize and remap. The first byte survives the move to a larger block.

const std = @import("std");
const Counter = struct {
    inner: std.mem.Allocator,
    bytes: usize = 0,
    allocs: usize = 0,
    fn allocFn(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        const ptr = self.inner.rawAlloc(len, alignment, ret_addr) orelse return null;
        self.bytes += len;
        self.allocs += 1;
        return ptr;
    }
    fn resizeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        return self.inner.rawResize(memory, alignment, new_len, ret_addr);
    }
    fn remapFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        return self.inner.rawRemap(memory, alignment, new_len, ret_addr);
    }
    fn freeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        self.bytes -= memory.len;
        self.inner.rawFree(memory, alignment, ret_addr);
    }
    fn allocator(self: *Counter) std.mem.Allocator {
        return .{ .ptr = self, .vtable = &.{ .alloc = allocFn, .resize = resizeFn, .remap = remapFn, .free = freeFn } };
    }
};
pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    var counter = Counter{ .inner = gpa.allocator() };
    const alloc = counter.allocator();
    const a = try alloc.alloc(u8, 4);
    a[0] = 7;
    const b = try alloc.realloc(a, 8);
    try out.print("{d} {d}\n", .{ b.len, b[0] });
    alloc.free(b);
}

Tracking freed memory

The free function subtracts the freed length from the live total. After both frees, zero bytes remain live.

const std = @import("std");
const Counter = struct {
    inner: std.mem.Allocator,
    bytes: usize = 0,
    allocs: usize = 0,
    fn allocFn(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        const ptr = self.inner.rawAlloc(len, alignment, ret_addr) orelse return null;
        self.bytes += len;
        self.allocs += 1;
        return ptr;
    }
    fn resizeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        return self.inner.rawResize(memory, alignment, new_len, ret_addr);
    }
    fn remapFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        return self.inner.rawRemap(memory, alignment, new_len, ret_addr);
    }
    fn freeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        self.bytes -= memory.len;
        self.inner.rawFree(memory, alignment, ret_addr);
    }
    fn allocator(self: *Counter) std.mem.Allocator {
        return .{ .ptr = self, .vtable = &.{ .alloc = allocFn, .resize = resizeFn, .remap = remapFn, .free = freeFn } };
    }
};
pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    var counter = Counter{ .inner = gpa.allocator() };
    const alloc = counter.allocator();
    const a = try alloc.alloc(u8, 10);
    const b = try alloc.alloc(u8, 5);
    alloc.free(a);
    alloc.free(b);
    try out.print("{d} {d}\n", .{ counter.allocs, counter.bytes });
}

Example

const std = @import("std");
const Counter = struct {
    inner: std.mem.Allocator,
    bytes: usize = 0,
    allocs: usize = 0,
    fn allocFn(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        const ptr = self.inner.rawAlloc(len, alignment, ret_addr) orelse return null;
        self.bytes += len;
        self.allocs += 1;
        return ptr;
    }
    fn resizeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        return self.inner.rawResize(memory, alignment, new_len, ret_addr);
    }
    fn remapFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        return self.inner.rawRemap(memory, alignment, new_len, ret_addr);
    }
    fn freeFn(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
        const self: *Counter = @ptrCast(@alignCast(ctx));
        self.bytes -= memory.len;
        self.inner.rawFree(memory, alignment, ret_addr);
    }
    fn allocator(self: *Counter) std.mem.Allocator {
        return .{ .ptr = self, .vtable = &.{ .alloc = allocFn, .resize = resizeFn, .remap = remapFn, .free = freeFn } };
    }
};
pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    var counter = Counter{ .inner = gpa.allocator() };
    const alloc = counter.allocator();
    const nums = try alloc.alloc(u8, 5);
    for (0..5) |i| nums[i] = @intCast(i + 1);
    var sum: usize = 0;
    for (nums) |v| sum += v;
    alloc.free(nums);
    try out.print("{d} {d}\n", .{ sum, counter.bytes });
}

Key points

  • Vtable: a vtable is a small table of function pointers that tells an allocator how to allocate, resize, remap, and free.
  • The allocFn function receives the requested length and alignment and returns a pointer or null.
  • Returning null from allocFn makes the caller report OutOfMemory.
  • The resizeFn function tries to change a block size in place and reports success or failure.
  • The remapFn function may move a block to a new place when resizing in place is impossible.
  • A counting allocator wraps a second allocator and records each allocation it passes through.
  • The freeFn function releases a block and must use the same alignment as the allocation.

The counting allocator wraps the general purpose allocator. The program fills five bytes with the values 1 through 5. Their sum is 15. Freeing the block drops the live byte count to zero.

Try it Yourself

Edit the code, then press Run.

Running is simulated

Exercise

A custom allocator's allocFn returns null. What does the caller receive?