Generators in Zig 0.6.0

generator.zig

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

const builtin = @import("builtin");
const TypeInfo = builtin.TypeInfo;
const TypeId = builtin.TypeId;  

/// Iterator based on async functions. Equivalent to generators in Python,
/// Javascript, etc.
///
/// Create a generator type by passing a struct with a public "generate" method
/// to Generator. The method should take a self-pointer and a pointer to the
/// type of item the generator emits. It should return a void error union. To
/// yield an item, it should write the value to the memory pointed at by the
/// item pointer and suspend itself.
///
///     const SliceIter = Generator(struct {
///         slice: []i32,
///
///         pub fn generator(ctx: *@This(), item: *i32) !void {
///             for (ctx.slice) |x| {
///                 item.* = x;
///                 suspend;
///             }
///         }
///     });
///
/// A generator object can be created using the init method, passing an
/// allocator and a Gen.Args struct. This is a name given to the struct passed
/// to Generator. Then the caller can loop through items using the "next"
/// method.
///
///     var iter = SliceIter.init(.{
///         .slice = []i32 { 1, 2, 3, 4, 5 },
///     });
///     while (try iter.next()) |item| {
///         warn("{}\n", item);
///     }
pub fn Generator(comptime Ctx: type) type {
    comptime {
        const ctx_tinfo = @typeInfo(Ctx);

        debug.assert(ctx_tinfo == .Struct);
        debug.assert(@hasDecl(Ctx, "generate"));

        const gen_fn = Ctx.generate;
        const fn_tinfo = @typeInfo(@TypeOf(gen_fn));

        debug.assert(fn_tinfo == .Fn);
        debug.assert(fn_tinfo.Fn.args.len == 2);
        
        const arg1_tinfo = @typeInfo(fn_tinfo.Fn.args[0].arg_type.?);
        const arg2_tinfo = @typeInfo(fn_tinfo.Fn.args[1].arg_type.?);
        const ret_tinfo = @typeInfo(fn_tinfo.Fn.return_type.?);

        debug.assert(arg1_tinfo == .Pointer);
        debug.assert(arg1_tinfo.Pointer.child == Ctx);
        debug.assert(arg2_tinfo == .Pointer);
        debug.assert(ret_tinfo == .ErrorUnion);
        debug.assert(ret_tinfo.ErrorUnion.payload == void);

        const ExitStatus = ret_tinfo.ErrorUnion.error_set;
        const Item = arg2_tinfo.Pointer.child;

        return struct {
            pub const Args = Ctx;
            pub const Error = ExitStatus;
            pub const Item = Item;

            state: enum { Start, Alive, End },
            err: ?ExitStatus,
            gen_frame: @Frame(gen_fn),
            item: Item,
            ctx: Ctx,

            pub fn init(ctx: Ctx) @This() {
                return .{
                    .state = .Start,
                    .err = null,
                    .item = undefined,
                    .gen_frame = undefined,
                    .ctx = ctx,
                };
            }   

            pub fn next(self: *@This()) Error!?Item {
                switch (self.state) {
                    .Start => _ = async self.returnValueAwaiter(),
                    .Alive => resume self.gen_frame,
                    .End   => return null,
                }

                if (self.state == .End) {
                    if (self.err) |err| {
                        return err;
                    }
                    return null;
                }

                return self.item; // results in a copy
            }

            fn returnValueAwaiter(self: *@This()) void {
                self.state = .Alive;
                defer self.state = .End;

                self.gen_frame = async gen_fn(&self.ctx, &self.item);

                await self.gen_frame catch |err| {
                    self.err = err;
                };
            }
        };
    }
}

iterator.zig

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

const builtin = @import("builtin");
const vtable = @import("vtable.zig");

pub fn Iter(comptime ErrorType: type, comptime ItemType: type) type {
    return struct {
        const VTable = struct {
            pub const Impl = @OpaqueType();
            next: fn (*Impl) ErrorType!?ItemType,
        };

        vtable: *const VTable,
        impl: *VTable.Impl,

        pub fn init(iter: var) @This() {
            const T = @TypeOf(iter).Child;
            return .{
                .vtable = comptime vtable.populate(VTable, T, T),
                .impl = @ptrCast(*VTable.Impl, iter),
            };
        }

        pub fn next(self: @This()) ErrorType!?ItemType {
            return self.vtable.next(self.impl);
        }
    };
}

main.zig

const std = @import("std");

const generator = @import("generator.zig");
const Generator = generator.Generator;

const iterator = @import("iterator.zig");
const Iter = iterator.Iter;

pub fn ByteScanner(comptime InStream: type) type {
    return Generator(struct {
        in: *InStream,

        pub fn generate(self: *@This(), item: *u8) !void {
            var buffer: [4096]u8 = undefined;

            while (true) {
                var nbytes = try self.in.read(buffer[0..]);
                if (nbytes == 0) {
                    return;
                }

                for (buffer[0..nbytes]) |byte| {
                    item.* = byte;
                    suspend;
                }
            }
        }
    });
}

