Skip to content

Zig type hacker and memory management

Joel Reymont
Joel Reymont
3 min read

You can map and fold in Zig but lack of closures makes it unergonomic. Here's a completely silly example that shows Zig type hackery as well as memory management.

const std = @import("std");
const mem = std.mem;
const testing = std.testing;

const Allocator = mem.Allocator;
pub const Error = Allocator.Error || error{};

pub const Expr = union(enum) {
    binary: Binary,
    int: usize,

    pub const Binary = struct {
        op: BinaryOp,
        lhs: *Expr,
        rhs: *Expr,
    };

    pub const BinaryOp = enum {
        @"+",
    };
};

pub fn exprWalker(T: type, E: type) type {
    return struct {
        context: T,
        walker: Walker,

        const Self = @This();

        pub const Walker = *const fn (*Self, *Expr) E!bool;

        pub fn run(self: *Self, node: *Expr) E!void {
            if (try self.walker(self, node))
                return;
            switch (node.*) {
                .binary => |binary| {
                    try self.run(binary.lhs);
                    try self.run(binary.rhs);
                },
                else => {},
            }
        }
    };
}

fn sumup(node: *Expr) Error!usize {
    const Context = usize;
    const Walker = exprWalker(Context, Error);
    const walk = struct {
        fn walk(self: *Walker, expr: *Expr) Error!bool {
            switch (expr.*) {
                .int => |i| {
                    self.context += i;
                    return true;
                },
                else => {},
            }
            return false;
        }
    }.walk;

    var walker = Walker{
        .context = 0,
        .walker = walk,
    };
    try walker.run(node);
    return walker.context;
}

pub fn clone(alloc: Allocator, value: anytype) Error!*@TypeOf(value) {
    const result = try alloc.create(@TypeOf(value));
    result.* = value;
    return result;
}

pub fn main() !void {
    const page_alloc = std.heap.page_allocator;
    var arena = std.heap.ArenaAllocator.init(page_alloc);
    defer arena.deinit();
    const alloc = arena.allocator();
    const lhs = Expr{ .int = 1 };
    const rhs = Expr{ .int = 2 };
    var expr = Expr{
        .binary = .{
            .op = .@"+",
            .lhs = try clone(alloc, lhs),
            .rhs = try clone(alloc, rhs),
        },
    };
    std.debug.print("sum = {d}\n", .{try sumup(&expr)});
}

Let's take it apart...

Recursive data structures require pointers so both left-hand side (lhs) and right-hand side (rhs) of the expression need to be pointers! Also, tagged unions (see union enum) are analogous to OCaml variants.

pub const Expr = union(enum) {
    binary: Binary,
    int: usize,

    pub const Binary = struct {
        op: BinaryOp,
        lhs: *Expr,
        rhs: *Expr,
    };

    pub const BinaryOp = enum {
        @"+",
    };
};

Memory allocation in Zig is very much explicit! I started my compiler project manually allocating and deallocating everything but quickly went with the arena allocator. This lets me free up all the data structures in one fell swoop once the compiler exits. Defer is awesome and will run when we exit the current scope. Think RAII in C++.

const page_alloc = std.heap.page_allocator;
var arena = std.heap.ArenaAllocator.init(page_alloc);
defer arena.deinit();
const alloc = arena.allocator();

This handy little function will make a copy of any value given to it by first allocating a chunk of memory the size of the given value and then assigning the value to the "dereferenced pointer".

Zig duck typing is great! Notice that I'm not specifying the type of value (anytype) and telling the Zig compiler to use the type of the value given ( TypeOf(value))

pub fn exprWalker(T: type, E: type) type {
    return struct {
        context: T,
        walker: Walker,

        const Self = @This();

        pub const Walker = *const fn (*Self, *Expr) E!bool;

        pub fn run(self: *Self, node: *Expr) E!void {
            if (try self.walker(self, node))
                return;
            switch (node.*) {
                .binary => |binary| {
                    try self.run(binary.lhs);
                    try self.run(binary.rhs);
                },
                else => {},
            }
        }
    };
}

Zig comptime generics at work here.

The exprWalker function returns a new type parameterized on the context and error types supplied to it.

The Walker type is a function that will do the heavy lifting. The expression walking machinery will abruptly stop if the function tells it that this expression has been handled. Otherwise, It will happily recurse into nested expressions.

Zig pattern matching using switch is not as advanced as in OCaml but perfectly serviceable. Finally, {} is an empty block.

fn sumup(node: *Expr) Error!usize {
    const Context = usize;
    const Walker = exprWalker(Context, Error);
    const walk = struct {
        fn walk(self: *Walker, expr: *Expr) Error!bool {
            switch (expr.*) {
                .int => |i| {
                    self.context += i;
                    return true;
                },
                else => {},
            }
            return false;
        }
    }.walk;

    var walker = Walker{
        .context = 0,
        .walker = walk,
    };
    try walker.run(node);
    return walker.context;
}

Here const Walker = exprWalker(Context, Error); instantiates a new Walker that's parameterized over the Context and Error types.

The walk function then looks for ints and sums them up, storing the sum in the context.

There are no closures in Zig but you can sort of kind of emulate them by returning a struct with a function inside it.

That's it!

Zig