rkgk/crates/haku2/src/vm.zig

const std = @import("std");
const mem = std.mem;
const debug = std.debug;
const log = std.log.scoped(.vm);
const testAllocator = std.testing.allocator;

const bytecode = @import("bytecode.zig");
const Canvas = @import("canvas.zig");
const system = @import("system.zig");
const value = @import("value.zig");
const Value = value.Value;

const Vm = @This();

stack: []Value,
stack_top: u32 = 0,
call_stack: []CallFrame,
call_stack_top: u32 = 0,
defs: []Value,
fuel: u32,
exception: ?Exception = null,

pub const Limits = struct {
    stack_capacity: usize = 256,
    call_stack_capacity: usize = 256,
    fuel: u32 = 63336,
};

pub const CallFrame = struct {
    closure: *const value.Closure,
    ip: [*]const u8,
    bottom: u32,
};

pub const Exception = struct {
    len: usize,
    format: *const fn (buf: []u8, args: *align(64) const anyopaque) []u8,
    args: [40]u8 align(64), // increase the size if we ever throw a larger exception
};

/// All errors coming from inside the VM get turned into a single Exception type, which signals
/// any kind of error that prevented the machine from running any further.
/// Details about the exception can be found in the VM's `exception` field.
pub const Error = error{Exception};

// NOTE: A VM is only ever initialized. There is no function for deallocating its resources.
// The intent is to use it with a freshly-reset Scratch, and deallocate its resources through that.
pub fn init(a: mem.Allocator, defs: *const bytecode.Defs, limits: *const Limits) !Vm {
    return .{
        .stack = try a.alloc(Value, limits.stack_capacity),
        .call_stack = try a.alloc(CallFrame, limits.call_stack_capacity),
        .defs = try a.alloc(Value, defs.num_defs),
        .fuel = limits.fuel,
    };
}

pub fn throw(vm: *Vm, comptime fmt: []const u8, args: anytype) Error {
    log.debug("throw: fmt={s}", .{fmt});

    const Args = @TypeOf(args);
    const max_args_size = @sizeOf(@TypeOf(vm.exception.?.args));
    if (@sizeOf(Args) > max_args_size) {
        @compileError(std.fmt.comptimePrint(
            "format arguments are too large; size={}, max={}",
            .{ @sizeOf(Args), max_args_size },
        ));
    }

    const Formatter = struct {
        fn format(buf: []u8, erased_args: *align(64) const anyopaque) []u8 {
            return std.fmt.bufPrint(buf, fmt, @as(*const Args, @ptrCast(erased_args)).*) catch unreachable;
        }
    };

    var exn = Exception{
        .len = @truncate(std.fmt.count(fmt, args)),
        .format = Formatter.format,
        .args = undefined,
    };
    @memcpy(exn.args[0..@sizeOf(Args)], mem.asBytes(&args));
    vm.exception = exn;

    return error.Exception;
}

pub fn outOfMemory(vm: *Vm) Error {
    return vm.throw("out of memory", .{});
}

/// Debug assertion for bytecode validity.
/// In future versions, this may become disabled in release builds.
fn validateBytecode(vm: *Vm, ok: bool, comptime fmt: []const u8, args: anytype) Error!void {
    if (!ok) {
        return vm.throw("corrupted bytecode: " ++ fmt, args);
    }
}

pub fn consumeFuel(vm: *Vm, fuel: *u32, amount: u32) Error!void {
    const new_fuel, const overflow = @subWithOverflow(fuel.*, amount);
    if (overflow > 0) {
        return vm.throw("code ran for too long (out of fuel!)", .{});
    }
    fuel.* = new_fuel;
}

pub fn push(vm: *Vm, val: Value) Error!void {
    if (vm.stack_top >= vm.stack.len) {
        return vm.throw("too many live temporary values (local variables and expression operands)", .{});
    }
    log.debug("PUSH {any} <- {}", .{ vm.stack[0..vm.stack_top], val });
    vm.stack[vm.stack_top] = val;
    vm.stack_top += 1;
}

pub fn pop(vm: *Vm) Error!Value {
    try vm.validateBytecode(vm.stack_top > 0, "value stack underflow", .{});
    vm.stack_top -= 1;
    const result = vm.stack[vm.stack_top];
    log.debug("POP  {any} -> {}", .{ vm.stack[0..vm.stack_top], result });
    return vm.stack[vm.stack_top];
}

