wasmslime/src/func.zig

const std = @import("std");
const assert = std.debug.assert;
const panic = std.debug.panic;

const Store = @import("./store.zig").Store;
const Extern = @import("./extern.zig").Extern;
const Memory = @import("./memory.zig").Memory;
const Diagnostics = @import("./diagnostics.zig").Diagnostics;
const Error = @import("./error.zig").Error;
const Trap = @import("./trap.zig").Trap;

const _val = @import("./val.zig");
const Val = _val.Val;
const ValType = _val.ValType;
const ValTypeVec = _val.ValTypeVec;

/// Structure used to learn about the caller of a host-defined function.
///
/// This structure is an argument to `Func.Callback`. Its purpose is to
/// acquire a `Store.Context` pointer to interact with objects, but it can
/// also be used to inspect the state of the caller (such as getting memories
/// and functions) with `getExport`.
pub const Caller = opaque {
    /// Returns the store context of the caller object.
    pub fn getContext(self: *Caller) *Store.Context {
        return wasmtime_caller_context(self);
    }

    /// Get an exported `Extern` from the caller's context.
    /// Returns `error.ExportNotFound` if export is not found.
    pub fn getExport(self: *Caller, name: []const u8) !Extern {
        var result: Extern = undefined;
        return if (wasmtime_caller_export_get(self, name.ptr, name.len, &result)) result else error.ExportNotFound;
    }

    /// Get an exported `Func` from the caller's context.
    /// Returns `error.ExportNotFound` if export is not found.
    /// Returns `error.ExportIncorrectType` if the export isn't a `Func`.
    pub fn getFunc(self: *Caller, name: []const u8) !Func {
        return (try self.getExport(name)).asFunc() orelse error.ExportIncorrectType;
    }

    /// Get an exported `Memory` from the caller's context.
    /// Returns `error.ExportNotFound` if export is not found.
    /// Returns `error.ExportIncorrectType` if the export isn't a `Memory`.
    pub fn getMemory(self: *Caller, name: []const u8) !Memory {
        return (try self.getExport(name)).asMemory() orelse error.ExportIncorrectType;
    }

    /// Get an exported `Memory`'s data slice from the caller's context.
    /// Returns `error.ExportNotFound` if export is not found.
    /// Returns `error.ExportIncorrectType` if the export isn't a `Memory`.
    pub fn getMemoryData(self: *Caller, name: []const u8) ![]u8 {
        return (try self.getMemory(name)).getData(self.getContext());
    }

    extern "c" fn wasmtime_caller_context(*Caller) *Store.Context;
    extern "c" fn wasmtime_caller_export_get(*Caller, [*]const u8, usize, *Extern) bool;
};

