home

Closures in Zig

Sep 10, 2024

Zig doesn't have a simple way to create closures because closures require allocation, and Zig has a "no hidden allocation" policy. The values that you want to capture need to live somewhere (i.e. the heap), but there's no way for Zig (the language) to create the necessary heap memory. However, outside of the language, either in the standard library, or our own code, we can write code that takes an allocator and can thus create a closure.

To get our bearings, we'll start with a simple explicit example. We'll worry about making it more generic later. First, we'll create a job:

const Job = struct {
    id: usize,

    pub fn run(self: *Job, message: []const u8) void {
        std.debug.print("job {d} says: {s}\n", .{self.id, message});
    }
};

If we wanted to queue a job up, possibly to be run in a separate thread, we could easily create an array of pending jobs and store the job and message. We can start with something like:

const JobQueue = struct {
    len: usize = 0,
    pending: [5]Data = undefined,

    const Data = struct {
        job: *Job,
        message: []const u8,
    };

    pub fn add(self: *JobQueue, job: *Job, message: []const u8) void {
        // TODO make sure pending isn't full
        // Maybe we should be using a std.DoublyLinkedList instead
        self.pending[self.len] = .{.job = job, .message = message};
        self.len += 1;
    }

    pub fn run() {
      // TODO
    }
};

The above has some lifetime concerns: both job and message need to remain valid until the job runs. But, that isn't particular germane, so let's make our example a little less trivial.

We'll change the signature of Job.run to take an anytype rather than a []const u8:

const Job = struct {
    id: usize,

    pub fn run(self: *Job, message: anytype) void {
        std.debug.print("job {d} says: {any}\n", .{self.id, message});
    }
};

This introduces some issues with our JobQueue. The Data we were capturing was a job: *Job and a message: []const. What should it be now? We could try changing message to anytype:

const JobQueue = struct {
    len: usize = 0,
    pending: [5]Data = undefined,

    const Data = struct {
        job: *Job,
        message: anytype,
    };

    pub fn add(self: *JobQueue, job: *Job, message: anytype) void {
        self.pending[self.len] = .{.job = job, .message = message};
        self.len += 1;
    }
};

But this gives us a compile time error: error: expected type expression, found 'anytype'. Zig doesn't let us declare a field as anytype. To solve this we need to create a struct specific to the type of message being passed to add, which can then hold a message. Even to someone with reasonable experience in Zig, the following can be a little confusion:

pub fn add(self: *JobQueue, job: *Job, message: anytype) void {
    _ = self; // we'll use this later

    const Closure = struct {
        job: *Job,
        message: @TypeOf(message),
    };
    const closure = Closure{.job = job, .message = message};
    std.debug.print("{any}\n", .{closure.message});
}

This is a small but important step. Whenever a function has an anytype parameter, Zig creates a distinct implementation for every type used. It acts as a template. So if we call add with a []const u8 message, as before, we should really imagine the code expanding to:

pub fn add(self: *JobQueue, job: *Job, message: []const u8) void {
    _ = self; // we'll use this later

    const Closure = struct {
        job: *Job,
        message: []const u8,
    };
    const closure = Closure{.job = job, .message = message};
    std.debug.print("{any}\n", .{closure.message});
}

And if we were to call add with a u32 or a User we should imagine two additional methods being created, replacing the above []const u8 with a u32 and User respectively.

To me, it's strange to declare a struct within a function. When we see a function, we typically think of our running program executing the code within. But part of add is executed at runtime and part of it is executed at compile time (comptime). In Zig, types always have to be known at comptime, so we know that the creation of the Closure is happening at comptime. If you were to print it's name, via @typeName(Closure), you'd get something like sample.JobQueue.add__anon_1600.Closure. If we call add with three different message types, we'd get three different struct with three different names.

While the struct declaration happens at comptime, the rest of the code is just normal runtime code. We create an instance of a struct, assigning it a field, and [as a placeholder] print the field. Even if we created multiple versions of Closure (by calling add with different message types), when our code run, there's no ambiguity about which "Closure" we want, because normal name resolution takes place. In other words, we can further expand the above code, like Zig's compiler does, to:

const closure = sample.JobQueue.add__anon_1600.Closure{.job = job, .message = message};
    std.debug.print("{any}\n", .{closure.message});
}

We now have a variable closure which captures the data we need. But it might not be clear how that's useful. Our issue was that we didn't know what type pending should be:

const JobQueue = struct {
    len: usize = 0,
    pending: [5]??? = undefined;
    // ...
};

And that still isn't clear. It can't be Closure, because there isn't a single Closure type - there's a Closure for each type of message. To solve this, we need to introduce another type, an interface:

const Runnable = struct {
    ptr: *anyopaque;
    runFn: *const fn(ptr: *anyopaque) void,
};