pub fn top(vm: *const Vm) Value {
    if (vm.stack_top > 0) {
        return vm.stack[vm.stack_top - 1];
    } else {
        return .nil;
    }
}

pub fn pushCall(vm: *Vm, frame: CallFrame) Error!void {
    if (vm.call_stack_top >= vm.call_stack.len) {
        return vm.throw("too much recursion", .{});
    }
    log.debug("PUSH CALL {}", .{frame});
    vm.call_stack[vm.call_stack_top] = frame;
    vm.call_stack_top += 1;
}

pub fn popCall(vm: *Vm) Error!CallFrame {
    try vm.validateBytecode(vm.call_stack_top > 0, "call stack underflow", .{});
    log.debug("POP CALL", .{});
    vm.call_stack_top -= 1;
    return vm.call_stack[vm.call_stack_top];
}

pub fn local(vm: *Vm, bottom: u32, offset: u8) Error!*Value {
    const index = bottom + offset;
    try vm.validateBytecode(index < vm.stack_top, "stack index out of bounds", .{});
    return &vm.stack[bottom + offset];
}

pub fn def(vm: *Vm, index: u16) Error!*Value {
    try vm.validateBytecode(index < vm.defs.len, "def index out of bounds", .{});
    return &vm.defs[index];
}

// Utility struct for storing and restoring the VM's state variables across FFI boundaries.
const Context = struct {
    fuel: u32,
};

fn restoreContext(vm: *Vm) Context {
    return .{
        .fuel = vm.fuel,
    };
}

fn storeContext(vm: *Vm, context: Context) void {
    vm.fuel = context.fuel;
}

inline fn read(comptime T: type, ip: *[*]const u8) T {
    const result = mem.readInt(T, ip.*[0..@sizeOf(T)], std.builtin.Endian.little);
    ip.* = ip.*[@sizeOf(T)..];
    return result;
}

inline fn readOpcode(ip: *[*]const u8) bytecode.Opcode {
    const opcode: bytecode.Opcode = @enumFromInt(read(u8, ip));
    log.debug("OP {*} {}", .{ ip.*, opcode });
    return opcode;
}