/// Representation of a function in Wasmtime.
pub const Func = extern struct {
    /// Internal identifier of what store this belongs to, never zero.
    store_id: u64,
    /// Internal index within the store.
    index: usize,

    /// Inserts a host-defined function into the `store` provided, which can
    /// be used to then instantiate a module with or define within a `Linker`.
    /// The returned function can only be used with the specified `store`.
    pub fn init(
        context: *Store.Context,
        func_type: *FuncType,
        callback: *const Callback,
        env: ?*anyopaque,
        env_finalizer: ?*const Finalizer,
    ) Func {
        var result: Func = undefined;
        wasmtime_func_new(context, func_type, callback, env, env_finalizer, &result);
        return result;
    }

    pub fn initFromFn(context: *Store.Context, func: anytype) Func {
        const func_type = FuncType.fromFn(func);
        const wrapped = Wrapped.init(func);
        return init(context, func_type, wrapped.callback, wrapped.env, wrapped.env_finalizer);
    }

    pub const Wrapped = struct {
        env: ?*anyopaque,
        env_finalizer: ?*const Finalizer,
        callback: *const Callback,

        pub fn init(func: anytype) Wrapped {
            return .{
                // Use the function's pointer itself as the "env" variable
                // that'll be passed to the callback. This way we can call
                // the original function from our custom callback.
                .env = @ptrCast(@constCast(&func)),
                .env_finalizer = null, // No finalizer necessary.
                .callback = &(struct {
                    const meta = @typeInfo(@TypeOf(func)).Fn;
                    const ArgsType = std.meta.ArgsTuple(@TypeOf(func));
                    const ResultType = meta.return_type orelse void;

                    fn callback(
                        env: ?*anyopaque,
                        caller: *Caller,
                        args: [*]const Val,
                        num_args: usize,
                        results: [*]Val,
                        num_results: usize,
                    ) ?*Trap {
                        var func_args: ArgsType = undefined;

                        comptime var offset = 0;
                        // If the first function parameter accepts a *Caller, pass that in!
                        if (meta.params.len > 0 and meta.params[0].type.? == *Caller) {
                            func_args[0] = caller;
                            offset = 1;
                        }

                        assert(num_args == meta.params.len - offset);
                        assert(num_results == getNumResultsFromResultType(ResultType));

                        // Turn the "env" parameter back into a function pointer.
                        const func_ptr: *const @TypeOf(func) = @ptrCast(env);

                        // Build argument tuple from the arguments WASM has given us.
                        inline for (meta.params[offset..], args, offset..) |p, arg, i|
                            func_args[i] = arg.toValue(p.type.?) catch
                                panic("Cannot convert {} to {s}", .{ arg, @typeName(p.type.?) });

                        // Call the function with the arguments, get the results.
                        const func_results = @call(.auto, func_ptr, func_args);

                        // Pass the result(s) back to WASM.
                        if (comptime std.meta.trait.isTuple(ResultType)) {
                            inline for (func_results, 0..) |res, i|
                                results[i] = Val.fromValue(@TypeOf(res), res);
                        } else if (ResultType != void)
                            results[0] = Val.fromValue(@TypeOf(func_results), func_results);

                        return null;
                    }
                }.callback),
            };
        }
    };

    pub fn getType(self: *const Func, context: *Store.Context) *FuncType {
        return wasmtime_func_type(context, self);
    }

    pub fn call(
        self: *const Func,
        context: *const Store.Context,
        args: anytype,
        comptime TResult: type,
        diag: ?*Diagnostics,
    ) !TResult {
        if (!comptime isTupleOrEmpty(@TypeOf(args)))
            @compileError("Expected tuple, but found '" ++ @typeName(@TypeOf(args)) ++ "'");

        // Set up function arguments.
        var args_array: [args.len]Val = undefined;
        inline for (&args_array, args) |*val, arg|
            val.* = Val.fromValue(@TypeOf(arg), arg);

        // Set up function results.
        const num_results = comptime getNumResultsFromResultType(TResult);
        var results_array: [num_results]Val = undefined;

        var trap: ?*Trap = null;
        const err = wasmtime_func_call(context, self, &args_array, args_array.len, &results_array, results_array.len, &trap);
        try Diagnostics.handleErrorOrTrap(err, error.FuncCall, trap, {}, diag);

        if (comptime std.meta.trait.isTuple(TResult)) {
            var result: TResult = undefined;
            inline for (TResult, results_array, result) |T, val, *res| {
                defer val.deinit();
                res.* = val.toValue(T) catch
                    panic("Cannot convert {} to {s}", .{ val, @typeName(T) });
            }
            return result;
        } else if (TResult != void) {
            defer results_array[0].deinit();
            return results_array[0].toValue(TResult) catch
                panic("Cannot convert {} to {s}", .{ results_array[0], @typeName(TResult) });
        } else return {};
    }

    pub fn toExtern(self: Func) Extern {
        return Extern.fromFunc(self);
    }

    fn isTupleOrEmpty(comptime T: type) bool {
        return switch (@typeInfo(T)) {
            .Struct => |s| s.is_tuple or s.fields.len == 0,
            else => false,
        };
    }

    fn getNumResultsFromResultType(comptime T: type) usize {
        return comptime if (std.meta.trait.isTuple(T))
            @typeInfo(T).Struct.fields.len
        else if (T != void) 1 else 0;
    }

    pub const Callback = fn (
        env: ?*anyopaque,
        caller: *Caller,
        args: [*]const Val,
        num_args: usize,
        results: [*]Val,
        num_results: usize,
    ) ?*Trap;

    pub const Finalizer = fn (env: *anyopaque) void;

    extern "c" fn wasmtime_func_new(
        context: *Store.Context,
        func_type: *const FuncType,
        callback: *const Callback,
        env: ?*anyopaque,
        env_finalizer: ?*const Finalizer,
        result: *Func,
    ) void;

    extern "c" fn wasmtime_func_type(
        context: *const Store.Context,
        func: *const Func,
    ) *FuncType;

    extern "c" fn wasmtime_func_call(
        context: *const Store.Context,
        func: *const Func,
        args: [*]const Val,
        num_args: usize,
        results: [*]Val,
        num_results: usize,
        trap: *?*Trap,
    ) ?*Error;
};

