const std = @import("std");

/// stores some state for each coroutine
/// as well as a pointer to the frame to be resumed
const CoroutineState = struct {
    pub const Self = @This();

    frame: anyframe = undefined,
    quit: bool = false,
    active: bool = false,

    /// must be called from within the coroutine
    /// as soon as the coroutine will exit.
    /// this is best done by using `defer coro.exit();`.
    pub fn exit(self: *Self) void {
        self.active = false;
    }

    /// yields the coroutine.
    /// will return an error if the coroutine has been quit.
    fn yield(coro: *Self) !void {
        suspend {
            coro.frame = @frame();
        }
        if (coro.quit)
            return error.Cancelled;
    }
};

var coroutines = [_]CoroutineState{CoroutineState{}} ** 3;

/// our coroutine
fn loop(coro: *CoroutineState, prefix: []const u8, limit: i32) !void {
    defer coro.exit(); // must be called to recognize that the coroutine was stopped

    std.debug.warn("{}: startup\n", prefix);
    defer std.debug.warn("{}: shutdown\n", prefix);

    var repetition: i32 = 0;
    while (repetition < limit) : (repetition += 1) {
        std.debug.warn("{}: loop {}\n", prefix, repetition);
        try coro.yield(); // yield may throw error.Cancelled
    }
}

/// starts a new coroutine with the given parameters
fn start_new(prefix: []const u8, limit: i32) !@Frame(loop) {
    for (coroutines) |*coro| {
        if (coro.active)
            continue;
        coro.active = true;
        coro.quit = false;
        coro.frame = undefined;
        return async loop(coro, prefix, limit);
    }
    return error.OutOfMemory;
}

pub fn main() !void {
    // this uses result location semantics to
    // store our coroutine frames in the main
    // function. in a proper implementation,
    // it should use an allocator
    _ = try start_new("o", 5);
    _ = try start_new("x", 10);
    _ = try start_new("~", 20);

    // schedule our coroutines here 15 times
    var rep: u32 = 0;
    while (rep < 15) : (rep += 1) {
        var any_alive = false;
        for (coroutines) |coro| {
            // resume all alive coroutines
            if (!coro.active)
                continue;
            any_alive = true;
            resume coro.frame;
        }
        if (!any_alive)
            break;
    }

    // kill all coroutines that are still alive
    for (coroutines) |*coro, index| {
        if (coro.active) {
            coro.quit = true;
            resume coro.frame;
        }
    }
}