Generators in Zig: complete, simple, and no allocator needed.

gen.zig

const std = @import("std");
const assert = std.debug.assert;

pub fn IterOf(comptime Result: type) type {
    return struct {
        // Invariant: self._yield_frame != null or self._value == null
        _value: ?Result = null,
        _yield_frame: ?anyframe = undefined,

        const _the_result_type = Result;

        pub fn yield(self: *@This(), _value: Result) void {
            self._value = _value;
            self._yield_frame = @frame();
            suspend {}
            self._value = null;
            self._yield_frame = null;
        }

        pub fn next(self: *@This()) ?Result {
            if (self._yield_frame != null and self._value == null) {
                resume self._yield_frame.?;
            }
            defer self._value = null;
            return self._value;
        }
    };
}

/// Pinned! That is, an instance of this struct must never move to a different memory location.
pub fn GenIterState(comptime Runner: type) type {
    // find the result type from type Runner (inline comments show up in compiler error messages)
    const runner_struct_decls = @typeInfo(Runner).Struct.decls; // Runner must be a struct
    assert(runner_struct_decls.len == 1); // Runner must have exactly one declaration; later: allow non-function declarations
    const run_function_args = @typeInfo(runner_struct_decls[0].data.Fn.fn_type).Fn.args; // Runner must have a function declaration
    assert(run_function_args.len == 2); // the Runner function must have exactly two arguments
    const run_function_second_arg_pointer_child = @typeInfo(run_function_args[1].arg_type.?).Pointer.child; // Runner function second arg must be a pointer
    const Result = run_function_second_arg_pointer_child._the_result_type; // Runner function second arg must be *IterOf()

    return struct {
        _iter: IterOf(Result) = undefined,
        _frame: @Frame(Runner.run) = undefined,

        pub fn iterator(self: *@This(), runner: Runner) *IterOf(Result) {
            self._iter = IterOf(Result){};
            self._frame = async runner.run(&(self._iter));
            return &(self._iter);
        }
    };
}

// TESTS AND EXAMPLES
// ------------------

const expectEqual = std.testing.expectEqual;

const Bits = struct {
    pub fn run(_: *const @This(), iter: *IterOf(usize)) void {
        iter.yield(0);
        iter.yield(1);
    }
};

test "generate all bits, finite iterator" {
    var iter_state = GenIterState(Bits){};
    const iter = iter_state.iterator(Bits{});

    try expectEqual(@as(?usize, 0), iter.next());
    try expectEqual(@as(?usize, 1), iter.next());
    try expectEqual(@as(?usize, null), iter.next());
    try expectEqual(@as(?usize, null), iter.next());
}

const Nats = struct {
    below: usize,

    pub fn run(self: *const @This(), iter: *IterOf(usize)) void {
        var i: usize = 0;
        while (i < self.below) : (i += 1) {
            iter.yield(i);
        }
    }
};

test "generate all bits, bounded iterator" {
    var iter_state = GenIterState(Nats){};
    const iter = iter_state.iterator(Nats{ .below = 2 });

    try expectEqual(@as(?usize, 0), iter.next());
    try expectEqual(@as(?usize, 1), iter.next());
    try expectEqual(@as(?usize, null), iter.next());
    try expectEqual(@as(?usize, null), iter.next());
}