/// Before calling this, vm.stack_top should be saved and the appropriate amount of parameters must
/// be pushed onto the stack to execute the provided closure.
/// The saved vm.stack_top must be passed to init_bottom.
///
/// allocator should be a scratch buffer; there is no way to free the memory allocated by a VM.
pub fn run(
    vm: *Vm,
    allocator: mem.Allocator,
    init_closure: *const value.Closure,
    init_bottom: u32,
) Error!void {
    log.debug("BEGIN RUN {}", .{init_closure});

    var closure = init_closure;
    var ip: [*]const u8 = closure.chunk.bytecode[closure.start..].ptr;
    var bottom = init_bottom;
    var fuel = vm.fuel;

    for (0..closure.local_count) |_| {
        try vm.push(.nil);
    }

    const call_bottom = vm.call_stack_top;
    try vm.pushCall(.{
        .closure = closure,
        .ip = ip,
        .bottom = bottom,
    });

    next: switch (readOpcode(&ip)) {
        .nil => {
            try vm.consumeFuel(&fuel, 1);
            try vm.push(.nil);
            continue :next readOpcode(&ip);
        },

        .false => {
            try vm.consumeFuel(&fuel, 1);
            try vm.push(.false);
            continue :next readOpcode(&ip);
        },

        .true => {
            try vm.consumeFuel(&fuel, 1);
            try vm.push(.true);
            continue :next readOpcode(&ip);
        },

        .tag => {
            try vm.consumeFuel(&fuel, 1);
            const tag_id: value.TagId = @enumFromInt(read(u16, &ip));
            try vm.push(.{ .tag = tag_id });
            continue :next readOpcode(&ip);
        },

        .number => {
            try vm.consumeFuel(&fuel, 1);
            const number: f32 = @bitCast(read(u32, &ip));
            try vm.push(.{ .number = number });
            continue :next readOpcode(&ip);
        },

        .rgba => {
            try vm.consumeFuel(&fuel, 1);
            const r = read(u8, &ip);
            const g = read(u8, &ip);
            const b = read(u8, &ip);
            const a = read(u8, &ip);
            try vm.push(.{ .rgba = value.rgbaFrom8(value.Rgba8{ r, g, b, a }) });
            continue :next readOpcode(&ip);
        },

        .local => {
            try vm.consumeFuel(&fuel, 1);
            const index = read(u8, &ip);
            const l = try vm.local(bottom, index);
            try vm.push(l.*);
            continue :next readOpcode(&ip);
        },

        .set_local => {
            try vm.consumeFuel(&fuel, 1);
            const index = read(u8, &ip);
            const new = try vm.pop();
            const l = try vm.local(bottom, index);
            l.* = new;
            continue :next readOpcode(&ip);
        },

        .capture => {
            try vm.consumeFuel(&fuel, 1);
            const index = read(u8, &ip);
            try vm.validateBytecode(index < closure.captures.len, "capture index out of bounds", .{});
            const capture = closure.captures[index];
            try vm.push(capture);
            continue :next readOpcode(&ip);
        },

        .def => {
            try vm.consumeFuel(&fuel, 1);
            const index = read(u16, &ip);
            const d = try vm.def(index);
            try vm.push(d.*);
            continue :next readOpcode(&ip);
        },

        .set_def => {
            try vm.consumeFuel(&fuel, 1);
            const index = read(u16, &ip);
            const new = try vm.pop();
            const d = try vm.def(index);
            d.* = new;
            continue :next readOpcode(&ip);
        },

        .list => {
            const len = read(u16, &ip);
            try vm.consumeFuel(&fuel, len);
            const list_end = vm.stack_top;
            vm.stack_top -= len;
            const elements = vm.stack[vm.stack_top..list_end];
            const list = allocator.dupe(Value, elements) catch return vm.outOfMemory();
            const ref = allocator.create(value.Ref) catch return vm.outOfMemory();
            ref.* = .{ .list = list };
            try vm.push(.{ .ref = ref });
            continue :next readOpcode(&ip);
        },

        .function => {
            try vm.consumeFuel(&fuel, 1);
            const param_count = read(u8, &ip);
            const then = read(u16, &ip);
            const body = ip;
            ip = closure.chunk.bytecode[then..].ptr;

            const local_count = read(u8, &ip);
            const capture_count = read(u8, &ip);

            const captures = allocator.alloc(Value, capture_count) catch return vm.outOfMemory();
            for (captures) |*capture| {
                const capture_kind = read(u8, &ip);
                const index = read(u8, &ip);
                capture.* = switch (capture_kind) {
                    bytecode.capture_local => (try vm.local(bottom, index)).*,
                    bytecode.capture_capture => blk: {
                        try vm.validateBytecode(index < closure.captures.len, "capture index out of bounds", .{});
                        break :blk closure.captures[index];
                    },
                    else => .nil,
                };
            }

            const new_closure = value.Closure{
                .chunk = closure.chunk,
                .start = @truncate(body - closure.chunk.bytecode.ptr),
                .param_count = param_count,
                .local_count = local_count,
                .captures = captures,
            };
            const ref = allocator.create(value.Ref) catch return vm.outOfMemory();
            ref.* = .{ .closure = new_closure };
            try vm.push(.{ .ref = ref });

            continue :next readOpcode(&ip);
        },

        .jump => {
            try vm.consumeFuel(&fuel, 1);
            const offset = read(u16, &ip);
            ip = closure.chunk.bytecode[offset..].ptr;
            continue :next readOpcode(&ip);
        },

        .jump_if_not => {
            try vm.consumeFuel(&fuel, 1);
            const offset = read(u16, &ip);
            const condition = try vm.pop();
            if (!condition.isTruthy()) {
                ip = closure.chunk.bytecode[offset..].ptr;
            }
            continue :next readOpcode(&ip);
        },

        .field => {
            try vm.consumeFuel(&fuel, 1);

            const tag_count = read(u8, &ip);
            const field_tag = try vm.pop();
            if (field_tag != .tag) {
                return vm.throw("name of data field to look up must be a tag (starting with an uppercase letter)", .{});
            }
            const field_tag_id = field_tag.tag;

            var found_index: ?usize = null;
            for (0..tag_count) |i| {
                const tag_id: value.TagId = @enumFromInt(read(u16, &ip));
                if (tag_id == field_tag_id) {
                    found_index = i;
                }
            }

            if (found_index) |index| {
                try vm.validateBytecode(index < closure.captures.len, "field index out of bounds", .{});
                const field = closure.captures[index];
                try vm.push(field);
            } else {
                return vm.throw("field with this name does not exist", .{});
            }

            continue :next readOpcode(&ip);
        },

        .call => {
            try vm.consumeFuel(&fuel, 1);

            const arg_count = read(u8, &ip);

            const function_value = try vm.pop();
            if (function_value != .ref or function_value.ref.* != .closure) {
                return vm.throw("attempt to call a value that is not a function", .{});
            }
            const called_closure = &function_value.ref.closure;

            if (arg_count != called_closure.param_count) {
                return vm.throw(
                    "function expects {} arguments, but it received {}",
                    .{ called_closure.param_count, arg_count },
                );
            }

            const call_frame = CallFrame{
                .closure = closure,
                .ip = ip,
                .bottom = bottom,
            };

            const new_bottom, const overflow = @subWithOverflow(vm.stack_top, arg_count);
            try vm.validateBytecode(overflow == 0, "not enough values on the stack for arguments", .{});

            closure = called_closure;
            ip = closure.chunk.bytecode[closure.start..].ptr;
            bottom = new_bottom;

            for (0..closure.local_count) |_| {
                try vm.push(.nil);
            }

            try vm.pushCall(call_frame);

            continue :next readOpcode(&ip);
        },

        .system => {
            try vm.consumeFuel(&fuel, 1);

            const index = read(u8, &ip);
            const arg_count = read(u8, &ip);
            const system_fn = system.fns[index];

            log.debug("system index={} arg_count={} system_fn={p}", .{ index, arg_count, system_fn });

            vm.storeContext(.{ .fuel = fuel });
            const result = try system_fn(.{
                .vm = vm,
                .allocator = allocator,
                .args = vm.stack[vm.stack_top - arg_count .. vm.stack_top],
            });
            const context = vm.restoreContext();
            fuel = context.fuel;

            vm.stack_top -= arg_count;
            try vm.push(result);

            continue :next readOpcode(&ip);
        },

        .ret => {
            try vm.consumeFuel(&fuel, 1);

            const result = try vm.pop();
            const call_frame = try vm.popCall();

            try vm.validateBytecode(
                bottom <= vm.stack_top,
                "called function popped too many values. bottom={} stack_top={}",
                .{ bottom, vm.stack_top },
            );
            try vm.validateBytecode(
                call_bottom <= vm.call_stack_top,
                "called function popped too many call frames. call_bottom={} call_stack_top={}",
                .{ call_bottom, vm.call_stack_top },
            );

            vm.stack_top = bottom;
            try vm.push(result);

            if (vm.call_stack_top == call_bottom) {
                // If this is the last call frame from this `run` instance, return from the function.
                vm.storeContext(Context{ .fuel = fuel });
                break :next;
            }

            closure = call_frame.closure;
            ip = call_frame.ip;
            bottom = call_frame.bottom;

            continue :next readOpcode(&ip);
        },

        _ => |invalid_opcode| {
            try vm.consumeFuel(&fuel, 1);
            // NOTE: Not a validateBytecode call because this is zero-cost on the happy path,
            // so we don't need to disable it on release builds.
            return vm.throw("corrupted bytecode: invalid opcode {}", .{invalid_opcode});
        },
    }

    log.debug("END RUN", .{});
}

pub const Dotter = struct {
    from: value.Vec4,
    to: value.Vec4,
    num: f32,
};

/// NOTE: Assumes the value at the top is a dotter reticle.
pub fn runDotter(
    vm: *Vm,
    allocator: mem.Allocator,
    from: value.Vec4,
    to: value.Vec4,
    num: f32,
) Error!void {
    const reticle = try vm.pop();
    const draw = reticle.ref.reticle.dotter.draw; // parameter count checked on construction

    const data = allocator.dupe(Value, &[_]Value{
        .{ .vec4 = from },
        .{ .vec4 = to },
        .{ .number = num },
    }) catch return vm.outOfMemory();
    const ref = allocator.create(value.Ref) catch return vm.outOfMemory();
    ref.* = value.Ref{ .closure = .{
        .chunk = &system.record_dotter,
        .start = 0,
        .param_count = 1,
        .local_count = 0,
        .captures = data,
    } };

    const bottom = vm.stack_top;
    try vm.push(.{ .ref = ref });
    try vm.run(allocator, draw, bottom);
}