Select

pub fn Select(comptime U: type) type {
    return struct {
        io: Io,
        group: Group,
        queue: Queue(U),

        const S = @This();

        pub const Union = U;

        pub const Field = std.meta.FieldEnum(U);

        pub fn init(io: Io, buffer: []U) S {
            return .{
                .io = io,
                .queue = .init(buffer),
                .group = .init,
            };
        }

        /// Calls `function` with `args` asynchronously. The resource spawned is
        /// owned by the select.
        ///
        /// `function` must have return type matching the `field` field of `Union`.
        ///
        /// `function` *may* be called immediately, before `async` returns.
        ///
        /// When this function returns, it is guaranteed that `function` has
        /// already been called and completed, or it has successfully been
        /// assigned a unit of concurrency.
        ///
        /// After this is called, `await` or `cancel` must be called before the
        /// select is deinitialized.
        ///
        /// Threadsafe.
        ///
        /// Related:
        /// * `Io.async`
        /// * `Group.async`
        pub fn async(
            s: *S,
            comptime field: Field,
            function: anytype,
            args: std.meta.ArgsTuple(@TypeOf(function)),
        ) void {
            const Context = struct {
                select: *S,
                args: @TypeOf(args),
                fn start(type_erased_context: *const anyopaque) void {
                    const context: *const @This() = @ptrCast(@alignCast(type_erased_context));
                    const result = @call(.auto, function, context.args);
                    const elem = @unionInit(U, @tagName(field), result);
                    context.select.queue.putOneUncancelable(context.select.io, elem) catch |err| switch (err) {
                        error.Closed => {},
                    };
                }
            };
            const context: Context = .{ .select = s, .args = args };
            s.io.vtable.groupAsync(s.io.userdata, &s.group, @ptrCast(&context), .of(Context), Context.start);
        }

        /// Calls `function` with `args` concurrently. The resource spawned is
        /// owned by the select.
        ///
        /// `function` must have return type matching the `field` field of `Union`.
        ///
        /// After this function returns successfully, it is guaranteed that
        /// `function` has been assigned a unit of concurrency, and `await` or
        /// `cancel` must be called before the select is deinitialized.
        ///
        ///
        /// Threadsafe.
        ///
        /// Related:
        /// * `Io.concurrent`
        /// * `Group.concurrent`
        pub fn concurrent(
            s: *S,
            comptime field: Field,
            function: anytype,
            args: std.meta.ArgsTuple(@TypeOf(function)),
        ) ConcurrentError!void {
            const Context = struct {
                select: *S,
                args: @TypeOf(args),
                fn start(type_erased_context: *const anyopaque) void {
                    const context: *const @This() = @ptrCast(@alignCast(type_erased_context));
                    const result = @call(.auto, function, context.args);
                    const elem = @unionInit(U, @tagName(field), result);
                    context.select.queue.putOneUncancelable(context.select.io, elem) catch |err| switch (err) {
                        error.Closed => {},
                    };
                }
            };
            const context: Context = .{ .select = s, .args = args };
            try s.io.vtable.groupConcurrent(s.io.userdata, &s.group, @ptrCast(&context), .of(Context), Context.start);
        }

        /// Blocks until another task of the select finishes.
        ///
        /// It is legal to call `async` and `concurrent` after this.
        ///
        /// Threadsafe.
        pub fn await(s: *S) Cancelable!U {
            return s.queue.getOne(s.io) catch |err| switch (err) {
                error.Canceled => |e| return e,
                error.Closed => unreachable,
            };
        }

        /// Blocks until at least `min` number of results have been copied
        /// into `buffer`.
        ///
        /// Asserts that `buffer.len >= min`.
        ///
        /// It is legal to call `async` and `concurrent` after this.
        ///
        /// Threadsafe.
        pub fn awaitMany(s: *S, buffer: []U, min: usize) Cancelable!usize {
            return s.queue.get(s.io, buffer, min) catch |err| switch (err) {
                error.Canceled => |e| return e,
                error.Closed => unreachable,
            };
        }

        /// Requests cancelation on all remaining tasks owned by the select,
        /// then blocks until they all finish. If the select was initialized
        /// with insufficient buffer space for all remaining tasks to finish, a
        /// deadlock occurs.
        ///
        /// If any of the select tasks allocate resources, those tasks may have
        /// completed, meaning that this function must be called in a loop
        /// until `null` is returned in order to deallocate those resources. If
        /// there is no possibility of resource leaks, `cancelDiscard` is
        /// preferable.
        ///
        /// It is illegal to call `await` or `awaitMany` after this.
        ///
        /// It is safe to call this multiple times, even after `null` is
        /// returned.
        ///
        /// Threadsafe.
        pub fn cancel(s: *S) ?U {
            const io = s.io;
            s.group.cancel(io);
            s.queue.close(io);
            return s.queue.getOneUncancelable(io) catch |err| switch (err) {
                error.Closed => return null,
            };
        }

        /// Requests cancelation on all remaining tasks owned by the select,
        /// then blocks until they all finish.
        ///
        /// All return values from outstanding tasks are discarded. This
        /// function is therefore inappropriate to call when a task can return
        /// an allocated resource. For that use case, see `cancel`.
        ///
        /// It is illegal to call `await` or `awaitMany` after this.
        ///
        /// It is safe to call this multiple times.
        ///
        /// Threadsafe.
        pub fn cancelDiscard(s: *S) void {
            const io = s.io;
            const token = s.group.token.load(.acquire) orelse return;
            s.queue.close(io);
            io.vtable.groupCancel(io.userdata, &s.group, token);
            assert(s.group.token.raw == null);
        }
    };
}