/// An opaque object representing the type of a function.
pub const FuncType = opaque {
    pub fn init(params: *ValTypeVec, results: *ValTypeVec) *FuncType {
        return wasm_functype_new(params, results);
    }

    pub fn fromFn(func: anytype) *FuncType {
        const meta = @typeInfo(@TypeOf(func)).Fn;

        comptime var offset = 0;
        // If the first function parameter accepts a *Caller, skip it!
        if (meta.params.len > 0 and meta.params[0].type.? == *Caller) offset = 1;

        var params = ValTypeVec.initUninitialized(meta.params.len - offset);
        inline for (meta.params[offset..], 0..) |p, i|
            params.data[i] = ValType.fromType(p.type.?);

        const ResultType = meta.return_type orelse void;
        var results = if (comptime std.meta.trait.isTuple(ResultType)) blk: {
            var res = ValTypeVec.initUninitialized(ResultType.len);
            inline for (ResultType, res.toSlice()) |T, *r|
                r.* = ValType.fromType(T);
            break :blk res;
        } else if (ResultType != void) blk: {
            var res = ValTypeVec.initUninitialized(1);
            res.data[0] = ValType.fromType(ResultType);
            break :blk res;
        } else ValTypeVec.initEmpty(0);

        return init(&params, &results);
    }

    pub fn copy(self: *const FuncType) *FuncType {
        return wasm_functype_copy(self);
    }

    pub fn deinit(self: *FuncType) void {
        wasm_functype_delete(self);
    }

    pub fn getParams(self: *const FuncType) *const ValTypeVec {
        return wasm_functype_params(self);
    }

    pub fn getResults(self: *const FuncType) *const ValTypeVec {
        return wasm_functype_results(self);
    }

    extern "c" fn wasm_functype_new(*ValTypeVec, *ValTypeVec) *FuncType;
    extern "c" fn wasm_functype_copy(*const FuncType) *FuncType;
    extern "c" fn wasm_functype_delete(*FuncType) void;
    extern "c" fn wasm_functype_params(*const FuncType) *const ValTypeVec;
    extern "c" fn wasm_functype_results(*const FuncType) *const ValTypeVec;
};

// pub const FuncTypeVec = Vec(*FuncType, wasm_functype_vec_new_empty, wasm_functype_vec_new_uninitialized, wasm_functype_vec_new, wasm_functype_vec_copy, wasm_functype_vec_delete);
// extern "c" fn wasm_functype_vec_new_empty(*anyopaque, usize) void;
// extern "c" fn wasm_functype_vec_new_uninitialized(*anyopaque, usize) void;
// extern "c" fn wasm_functype_vec_new(*anyopaque, usize, [*]const ?*FuncType) void;
// extern "c" fn wasm_functype_vec_copy(*anyopaque, *const anyopaque) void;
// extern "c" fn wasm_functype_vec_delete(*anyopaque) void;