pub fn Utf8Scanner(comptime ByteSrcError: type) type {
    return Generator(struct {
        bytesrc: Iter(ByteSrcError, u8),

        pub fn generate(self: *@This(), item: *u21) !void {
            while (try self.bytesrc.next()) |byte| {
                var nextra: u2 = undefined;

                switch (byte >> 3) {
                    0b00000 ... 0b01111 => {
                        item.* = @intCast(u21, byte);
                        nextra = 0;
                    },
                    0b11000 ... 0b11011 => {
                        item.* = @intCast(u21, byte & 0b00011111);
                        nextra = 1;
                    },
                    0b11100 ... 0b11101 => {
                        item.* = @intCast(u21, byte & 0b00001111);
                        nextra = 2;
                    },
                    0b11110 => {
                        item.* = @intCast(u21, byte & 0b00000111);
                        nextra = 3;
                    },
                    else => { return error.EncodingError; },
                }

                var i: u2 = 0;
                while (i < nextra) {
                    var extra = (try self.bytesrc.next())
                        orelse return error.EncodingError;
                    if (extra & 0b11000000 != 0b10000000) {
                        return error.EncodingError;
                    }

                    item.* <<= 6;
                    item.* |= @intCast(u21, extra & 0b00111111);

                    i += 1;
                }

                suspend;
            }
        }
    });
}

pub fn main() !void {
    var in_file = std.io.getStdIn();
    var in_stream = &in_file.inStream();

    const FileScanner = ByteScanner(std.fs.File.InStream);
    var scanner = ByteScanner(std.fs.File.InStream).init(.{ .in = in_stream });

    var utf8 = Utf8Scanner(FileScanner.Error).init(.{ .bytesrc = Iter(FileScanner.Error, u8).init(&scanner) });

    while (try utf8.next()) |cp| {
        std.debug.warn("{x}\n", .{ cp });
    }
}

vtable.zig

// https://github.com/Hejsil/zig-interface
//
// MIT License
//
// Copyright (c) 2018
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

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

const TypeInfo = builtin.TypeInfo;

/// A function for populating vtables with their implementation.
///
/// VTable is the vtables type. It has to be a struct, which has only function fields.
///  It is also required that VTable contains a definition of Impl, which should be an
///  OpaqueType. A pointer to this type denotes the 'self' parameter of each method.
///
/// Functions is a namespace with all the functions that should populate the vtable.
///  Functions can contain functions not in VTable. These will just be ignored.
///  When populate populates the VTable with the functions from Functions, it does
///  type checking to ensure that the population is safe. If VTable has a field of
///  type 'fn(*Impl, u8) []u8', then populate will assert that Functions contain
///  a function of the same name, with the type 'fn(*T, u8) []u8'.
///
/// T is the self parameter of all the functions in Functions.
///
/// The result will be a pointer to a global VTable that can be shared between all
///  instantiations of T.
pub fn populate(comptime VTable: type, comptime Functions: type, comptime T: type) *const VTable {
    const GlobalStorage = struct {
        const vtable = blk: {
            const Impl = VTable.Impl;

            var res: VTable = undefined;
            inline for (@typeInfo(VTable).Struct.fields) |field| {
                const Fn = @TypeOf(@field(res, field.name));
                const Expect = @typeInfo(Fn).Fn;
                const Actual = @typeInfo(@TypeOf(@field(Functions, field.name))).Fn;
                debug.assert(!Expect.is_generic);
                debug.assert(!Expect.is_var_args);
                debug.assert(Expect.args.len > 0);
                debug.assert(Expect.calling_convention == Actual.calling_convention);
                debug.assert(Expect.is_generic == Actual.is_generic);
                debug.assert(Expect.is_var_args == Actual.is_var_args);
                debug.assert(Expect.return_type.? == Actual.return_type.?);
                debug.assert(Expect.args.len == Actual.args.len);

                for (Expect.args) |expect_arg, i| {
                    const actual_arg = Actual.args[i];
                    debug.assert(!expect_arg.is_generic);
                    debug.assert(expect_arg.is_generic == actual_arg.is_generic);
                    debug.assert(expect_arg.is_noalias == actual_arg.is_noalias);

                    // For the first arg. We enforce that it is a pointer, and
                    // that the actual function takes *T.
                    if (i == 0) {
                        const expect_ptr = @typeInfo(expect_arg.arg_type.?).Pointer;
                        const actual_ptr = @typeInfo(actual_arg.arg_type.?).Pointer;
                        debug.assert(expect_ptr.size == TypeInfo.Pointer.Size.One);
                        debug.assert(expect_ptr.size == actual_ptr.size);
                        debug.assert(expect_ptr.is_const == actual_ptr.is_const);
                        debug.assert(expect_ptr.is_volatile == actual_ptr.is_volatile);
                        debug.assert(actual_ptr.child == T);
                        debug.assert(expect_ptr.child == Impl);
                    } else {
                        debug.assert(expect_arg.arg_type.? == actual_arg.arg_type.?);
                    }
                }

                @field(res, field.name) = @ptrCast(Fn, @field(T, field.name));
            }

            break :blk res;
        };
    };

    return &GlobalStorage.vtable;
}