tasks.zig

const std = @import("std");
const direct_allocator = std.heap.direct_allocator;

// The goal is to prototype a system where functions and their parameters are
// pushed in a global queue that feeds background worker threads.
// async frames seem like the perfect choice for this because then we don't
// need to have indirections or type erasure to store the parameters : they
// are in the frame. This implies that we can write the wrap function more
// or less as it is now.

pub fn main() void {
    defer tasks.deinit();
    _ = async offloadMain();
}

fn offloadMain() void {
    const f1 = offload(foo, @as(u32, 45), @as(u32, 65));
    const f2 = offload(bar, @as(u32, 45), @as(u32, 65));
    const f3 = offload(baz);

    // Simulate threads executing
    std.debug.warn("foo={}\n", f1.get());
    std.debug.warn("bar={}\n", f2.get());
}

/// ----------------------
/// Dummy functions we want to offload on worker threads.
fn foo(a: u32, b: u32) u32 {
    return a + b;
}

fn bar(a: u32, b: u32) u32 {
    return a * b;
}

fn baz() void {
    std.debug.warn("baz\n");
}

/// ----------------------
/// Holds a task execution context. Created to execute a function.
const Task = struct {
    const FrameAlign: u29 = 64;

    frame: anyframe,
    bytes: []align(FrameAlign) u8,
};

/// Represents a task queue used by multiple threads to pop tasks to execute.
var tasks = std.ArrayList(Task).init(direct_allocator);

/// Basically a ticket for an offloaded task.
/// Can be waited for termination or ignored, can be used to retrieve a result.
fn Future(comptime return_type: type) type {
    return struct {
        const Self = @This();
        const Return = return_type;

        id: usize,

        /// Emulates the fact that a worker thread would resume (execute) the task.
        /// Returns the result of the task.
        pub fn get(self: Self) Return {
            const task = &tasks.toSlice()[self.id];
            resume task.frame;
            return await @ptrCast(anyframe->Return, task.frame);
        }
    };
}

/// Entry point for users, give the system a function and its arguments to
/// push an asynchronous task that will be executed by a worker thread.
fn offload(comptime fun: var, args: ...) Future(@typeInfo(@typeOf(fun)).Fn.return_type.?) {
    comptime std.debug.assert(@alignOf(@Frame(fun)) <= Task.FrameAlign);

    // Obviously this doesn't work, it would need a way to get the type of
    // `wrap` instanciated with expanded `args`
    const frame_size = @sizeOf(@Frame(wrap));

    const bytes = direct_allocator.alignedAlloc(u8, 64, frame_size) catch unreachable;

    const idx = tasks.count();
    tasks.append(Task{ .frame = undefined, .bytes = bytes }) catch unreachable;

    wrap(fun, tasks.toSlice()[idx], args);

    std.debug.warn("offloaded {}\n", @typeName(@typeOf(fun)));

    return Future(@typeInfo(@typeOf(fun)).Fn.return_type.?){ .id = idx };
}

/// Wrapper to have the arguments stored in the task's frame. The goal is to
/// avoid storing them explicitely because ideally we don't need to.
fn wrap(comptime fun: var, frame: *anyframe, args: ...) void {
    frame.? = @Frame();
    suspend;
    fun(args);
}