An interface works by storing a typeless pointer (an *anyopque) along with the function(s) to execute. When the interface function is executed, the type is re-established allowing the real function to be called. In short, by erasing the type, we can store anything. Importantly, our interface has to store a pointer, i.e. it has to store an *anyopaque not an anyopaque (which you'll rarely, if ever, see). Why? because the size of everything has to be known at compile time, and the size of a pointer, no mater what it points to, is always usize. Thus, on a 64 bit platform, the size of Runnable is 16 bytes (two 64-bit pointers).

I say this is important because, at the top of this blog, I mentioned that our closure would require an allocation. We're ready to see that in action:

const JobQueue = struct {
    len: usize = 0,
    pending: [5]Runnable = undefined;

    pub fn add(self: *JobQueue, allocator: Allocator, job: *Job, message: anytype) !void {
        const Closure = struct {
            job: *Job,
            message: @TypeOf(message),
        };

        const closure = try allocator.create(Closure);
        closure.* = Closure{.job = job, .message = message};

        self.pending[self.len] = .{
            .ptr = closure,
            .runFn = undefined, //TODO
        };
        self.len += 1;
    }
};

const Runnable = struct {
    ptr: *anyopaque;
    runFn: *const fn(ptr: *anyopaque) void,
};

Our Closure instance, closure, automatically coerces to an *anyopaque, but only because we've turned into into a pointer by allocating memory on the heap for it. In other words, our Closure captures the values, but to be useful we need it to live somewhere beyond add's stack.

To complete our interface, we need to add a run method which will re-establish the type of the *anyopaque pointer. This isn't difficult, since the type of ptr is *Closure:

pub fn add(self: *JobQueue, allocator: Allocator, job: *Job, message: anytype) !void {
    const Closure = struct {
        job: *Job,
        message: @TypeOf(message),

        fn run(ptr: *anyopaque) void {
            var c: *@This() = @ptrCast(@alignCast(ptr));
            c.job.run(c.message);
        }
    };

    // ...
}

@This() returns the innermost struct (or union or enum) which, in this case, is our Closure. *@This() is a little weird, but it translate to *Closure, which is exactly what ptr is. We can't use Closure here because ...well, I'm not sure, but I guess the name resolution fails or is ambiguous due to how the structure is declared.

With our run method, we can almost finish our code. Previously, we had left the runFn of our Runnable as undefined. We now have something to set it to:

pub fn add(self: *JobQueue, allocator: Allocator, job: *Job, message: anytype) !void {

    // ...

    const closure = try allocator.create(Closure);
    closure.* = Closure{.job = job, .message = message};

    self.pending[self.len] = .{
        .ptr = closure,
        .runFn = Closure.run
    };
    self.len += 1;
}

The last thing that needs to be done is to free the memory we've allocated. Just like our closure captured (or closed oves) our job and message, we also have to capture the allocator. Here's a complete runnable example:

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

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const allocator = gpa.allocator();
    defer _  = gpa.detectLeaks();

    var queue = JobQueue{};

    var job = Job{.id = 1};
    try queue.add(allocator, &job, "hello");

    // without the @as, 123 is a comptime_int, which will not work
    try queue.add(allocator, &job, @as(u32, 123));

    queue.runOne();
    queue.runOne();
}

const Job = struct {
    id: usize,

    pub fn run(self: *Job, message: anytype) void {
        std.debug.print("job {d} says: {any}\n", .{self.id, message});
    }
};

const JobQueue = struct {
    len: usize = 0,
    pending: [5]Runnable = undefined,

    pub fn add(self: *JobQueue, allocator: Allocator, job: *Job, message: anytype) !void {
        const Closure = struct {
            job: *Job,
            allocator: Allocator,
            message: @TypeOf(message),

            fn run(ptr: *anyopaque) void {
                var c: *@This() = @ptrCast(@alignCast(ptr));
                defer c.allocator.destroy(c);
                c.job.run(c.message);
            }
        };

        const closure = try allocator.create(Closure);
        closure.* = Closure{.job = job, .allocator = allocator, .message = message};

        self.pending[self.len] = .{
            .ptr = closure,
            .runFn = Closure.run,
        };
        self.len += 1;
    }

    pub fn runOne(self: *JobQueue) void {
        const last = self.len - 1;
        const runnable = self.pending[last];
        runnable.runFn(runnable.ptr);
        self.len = last;
    }
};

const Runnable = struct {
    ptr: *anyopaque,
    runFn: *const fn(ptr: *anyopaque) void,
};

Conclusion

If you want to see a real example of this, checkout the spawn method of std.Thread.Pool. That implementation uses the @fieldParentPtr builtin, which makes it harder to follow, but which is also worth understanding. If you don't know what @fieldParentPtr does, you might want to first read Zig's @fieldParentPtr for dumbos like me.

Next time you're using a language with syntactical sugar for closures, you can visualize the auto-generated structure that the compiler creates and fields of an instance are set to the values you're capturing.