Skip to content

Instantly share code, notes, and snippets.

@adrusi
Last active March 10, 2024 15:09
Show Gist options
  • Save adrusi/0dfebca7ff79a8e9ff0c94259d89146d to your computer and use it in GitHub Desktop.
Save adrusi/0dfebca7ff79a8e9ff0c94259d89146d to your computer and use it in GitHub Desktop.
Generators in Zig 0.6.0
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;
};
}
};
}
}
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);
}
};
}
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 });
}
}
// 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;
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment