├── .gitattributes ├── .github └── workflows │ └── zig.yml ├── .gitignore ├── LICENSE ├── README.md ├── benchmarks.zig ├── build.zig ├── src ├── _.zig ├── generator.zig ├── join.zig ├── lib.zig ├── map.zig └── tests.zig ├── zig.mod └── zigmod.lock /.gitattributes: -------------------------------------------------------------------------------- 1 | * text=auto 2 | *.zig text eol=lf 3 | zig.mod text eol=lf 4 | zigmod.* text eol=lf 5 | zig.mod linguist-language=YAML 6 | -------------------------------------------------------------------------------- /.github/workflows/zig.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | on: 3 | push: 4 | branches: [master] 5 | pull_request: 6 | branches: [master] 7 | 8 | jobs: 9 | test: 10 | strategy: 11 | fail-fast: false 12 | matrix: 13 | os: [ubuntu-latest, macos-latest, windows-latest] 14 | runs-on: ${{matrix.os}} 15 | steps: 16 | - uses: actions/checkout@v2 17 | - uses: goto-bus-stop/setup-zig@v1 18 | with: 19 | version: master 20 | 21 | - run: zig version 22 | - run: zig env 23 | - uses: yrashk/actions-setup-zigmod@v1_1 24 | with: 25 | version: 91 26 | - run: zigmod version 27 | - run: zigmod ci 28 | - run: zig build test 29 | lint: 30 | runs-on: ubuntu-latest 31 | steps: 32 | - uses: actions/checkout@v2 33 | - uses: goto-bus-stop/setup-zig@v1 34 | with: 35 | version: master 36 | - run: zig fmt --check src/*.zig 37 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.swp 2 | zig-* 3 | .zigmod 4 | deps.zig 5 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Yurii Rashkovskii 4 | 5 | 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: 6 | 7 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 8 | 9 | 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. 10 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Zig Generator 2 | 3 | This library provides a way to write generators in Zig. 4 | 5 | Features: 6 | 7 | * Supports async generator functions 8 | * Propagates generator errors 9 | * Return value capture 10 | 11 | ## Usage 12 | 13 | Here's an example of its basic usage: 14 | 15 | ```zig 16 | const std = @import("std"); 17 | const gen = @import("generator"); 18 | 19 | const Ty = struct { 20 | pub fn generate(_: *@This(), handle: *gen.Handle(u8)) !u8 { 21 | try handle.yield(0); 22 | try handle.yield(1); 23 | try handle.yield(2); 24 | return 3; 25 | } 26 | }; 27 | 28 | const G = gen.Generator(Ty, u8); 29 | 30 | pub const io_mode = .evented; 31 | 32 | pub fn main() !void { 33 | var g = G.init(Ty{}); 34 | 35 | std.debug.assert((try g.next()).? == 0); 36 | std.debug.assert((try g.next()).? == 1); 37 | std.debug.assert((try g.next()).? == 2); 38 | std.debug.assert((try g.next()) == null); 39 | std.debug.assert(g.state.Returned == 3); 40 | } 41 | ``` 42 | 43 | -------------------------------------------------------------------------------- /benchmarks.zig: -------------------------------------------------------------------------------- 1 | const std = @import("std"); 2 | 3 | const generator = @import("generator"); 4 | const Handle = generator.Handle; 5 | const Generator = generator.Generator; 6 | const Map = generator.Map; 7 | 8 | pub const io_mode = .evented; 9 | 10 | const expect = std.testing.expect; 11 | 12 | pub fn main() !void { 13 | try benchDrain(); 14 | try benchGeneratorVsCallback(); 15 | } 16 | 17 | pub fn benchDrain() !void { 18 | std.debug.print("\n=== Benchmark: generator draining\n", .{}); 19 | 20 | const ty = struct { 21 | n: usize, 22 | pub fn generate(self: *@This(), handle: *Handle(u8)) !u8 { 23 | while (self.n > 0) { 24 | try handle.yield(1); 25 | self.n -= 1; 26 | } 27 | return 3; 28 | } 29 | }; 30 | 31 | // ensure they behave as expected 32 | const G = Generator(ty, u8); 33 | var g = G.init(ty{ .n = 1 }); 34 | 35 | try g.drain(); 36 | try expect(g.state.Returned == 3); 37 | 38 | // measure performance 39 | 40 | const bench = @import("bench"); 41 | try bench.benchmark(struct { 42 | pub const args = [_]usize{ 1, 2, 3, 5, 10, 100 }; 43 | pub const arg_names = [_][]const u8{ "1", "2", "3", "5", "10", "100" }; 44 | 45 | pub fn return_value(n: usize) !void { 46 | var gen = G.init(ty{ .n = n }); 47 | try gen.drain(); 48 | try expect(gen.state.Returned == 3); 49 | } 50 | }); 51 | } 52 | 53 | pub fn benchGeneratorVsCallback() !void { 54 | const W = fn (u8) callconv(.Async) anyerror!void; 55 | 56 | const busy_work = struct { 57 | fn do(_: u8) callconv(.Async) !void { 58 | std.os.nanosleep(0, 10); 59 | } 60 | }; 61 | 62 | const no_work = struct { 63 | fn do(_: u8) callconv(.Async) !void {} 64 | }; 65 | _ = no_work; 66 | 67 | std.debug.print("\n=== Benchmark: generator vs callback\n", .{}); 68 | 69 | const ty = struct { 70 | pub fn generate(_: *@This(), handle: *Handle(u8)) !u8 { 71 | try handle.yield(0); 72 | try handle.yield(1); 73 | try handle.yield(2); 74 | return 3; 75 | } 76 | }; 77 | 78 | const tyc = struct { 79 | pub fn run(_: *@This(), cb: fn (u8) callconv(.Async) anyerror!void) !u8 { 80 | var frame_buffer: [64]u8 align(@alignOf(@Frame(busy_work.do))) = undefined; 81 | var result: anyerror!void = undefined; 82 | suspend { 83 | resume @frame(); 84 | } 85 | try await @asyncCall(&frame_buffer, &result, cb, .{0}); 86 | suspend { 87 | resume @frame(); 88 | } 89 | try await @asyncCall(&frame_buffer, &result, cb, .{1}); 90 | suspend { 91 | resume @frame(); 92 | } 93 | try await @asyncCall(&frame_buffer, &result, cb, .{2}); 94 | return 3; 95 | } 96 | }; 97 | 98 | // ensure they behave as expected 99 | const G = Generator(ty, u8); 100 | var g = G.init(ty{}); 101 | 102 | try g.drain(); 103 | try expect(g.state.Returned == 3); 104 | 105 | // measure performance 106 | 107 | const bench = @import("bench"); 108 | try bench.benchmark(struct { 109 | pub const args = [_]W{ 110 | no_work.do, 111 | busy_work.do, 112 | }; 113 | 114 | pub const arg_names = [_][]const u8{ 115 | "no work", 116 | "busy work", 117 | }; 118 | 119 | pub fn generator(w: W) !void { 120 | var gen = G.init(ty{}); 121 | var frame_buffer: [64]u8 align(@alignOf(@Frame(busy_work.do))) = undefined; 122 | var result: anyerror!void = undefined; 123 | while (try gen.next()) |v| { 124 | try await @asyncCall(&frame_buffer, &result, w, .{v}); 125 | } 126 | try expect(gen.state.Returned == 3); 127 | } 128 | 129 | pub fn callback(w: W) !void { 130 | var c = tyc{}; 131 | try expect((try await async c.run(w)) == 3); 132 | } 133 | }); 134 | } 135 | -------------------------------------------------------------------------------- /build.zig: -------------------------------------------------------------------------------- 1 | const deps = @import("./deps.zig"); 2 | const std = @import("std"); 3 | 4 | pub fn build(b: *std.build.Builder) void { 5 | // Standard release options allow the person running `zig build` to select 6 | // between Debug, ReleaseSafe, ReleaseFast, and ReleaseSmall. 7 | const mode = b.standardReleaseOptions(); 8 | 9 | const lib = b.addStaticLibrary("zig-generator", "src/lib.zig"); 10 | lib.setBuildMode(mode); 11 | lib.install(); 12 | 13 | const tests = b.addTest("src/tests.zig"); 14 | tests.test_evented_io = true; 15 | tests.setBuildMode(mode); 16 | 17 | const test_step = b.step("test", "Run library tests"); 18 | test_step.dependOn(&tests.step); 19 | 20 | const benchmarks = b.addExecutable("zig-generator-benchmarks", "benchmarks.zig"); 21 | benchmarks.setBuildMode(.ReleaseFast); 22 | benchmarks.install(); 23 | 24 | const run_benchmarks = benchmarks.run(); 25 | run_benchmarks.step.dependOn(b.getInstallStep()); 26 | 27 | const benchmarks_step = b.step("bench", "Run benchmarks"); 28 | benchmarks_step.dependOn(&run_benchmarks.step); 29 | 30 | deps.addAllTo(lib); 31 | deps.addAllTo(tests); 32 | deps.addAllTo(benchmarks); 33 | } 34 | -------------------------------------------------------------------------------- /src/_.zig: -------------------------------------------------------------------------------- 1 | const std = @import("std"); 2 | 3 | /// Internal use: 4 | /// 5 | /// Finds and extracts any function in a given struct `T` 6 | /// that has a type `Sig` 7 | pub fn extractFn(comptime T: type, comptime Sig: type) Sig { 8 | switch (@typeInfo(T)) { 9 | .Struct => { 10 | const decls = std.meta.declarations(T); 11 | inline for (decls) |decl| { 12 | if (decl.is_pub) { 13 | switch (decl.data) { 14 | .Fn => |fn_decl| { 15 | if (fn_decl.fn_type == Sig) { 16 | return @field(T, decl.name); 17 | } 18 | }, 19 | else => {}, 20 | } 21 | } 22 | } 23 | @compileError("no public functions found"); 24 | }, 25 | else => @compileError("only structs are allowed"), 26 | } 27 | } 28 | 29 | test "extractFn" { 30 | const f = extractFn(struct { 31 | pub fn check(_: u8, _: u8) void {} 32 | pub fn add(a: u8, b: u8) u8 { 33 | return a + b; 34 | } 35 | }, fn (u8, u8) u8); 36 | try std.testing.expect(f(1, 2) == 3); 37 | } 38 | -------------------------------------------------------------------------------- /src/generator.zig: -------------------------------------------------------------------------------- 1 | const std = @import("std"); 2 | const assert = std.debug.assert; 3 | 4 | pub const Cancellation = error{ 5 | GeneratorCancelled, 6 | }; 7 | 8 | pub const State = enum { Initialized, Started, Error, Returned, Cancelled }; 9 | 10 | /// Generator handle, to be used in Handle's Ctx type 11 | /// 12 | /// `T` is the type that the generator yields 13 | /// `Return` is generator's return type 14 | pub fn Handle(comptime T: type) type { 15 | return struct { 16 | const Self = @This(); 17 | 18 | const HandleState = enum { Working, Yielded, Cancel }; 19 | 20 | const Suspension = enum(u8) { Unsuspended, Suspended, Yielded }; 21 | 22 | frame: *@Frame(yield) = undefined, 23 | gen_frame: anyframe = undefined, 24 | gen_frame_suspended: std.atomic.Atomic(Suspension) = std.atomic.Atomic(Suspension).init(.Unsuspended), 25 | 26 | state: union(HandleState) { 27 | Working: void, 28 | Yielded: T, 29 | Cancel: void, 30 | } = .Working, 31 | 32 | /// Yields a value 33 | pub fn yield(self: *Self, t: T) error{GeneratorCancelled}!void { 34 | if (self.state == .Cancel) return error.GeneratorCancelled; 35 | 36 | suspend { 37 | self.state = .{ .Yielded = t }; 38 | self.frame = @frame(); 39 | if (self.gen_frame_suspended.swap(.Yielded, .SeqCst) == .Suspended) { 40 | resume self.gen_frame; 41 | } 42 | } 43 | if (self.state == .Cancel) return error.GeneratorCancelled; 44 | self.state = .Working; 45 | } 46 | }; 47 | } 48 | 49 | /// Generator type allows an async function to yield multiple 50 | /// values, and return an error or a result. 51 | /// 52 | /// Ctx type must be a struct and it must have the following function: 53 | /// 54 | /// * `generate(self: *@This(), handle: *generator.Handle(T)) !Return` 55 | /// 56 | /// where `T` is the type of value yielded by the generator and `Return` is 57 | /// the type of the return value. 58 | /// 59 | /// NOTE: In many cases it may be advisable to have `T` be a pointer to a type, 60 | /// particularly if the the yielded type is larger than a machine word. 61 | /// This will eliminate the unnecessary copying of the value and may have a positive 62 | /// impact on performance. 63 | /// This is also a critical consideration if the generator needs to be able to 64 | /// observe changes that occurred to the value during suspension. 65 | pub fn Generator(comptime Ctx: type, comptime T: type) type { 66 | const ty = @typeInfo(Ctx); 67 | comptime { 68 | assert(ty == .Struct); 69 | assert(@hasDecl(Ctx, "generate")); 70 | } 71 | 72 | const generate_fn = Ctx.generate; 73 | const generate_fn_info = @typeInfo(@TypeOf(generate_fn)); 74 | 75 | assert(generate_fn_info == .Fn); 76 | assert(generate_fn_info.Fn.args.len == 2); 77 | 78 | const arg1_tinfo = @typeInfo(generate_fn_info.Fn.args[0].arg_type.?); 79 | const arg2_tinfo = @typeInfo(generate_fn_info.Fn.args[1].arg_type.?); 80 | const ret_tinfo = @typeInfo(generate_fn_info.Fn.return_type.?); 81 | 82 | // context 83 | assert(arg1_tinfo == .Pointer); 84 | assert(arg1_tinfo.Pointer.child == Ctx); 85 | 86 | // Handle 87 | assert(arg2_tinfo == .Pointer); 88 | assert(arg2_tinfo.Pointer.child == Handle(T)); 89 | 90 | assert(ret_tinfo == .ErrorUnion); 91 | 92 | return struct { 93 | const Self = @This(); 94 | const Err = ret_tinfo.ErrorUnion.error_set; 95 | const CompleteErrorSet = Err || Cancellation; 96 | 97 | pub const Return = ret_tinfo.ErrorUnion.payload; 98 | pub const Context = Ctx; 99 | pub const GeneratorState = union(State) { 100 | Initialized: void, 101 | Started: void, 102 | Error: Err, 103 | Returned: Return, 104 | Cancelled: void, 105 | }; 106 | 107 | handle: Handle(T) = Handle(T){}, 108 | 109 | /// Generator's state 110 | /// 111 | /// * `.Initialized` -- it hasn't been started yet 112 | /// * `.Started` -- it has been started 113 | /// * `.Returned` -- it has returned a value 114 | /// * `.Error` -- it has returned an error 115 | /// * `.Cancelled` -- it has been cancelled 116 | state: GeneratorState = .Initialized, 117 | 118 | /// Generator's own structure 119 | context: Context, 120 | 121 | generator_frame: @Frame(generator) = undefined, 122 | 123 | /// Initializes a generator 124 | pub fn init(ctx: Ctx) Self { 125 | return Self{ 126 | .context = ctx, 127 | }; 128 | } 129 | 130 | fn generator(self: *Self) void { 131 | if (generate_fn(&self.context, &self.handle)) |val| { 132 | self.state = .{ .Returned = val }; 133 | } else |err| { 134 | switch (err) { 135 | error.GeneratorCancelled => { 136 | self.state = .Cancelled; 137 | }, 138 | else => |e| { 139 | self.state = .{ .Error = e }; 140 | }, 141 | } 142 | } 143 | if (self.handle.gen_frame_suspended.swap(.Unsuspended, .SeqCst) == .Suspended) { 144 | resume self.handle.gen_frame; 145 | } 146 | 147 | suspend {} 148 | unreachable; 149 | } 150 | 151 | /// Returns the next yielded value, or `null` if the generator returned or was cancelled. 152 | /// `next()` propagates errors returned by the generator function. 153 | /// 154 | /// .state.Returned union variant can be used to retrieve the return value of the generator 155 | /// .state.Cancelled indicates that the generator was cancelled 156 | /// .state.Error union variant can be used to retrieve the error 157 | /// 158 | pub fn next(self: *Self) Err!?T { 159 | switch (self.state) { 160 | .Initialized => { 161 | self.state = .Started; 162 | self.handle.gen_frame = @frame(); 163 | self.generator_frame = async self.generator(); 164 | }, 165 | .Started => { 166 | resume self.handle.frame; 167 | }, 168 | else => return null, 169 | } 170 | 171 | while (self.state == .Started and self.handle.state == .Working) { 172 | suspend { 173 | if (self.handle.gen_frame_suspended.swap(.Suspended, .SeqCst) == .Yielded) { 174 | resume @frame(); 175 | } 176 | } 177 | } 178 | self.handle.gen_frame_suspended.store(.Unsuspended, .SeqCst); 179 | 180 | switch (self.state) { 181 | .Started => { 182 | return self.handle.state.Yielded; 183 | }, 184 | .Error => |e| return e, 185 | else => return null, 186 | } 187 | } 188 | 189 | /// Drains the generator until it is done 190 | pub fn drain(self: *Self) !void { 191 | while (try self.next()) |_| {} 192 | } 193 | 194 | /// Cancels the generator 195 | /// 196 | /// It won't yield any more values and will run its deferred code. 197 | /// 198 | /// However, it may still continue working until it attempts to yield. 199 | /// This is possible if the generator is an async function using other async functions. 200 | /// 201 | /// NOTE that the generator must cooperate (or at least, not get in the way) with its cancellation. 202 | /// An uncooperative generator can catch `GeneratorCancelled` error and refuse to be terminated. 203 | /// In such case, the generator will be effectively drained upon an attempt to cancel it. 204 | pub fn cancel(self: *Self) void { 205 | self.handle.state = .Cancel; 206 | } 207 | }; 208 | } 209 | 210 | test "basic generator" { 211 | const expect = std.testing.expect; 212 | const ty = struct { 213 | pub fn generate(_: *@This(), handle: *Handle(u8)) !void { 214 | try handle.yield(0); 215 | try handle.yield(1); 216 | try handle.yield(2); 217 | } 218 | }; 219 | const G = Generator(ty, u8); 220 | var g = G.init(ty{}); 221 | 222 | try expect((try g.next()).? == 0); 223 | try expect((try g.next()).? == 1); 224 | try expect((try g.next()).? == 2); 225 | try expect((try g.next()) == null); 226 | try expect(g.state == .Returned); 227 | try expect((try g.next()) == null); 228 | } 229 | 230 | test "generator with async i/o" { 231 | const expect = std.testing.expect; 232 | const ty = struct { 233 | pub fn generate(_: *@This(), handle: *Handle(u8)) !void { 234 | const file = try std.fs.cwd() 235 | .openFile("README.md", std.fs.File.OpenFlags{ .read = true, .write = false }); 236 | const reader = file.reader(); 237 | 238 | while (true) { 239 | const byte = reader.readByte() catch return; 240 | try handle.yield(byte); 241 | } 242 | } 243 | }; 244 | const G = Generator(ty, u8); 245 | var g = G.init(ty{}); 246 | 247 | var bytes: usize = 0; 248 | 249 | while (try g.next()) |_| { 250 | bytes += 1; 251 | } 252 | 253 | try expect(bytes > 0); 254 | } 255 | 256 | test "generator with async await" { 257 | const expect = std.testing.expect; 258 | const ty = struct { 259 | fn doAsync() callconv(.Async) u8 { 260 | suspend { 261 | resume @frame(); 262 | } 263 | return 1; 264 | } 265 | 266 | pub fn generate(_: *@This(), handle: *Handle(u8)) !void { 267 | try handle.yield(await async doAsync()); 268 | } 269 | }; 270 | const G = Generator(ty, u8); 271 | var g = G.init(ty{}); 272 | 273 | try expect((try g.next()).? == 1); 274 | } 275 | 276 | test "context" { 277 | const expect = std.testing.expect; 278 | 279 | const ty = struct { 280 | a: u8 = 1, 281 | 282 | pub fn generate(_: *@This(), handle: *Handle(u8)) !void { 283 | try handle.yield(0); 284 | try handle.yield(1); 285 | try handle.yield(2); 286 | } 287 | }; 288 | 289 | const G = Generator(ty, u8); 290 | var g = G.init(ty{}); 291 | 292 | try expect(g.context.a == 1); 293 | } 294 | 295 | test "errors in generators" { 296 | const expect = std.testing.expect; 297 | const ty = struct { 298 | pub fn generate(_: *@This(), handle: *Handle(u8)) !void { 299 | try handle.yield(0); 300 | try handle.yield(1); 301 | return error.SomeError; 302 | } 303 | }; 304 | const G = Generator(ty, u8); 305 | var g = G.init(ty{}); 306 | 307 | try expect((try g.next()).? == 0); 308 | try expect((try g.next()).? == 1); 309 | _ = g.next() catch |err| { 310 | try expect(g.state == .Error); 311 | try expect((try g.next()) == null); 312 | switch (err) { 313 | error.SomeError => { 314 | return; 315 | }, 316 | else => { 317 | @panic("incorrect error has been captured"); 318 | }, 319 | } 320 | return; 321 | }; 322 | @panic("error should have been generated"); 323 | } 324 | 325 | test "return value in generator" { 326 | const expect = std.testing.expect; 327 | const ty = struct { 328 | complete: bool = false, 329 | pub fn generate(self: *@This(), handle: *Handle(u8)) !u8 { 330 | defer { 331 | self.complete = true; 332 | } 333 | try handle.yield(0); 334 | try handle.yield(1); 335 | try handle.yield(2); 336 | return 3; 337 | } 338 | }; 339 | const G = Generator(ty, u8); 340 | var g = G.init(ty{}); 341 | 342 | try expect((try g.next()).? == 0); 343 | try expect((try g.next()).? == 1); 344 | try expect((try g.next()).? == 2); 345 | try expect((try g.next()) == null); 346 | try expect(g.state == .Returned); 347 | try expect(g.state.Returned == 3); 348 | try expect(g.context.complete); 349 | } 350 | 351 | test "drain" { 352 | const expect = std.testing.expect; 353 | const ty = struct { 354 | pub fn generate(_: *@This(), handle: *Handle(u8)) !u8 { 355 | try handle.yield(0); 356 | try handle.yield(1); 357 | try handle.yield(2); 358 | return 3; 359 | } 360 | }; 361 | const G = Generator(ty, u8); 362 | var g = G.init(ty{}); 363 | 364 | try g.drain(); 365 | try expect(g.state.Returned == 3); 366 | } 367 | 368 | test "cancel" { 369 | const expect = std.testing.expect; 370 | const ty = struct { 371 | drained: bool = false, 372 | cancelled: bool = false, 373 | 374 | pub fn generate(self: *@This(), handle: *Handle(u8)) !u8 { 375 | errdefer |e| { 376 | if (e == error.GeneratorCancelled) { 377 | self.cancelled = true; 378 | } 379 | } 380 | try handle.yield(0); 381 | try handle.yield(1); 382 | try handle.yield(2); 383 | self.drained = true; 384 | return 3; 385 | } 386 | }; 387 | 388 | const G = Generator(ty, u8); 389 | 390 | // cancel before yielding 391 | var g = G.init(ty{}); 392 | g.cancel(); 393 | try expect((try g.next()) == null); 394 | try expect(g.state == .Cancelled); 395 | try expect(!g.context.drained); 396 | try expect(g.context.cancelled); 397 | 398 | // cancel after yielding 399 | g = G.init(ty{}); 400 | try expect((try g.next()).? == 0); 401 | g.cancel(); 402 | try expect((try g.next()) == null); 403 | try expect(g.state == .Cancelled); 404 | try expect(!g.context.drained); 405 | try expect(g.context.cancelled); 406 | } 407 | 408 | test "uncooperative cancellation" { 409 | const expect = std.testing.expect; 410 | const ty = struct { 411 | drained: bool = false, 412 | ignored_termination_0: bool = false, 413 | ignored_termination_1: bool = false, 414 | 415 | pub fn generate(self: *@This(), handle: *Handle(u8)) !void { 416 | handle.yield(0) catch { 417 | self.ignored_termination_0 = true; 418 | }; 419 | handle.yield(1) catch { 420 | self.ignored_termination_1 = true; 421 | }; 422 | self.drained = true; 423 | } 424 | }; 425 | 426 | const G = Generator(ty, u8); 427 | 428 | // Cancel before yielding 429 | var g = G.init(ty{}); 430 | g.cancel(); 431 | try expect((try g.next()) == null); 432 | try expect(g.state == .Returned); 433 | try expect(g.context.drained); 434 | try expect(g.context.ignored_termination_0); 435 | try expect(g.context.ignored_termination_1); 436 | 437 | // Cancel after yielding 438 | g = G.init(ty{}); 439 | try expect((try g.next()).? == 0); 440 | g.cancel(); 441 | try expect((try g.next()) == null); 442 | try expect(g.state == .Returned); 443 | try expect(g.context.drained); 444 | try expect(g.context.ignored_termination_0); 445 | try expect(g.context.ignored_termination_1); 446 | } 447 | -------------------------------------------------------------------------------- /src/join.zig: -------------------------------------------------------------------------------- 1 | const std = @import("std"); 2 | const generator = @import("./generator.zig"); 3 | 4 | // pending resolution of https://github.com/ziglang/zig/issues/10442, 5 | // this has to be a function separate from `Join` 6 | fn joinedGenerator(comptime g: type, comptime T: type, comptime allocating: bool) type { 7 | return struct { 8 | generator: g, 9 | state: enum { Next, Awaiting, Returned, Done } = .Next, 10 | frame: if (allocating) *@Frame(next) else @Frame(next) = undefined, 11 | fn next(self: *@This(), counter: *std.atomic.Atomic(usize), frame: anyframe) !?T { 12 | defer { 13 | self.state = .Returned; 14 | if (counter.fetchAdd(1, .SeqCst) == 0) { 15 | resume frame; 16 | } 17 | } 18 | return self.generator.next(); 19 | } 20 | }; 21 | } 22 | 23 | fn initializer( 24 | comptime Self: type, 25 | comptime generators: []const type, 26 | generator_fields: []const std.builtin.TypeInfo.StructField, 27 | comptime allocating: bool, 28 | ) type { 29 | return if (allocating) struct { 30 | pub fn init(g: std.meta.Tuple(generators), allocator: std.mem.Allocator) Self { 31 | var s = Self{ .allocator = allocator }; 32 | inline for (generator_fields) |_, i| { 33 | s.generators[i] = .{ .generator = g[i] }; 34 | } 35 | return s; 36 | } 37 | } else struct { 38 | pub fn init(g: std.meta.Tuple(generators)) Self { 39 | var s = Self{}; 40 | inline for (generator_fields) |_, i| { 41 | s.generators[i] = .{ .generator = g[i] }; 42 | } 43 | return s; 44 | } 45 | }; 46 | } 47 | 48 | /// Joins multiple generators into one and yields values as they come from 49 | /// either generator 50 | pub fn Join(comptime generators: []const type, comptime T: type, comptime allocating: bool) type { 51 | var generator_fields: [generators.len]std.builtin.TypeInfo.StructField = undefined; 52 | inline for (generators) |g, field_index| { 53 | const G = joinedGenerator(g, T, allocating); 54 | generator_fields[field_index] = .{ 55 | .name = std.fmt.comptimePrint("{d}", .{field_index}), 56 | .field_type = G, 57 | .default_value = @as(?G, null), 58 | .is_comptime = false, 59 | .alignment = @alignOf(G), 60 | }; 61 | } 62 | const generators_struct = std.builtin.TypeInfo{ 63 | .Struct = .{ 64 | .layout = .Auto, 65 | .fields = &generator_fields, 66 | .decls = &[0]std.builtin.TypeInfo.Declaration{}, 67 | .is_tuple = true, 68 | }, 69 | }; 70 | 71 | const generator_fields_const = generator_fields; 72 | 73 | return generator.Generator(struct { 74 | const Self = @This(); 75 | 76 | generators: @Type(generators_struct) = undefined, 77 | frame: *@Frame(generate) = undefined, 78 | allocator: if (allocating) std.mem.Allocator else void = undefined, 79 | 80 | pub usingnamespace initializer(Self, generators, &generator_fields_const, allocating); 81 | 82 | pub fn generate(self: *Self, handle: *generator.Handle(T)) !void { 83 | if (allocating) { 84 | inline for (generator_fields_const) |_, i| { 85 | var g = &self.generators[i]; 86 | g.frame = self.allocator.create(@Frame(@TypeOf(g.*).next)) catch |e| { 87 | @setEvalBranchQuota(generators.len * 1000); 88 | inline for (generator_fields_const) |_, i_| { 89 | if (i_ == i) return e; 90 | var g_ = &self.generators[i_]; 91 | self.allocator.destroy(g_.frame); 92 | } 93 | }; 94 | } 95 | } 96 | 97 | defer { 98 | if (allocating) { 99 | inline for (generator_fields_const) |_, i| { 100 | var g = &self.generators[i]; 101 | if (g.state != .Done) 102 | self.allocator.destroy(g.frame); 103 | } 104 | } 105 | } 106 | 107 | var counter = std.atomic.Atomic(usize).init(0); 108 | var active: usize = self.generators.len; 109 | var reported: usize = 0; 110 | while (true) { 111 | // If there are no new reports, suspend until resumed by one 112 | suspend { 113 | if (counter.swap(0, .SeqCst) == reported) { 114 | // run next() where needed 115 | inline for (generator_fields_const) |_, i| { 116 | var g = &self.generators[i]; 117 | if (g.state == .Next) { 118 | g.state = .Awaiting; 119 | if (allocating) 120 | g.frame.* = async g.next(&counter, @frame()) 121 | else 122 | g.frame = async g.next(&counter, @frame()); 123 | } 124 | } 125 | } else { 126 | reported = 0; 127 | resume @frame(); 128 | } 129 | } 130 | reported = counter.load(.SeqCst); 131 | 132 | while (true) { 133 | // check for returns 134 | var yielded: usize = 0; 135 | inline for (generator_fields_const) |_, i| { 136 | var g = &self.generators[i]; 137 | if (g.state == .Returned) { 138 | yielded += 1; 139 | const value = try await g.frame; 140 | if (value) |v| { 141 | try handle.yield(v); 142 | g.state = .Next; 143 | } else { 144 | if (allocating) 145 | self.allocator.destroy(g.frame); 146 | g.state = .Done; 147 | active -= 1; 148 | } 149 | } 150 | } 151 | // ...until we run out of reports 152 | if (yielded == 0) break; 153 | } 154 | if (active == 0) break; 155 | } 156 | } 157 | }, T); 158 | } 159 | 160 | test "basic" { 161 | const expect = std.testing.expect; 162 | const ty = struct { 163 | pub fn generate(_: *@This(), handle: *generator.Handle(u8)) !void { 164 | try handle.yield(1); 165 | try handle.yield(2); 166 | try handle.yield(3); 167 | } 168 | }; 169 | const ty1 = struct { 170 | pub fn generate(_: *@This(), handle: *generator.Handle(u8)) !void { 171 | try handle.yield(10); 172 | try handle.yield(20); 173 | try handle.yield(30); 174 | } 175 | }; 176 | 177 | const G0 = generator.Generator(ty, u8); 178 | const G1 = generator.Generator(ty1, u8); 179 | const G = Join(&[_]type{ G0, G1 }, u8, false); 180 | var g = G.init(G.Context.init(.{ G0.init(ty{}), G1.init(ty1{}) })); 181 | 182 | var sum: usize = 0; 183 | while (try g.next()) |v| { 184 | sum += v; 185 | } 186 | try expect(sum == 66); 187 | } 188 | 189 | test "with async i/o" { 190 | // determine file size 191 | const test_file = try std.fs.cwd() 192 | .openFile("README.md", std.fs.File.OpenFlags{ .read = true, .write = false }); 193 | const test_reader = test_file.reader(); 194 | 195 | var file_size: usize = 0; 196 | 197 | while (true) { 198 | _ = test_reader.readByte() catch break; 199 | file_size += 1; 200 | } 201 | 202 | const expect = std.testing.expect; 203 | 204 | // prepare reader type 205 | const ty = struct { 206 | pub fn generate(_: *@This(), handle: *generator.Handle(u8)) !void { 207 | const file = try std.fs.cwd() 208 | .openFile("README.md", std.fs.File.OpenFlags{ .read = true, .write = false }); 209 | const reader = file.reader(); 210 | 211 | while (true) { 212 | const byte = reader.readByte() catch return; 213 | try handle.yield(byte); 214 | } 215 | } 216 | }; 217 | const G0 = generator.Generator(ty, u8); 218 | const G = Join(&[_]type{ G0, G0 }, u8, false); 219 | var g = G.init(G.Context.init(.{ G0.init(ty{}), G0.init(ty{}) })); 220 | 221 | // test 222 | var size: usize = 0; 223 | while (try g.next()) |_| { 224 | size += 1; 225 | } 226 | 227 | try expect(size == file_size * 2); 228 | } 229 | 230 | test "memory impact of not allocating vs allocating frames" { 231 | const ty = struct { 232 | pub fn generate(_: *@This(), handle: *generator.Handle(u8)) !void { 233 | try handle.yield(1); 234 | try handle.yield(2); 235 | try handle.yield(3); 236 | } 237 | }; 238 | const G0 = generator.Generator(ty, u8); 239 | const GAllocating = Join(&[_]type{G0} ** 50, u8, true); 240 | const GNonAllocating = Join(&[_]type{G0} ** 50, u8, false); 241 | 242 | _ = GNonAllocating.init(GNonAllocating.Context.init(.{G0.init(ty{})})); 243 | _ = GAllocating.init(GAllocating.Context.init(.{G0.init(ty{})}, std.testing.allocator)); 244 | 245 | // The assertion below doesn't hold true for all number of joined generators 246 | // as the frame of the allocating Join generator can get larger than of the non-allocating one. 247 | // Could be related to this: 248 | // https://zigforum.org/t/unacceptable-memory-overhead-with-nested-async-function-call/407/5 249 | 250 | // try std.testing.expect(@sizeOf(GAllocating) < @sizeOf(GNonAllocating)); 251 | } 252 | 253 | test "allocating join" { 254 | const expect = std.testing.expect; 255 | 256 | const ty = struct { 257 | pub fn generate(_: *@This(), handle: *generator.Handle(u8)) !void { 258 | try handle.yield(1); 259 | try handle.yield(2); 260 | try handle.yield(3); 261 | } 262 | }; 263 | const G0 = generator.Generator(ty, u8); 264 | const G = Join(&[_]type{G0}, u8, true); 265 | 266 | var g = G.init(G.Context.init(.{G0.init(ty{})}, std.testing.allocator)); 267 | 268 | try expect((try g.next()).? == 1); 269 | try expect((try g.next()).? == 2); 270 | try expect((try g.next()).? == 3); 271 | try expect((try g.next()) == null); 272 | } 273 | -------------------------------------------------------------------------------- /src/lib.zig: -------------------------------------------------------------------------------- 1 | //! This module provides functionality to define async generator functions. 2 | //! 3 | //! It allows one to develop linear algorithms that yield values at certain 4 | //! points without having to maintain re-entrancy manually. 5 | //! 6 | //! ``` 7 | //! const std = @import("std"); 8 | //! const gen = @import("./src/lib.zig"); 9 | //! 10 | //! const Ty = struct { 11 | //! pub fn generate(_: *@This(), handle: *gen.Handle(u8)) !u8 { 12 | //! try handle.yield(0); 13 | //! try handle.yield(1); 14 | //! try handle.yield(2); 15 | //! return 3; 16 | //! } 17 | //! }; 18 | //! 19 | //! const G = gen.Generator(Ty, u8); 20 | //! 21 | //! pub const io_mode = .evented; 22 | //! 23 | //! pub fn main() !void { 24 | //! var g = G.init(Ty{}); 25 | //! 26 | //! std.debug.assert((try g.next()).? == 0); 27 | //! std.debug.assert((try g.next()).? == 1); 28 | //! std.debug.assert((try g.next()).? == 2); 29 | //! std.debug.assert((try g.next()) == null); 30 | //! std.debug.assert(g.state.Returned == 3); 31 | //! } 32 | //! ``` 33 | 34 | pub usingnamespace @import("./generator.zig"); 35 | 36 | pub const Join = @import("./join.zig").Select; 37 | pub const Map = @import("./map.zig").Map; 38 | -------------------------------------------------------------------------------- /src/map.zig: -------------------------------------------------------------------------------- 1 | const std = @import("std"); 2 | const generator = @import("./generator.zig"); 3 | const extractFn = @import("./_.zig").extractFn; 4 | 5 | fn initializer(comptime Self: type, comptime G: type, comptime F: type, comptime stateful: bool) type { 6 | return if (stateful) struct { 7 | pub const Mapper = F; 8 | 9 | pub fn init(inner: G, f_state: F) Self { 10 | return Self{ 11 | .inner = inner, 12 | .state = f_state, 13 | }; 14 | } 15 | } else struct { 16 | pub fn init(inner: G) Self { 17 | return Self{ 18 | .inner = inner, 19 | }; 20 | } 21 | }; 22 | } 23 | 24 | /// `Map` creates a generator that maps yielded values from wrapped generator `G` of type `I` to type `O` 25 | pub fn Map(comptime G: type, comptime I: type, comptime O: type, comptime F: type, comptime stateful: bool) type { 26 | const f = if (stateful) extractFn(F, fn (*F, I) O) else extractFn(F, fn (I) O); 27 | 28 | return generator.Generator(struct { 29 | pub const Inner = G; 30 | inner: G, 31 | 32 | state: if (stateful) F else void = undefined, 33 | pub usingnamespace initializer(@This(), G, F, stateful); 34 | 35 | pub fn generate(self: *@This(), handle: *generator.Handle(O)) !void { 36 | while (try self.inner.next()) |v| { 37 | if (stateful) try handle.yield(f(&self.state, v)) else try handle.yield(f(v)); 38 | } 39 | } 40 | }, I); 41 | } 42 | 43 | test { 44 | const expect = std.testing.expect; 45 | const ty = struct { 46 | pub fn generate(_: *@This(), handle: *generator.Handle(u8)) !void { 47 | try handle.yield(0); 48 | try handle.yield(1); 49 | try handle.yield(2); 50 | } 51 | }; 52 | const G = Map(generator.Generator(ty, u8), u8, u8, struct { 53 | pub fn incr(i: u8) u8 { 54 | return i + 1; 55 | } 56 | }, false); 57 | var g = G.init(G.Context.init(G.Context.Inner.init(ty{}))); 58 | 59 | try expect((try g.next()).? == 1); 60 | try expect((try g.next()).? == 2); 61 | try expect((try g.next()).? == 3); 62 | try expect((try g.next()) == null); 63 | try expect(g.state == .Returned); 64 | } 65 | 66 | test "stateful" { 67 | const expect = std.testing.expect; 68 | const ty = struct { 69 | pub fn generate(_: *@This(), handle: *generator.Handle(u8)) !void { 70 | try handle.yield(0); 71 | try handle.yield(1); 72 | try handle.yield(2); 73 | } 74 | }; 75 | const G = Map(generator.Generator(ty, u8), u8, u8, struct { 76 | n: u8 = 0, 77 | pub fn incr(self: *@This(), i: u8) u8 { 78 | self.n += 1; 79 | return i + self.n; 80 | } 81 | }, true); 82 | var g = G.init(G.Context.init(G.Context.Inner.init(ty{}), G.Context.Mapper{})); 83 | 84 | try expect((try g.next()).? == 1); 85 | try expect((try g.next()).? == 3); 86 | try expect((try g.next()).? == 5); 87 | try expect((try g.next()) == null); 88 | try expect(g.state == .Returned); 89 | } 90 | -------------------------------------------------------------------------------- /src/tests.zig: -------------------------------------------------------------------------------- 1 | test { 2 | _ = @import("./join.zig"); 3 | _ = @import("./map.zig"); 4 | } 5 | -------------------------------------------------------------------------------- /zig.mod: -------------------------------------------------------------------------------- 1 | id: 9xwrubjr5dq4etfef4rbwkccintq9ear6atnefhv2bffx4y6 2 | name: generator 3 | main: src/lib.zig 4 | license: MIT 5 | description: Async generator type 6 | dev_dependencies: 7 | - src: git https://github.com/Hejsil/zig-bench 8 | -------------------------------------------------------------------------------- /zigmod.lock: -------------------------------------------------------------------------------- 1 | 2 2 | git https://github.com/Hejsil/zig-bench commit-3810f2e56129374c48c54f6fde995e3c941d7976 3 | --------------------------------------------------------------------------------