🎓 A1Learner
🔍
0 of 46 lessons completed

Zig Custom Allocator

A custom allocator is an allocator a program defines itself by filling in the four operations of the allocator interface. It wraps another allocator and adds its own behavior.

Count allocations and bytes

Two allocations come through the counter, and only the first is freed.

const std = @import("std");

const CountingAllocator = struct {
    child: std.mem.Allocator,
    allocations: usize,
    bytes: usize,
    freed: usize,
};

fn countingAllocator(state: *CountingAllocator) std.mem.Allocator {
    return .{
        .ptr = state,
        .vtable = &.{
            .alloc = countAlloc,
            .resize = countResize,
            .remap = countRemap,
            .free = countFree,
        },
    };
}

fn countAlloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    const result = self.child.rawAlloc(len, alignment, ret_addr) orelse return null;
    self.allocations += 1;
    self.bytes += len;
    return result;
}

fn countResize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    return self.child.rawResize(memory, alignment, new_len, ret_addr);
}

fn countRemap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    return self.child.rawRemap(memory, alignment, new_len, ret_addr);
}

fn countFree(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    self.freed += memory.len;
    self.child.rawFree(memory, alignment, ret_addr);
}

pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var state = CountingAllocator{
        .child = std.heap.page_allocator,
        .allocations = 0,
        .bytes = 0,
        .freed = 0,
    };
    const alloc = countingAllocator(&state);
    const a = try alloc.alloc(u8, 10);
    const b = try alloc.alloc(u8, 20);
    alloc.free(a);
    _ = b;
    try out.print("allocations={d} bytes={d} freed={d} live={d}\n", .{ state.allocations, state.bytes, state.freed, state.bytes - state.freed });
}

Count over a fixed buffer

The counter wraps a fixed buffer, and the failed request adds nothing to the counts.

const std = @import("std");

const CountingAllocator = struct {
    child: std.mem.Allocator,
    allocations: usize,
    bytes: usize,
    freed: usize,
};

fn countingAllocator(state: *CountingAllocator) std.mem.Allocator {
    return .{
        .ptr = state,
        .vtable = &.{
            .alloc = countAlloc,
            .resize = countResize,
            .remap = countRemap,
            .free = countFree,
        },
    };
}

fn countAlloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    const result = self.child.rawAlloc(len, alignment, ret_addr) orelse return null;
    self.allocations += 1;
    self.bytes += len;
    return result;
}

fn countResize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    return self.child.rawResize(memory, alignment, new_len, ret_addr);
}

fn countRemap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    return self.child.rawRemap(memory, alignment, new_len, ret_addr);
}

fn countFree(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    self.freed += memory.len;
    self.child.rawFree(memory, alignment, ret_addr);
}

pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var buf: [64]u8 = undefined;
    var fba = std.heap.FixedBufferAllocator.init(&buf);
    var state = CountingAllocator{
        .child = fba.allocator(),
        .allocations = 0,
        .bytes = 0,
        .freed = 0,
    };
    const alloc = countingAllocator(&state);
    const first = try alloc.alloc(u8, 32);
    try out.print("first={d}\n", .{first.len});
    const second = alloc.alloc(u8, 100) catch null;
    if (second == null) {
        try out.print("second failed\n", .{});
    }
    try out.print("allocations={d} bytes={d}\n", .{ state.allocations, state.bytes });
}

Track freed bytes

Both slices are freed, so the live count falls back to zero.

const std = @import("std");

const CountingAllocator = struct {
    child: std.mem.Allocator,
    allocations: usize,
    bytes: usize,
    freed: usize,
};

fn countingAllocator(state: *CountingAllocator) std.mem.Allocator {
    return .{
        .ptr = state,
        .vtable = &.{
            .alloc = countAlloc,
            .resize = countResize,
            .remap = countRemap,
            .free = countFree,
        },
    };
}

fn countAlloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    const result = self.child.rawAlloc(len, alignment, ret_addr) orelse return null;
    self.allocations += 1;
    self.bytes += len;
    return result;
}

fn countResize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    return self.child.rawResize(memory, alignment, new_len, ret_addr);
}

fn countRemap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    return self.child.rawRemap(memory, alignment, new_len, ret_addr);
}

fn countFree(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    self.freed += memory.len;
    self.child.rawFree(memory, alignment, ret_addr);
}

pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var state = CountingAllocator{
        .child = std.heap.page_allocator,
        .allocations = 0,
        .bytes = 0,
        .freed = 0,
    };
    const alloc = countingAllocator(&state);
    const a = try alloc.alloc(u8, 10);
    const b = try alloc.alloc(u8, 20);
    alloc.free(a);
    alloc.free(b);
    try out.print("freed={d} live={d}\n", .{ state.freed, state.bytes - state.freed });
}

Example

const std = @import("std");

const CountingAllocator = struct {
    child: std.mem.Allocator,
    allocations: usize,
    bytes: usize,
    freed: usize,
};

fn countingAllocator(state: *CountingAllocator) std.mem.Allocator {
    return .{
        .ptr = state,
        .vtable = &.{
            .alloc = countAlloc,
            .resize = countResize,
            .remap = countRemap,
            .free = countFree,
        },
    };
}

fn countAlloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    const result = self.child.rawAlloc(len, alignment, ret_addr) orelse return null;
    self.allocations += 1;
    self.bytes += len;
    return result;
}

fn countResize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    return self.child.rawResize(memory, alignment, new_len, ret_addr);
}

fn countRemap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    return self.child.rawRemap(memory, alignment, new_len, ret_addr);
}

fn countFree(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
    const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
    self.freed += memory.len;
    self.child.rawFree(memory, alignment, ret_addr);
}

pub fn main() !void {
    const out = std.io.getStdOut().writer();
    var state = CountingAllocator{
        .child = std.heap.page_allocator,
        .allocations = 0,
        .bytes = 0,
        .freed = 0,
    };
    const alloc = countingAllocator(&state);
    const a = try alloc.alloc(u8, 5);
    const b = try alloc.alloc(u8, 15);
    const c = try alloc.alloc(u8, 10);
    alloc.free(b);
    _ = a;
    _ = c;
    try out.print("allocations={d} bytes={d} freed={d} live={d}\n", .{ state.allocations, state.bytes, state.freed, state.bytes - state.freed });
}

Key points

  • Allocator interface: an allocator interface is the set of four operations every allocator provides: allocate, resize, remap, and free.
  • CountingAllocator: a counting allocator is an allocator that counts how many allocations it served and how many bytes they took.
  • The counting state is a struct that holds the wrapped allocator and the counts.
  • Vtable: a vtable is a table of four functions, one per operation, that Zig calls through the allocator value.
  • Each counting function updates a count and then passes the work to the wrapped allocator.
  • The allocate function returns null when the wrapped allocator fails, and Zig turns that null into error.OutOfMemory.
  • A custom allocator lets a program measure memory use without changing the code that allocates.

The workload allocates three slices and frees one. The report shows three allocations and thirty bytes. Fifteen bytes stay live, since two slices are never freed.

Try it Yourself

Edit the code, then press Run.

Running is simulated

Exercise

What does the counting allocator's allocate function do before returning memory?