Compiler Explorer

Source code

const std = @import("std");

var resume_loc: anyframe = undefined;

// callconv() is necessary because of a compiler bug
fn generator(ctx: *GeneratorContext(u32), from: u32, to: u32) callconv(.Async) void
{
    var result : u32 = from;
    while(result < to) : (result += 1) {
        ctx.yield(result);
    }
}

pub fn main() void
{
    var gen = Generator(generator).init(.{
        .@"0" = @as(u32, 10),
        .@"1" = @as(u32, 20)
    });

var index : usize = 0;
    while(gen.next()) |value| {
        std.debug.warn("result[{}] = {}\n", .{ index, value });
        index += 1;
    }
}

fn GeneratorContext(comptime T: type) type {
    return struct {
        const Self = @This();

pub const Type = T;

buffer: T,
        frame: anyframe,

pub fn yield(self: *Self, result: T) void {
            
            self.buffer = result;

self.frame = @frame();
            suspend;
        }
    };
}

fn Generator(comptime function: var) type
{
    const Fun = @TypeOf(function);
    const info = @typeInfo(Fun).Fn;

const ArgTuple = TupleType(info.args[1..]);

return struct {
        const Self = @This();
        const Context = info.args[0].arg_type.?.Child;

context: Context,
        frame: @Frame(generatorWrapper),
        was_init: bool,
        is_finished: bool,
        args: ArgTuple,

fn generatorWrapper(self: *Self) callconv(.Async) void {
            await @call(
                .{ .modifier = .async_kw, },
                function,
                .{ &self.context } ++ self.args,
            );
            self.is_finished = true;
        }

pub fn init(args: ArgTuple) Self {
            return Self {
                .context = Context {
                    .buffer = undefined,
                    .frame = undefined,
                },
                .frame = undefined,
                .was_init = false,
                .is_finished = false,
                .args = args,
            };
        }

pub fn next(self: *Self) ?Context.Type {
            if(!self.is_finished) {
                if(self.was_init) {
                    resume self.context.frame;
                } else {
                    self.frame = async generatorWrapper(self);
                    self.was_init = true;
                }
            }
            return if(!self.is_finished)
                self.context.buffer
            else
                null;
        }
    };
}

fn MakeTuple1(comptime T0: type) type {
    const H = struct {
        var r0: T0 = undefined;
    };
    return @TypeOf(.{ H.r0 });
}

fn MakeTuple2(comptime T0: type, comptime T1: type) type {
    const H = struct {
        var r0: T0 = undefined;
        var r1: T1 = undefined;
    };
    return @TypeOf(.{ H.r0, H.r1 });
}

fn TupleType(comptime args: []std.builtin.TypeInfo.FnArg) type
{
    return switch(args.len) {
        0 => @TypeOf(.{}),
        1 => MakeTuple1(args[0].arg_type.?),
        2 => MakeTuple2(args[0].arg_type.?, args[1].arg_type.?),
        else => @compileError("Too many args for generator. Please use less!"),
    };
}