rkgk/crates/haku2/src/vm.zig

const std = @import("std");
const mem = std.mem;
const testAllocator = std.testing.allocator;

const bytecode = @import("bytecode.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_buffer: [1024]u8 = [_]u8{0} ** 1024, // buffer for exception message
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 {
    message: []const u8,
};

/// 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 {
    const message = std.fmt.bufPrint(vm.exception_buffer[0..], fmt, args) catch {
        vm.exception = .{ .message = "[exception message is too long; format string: " ++ fmt ++ "]" };
        return error.Exception;
    };
    vm.exception = .{ .message = message };
    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)", .{});
    }
    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;
    return vm.stack[vm.stack_top];
}

pub fn pushCall(vm: *Vm, frame: CallFrame) Error!void {
    if (vm.call_stack_top >= vm.call_stack.len) {
        return vm.throw("too much recursion", .{});
    }
    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", .{});
    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 {
    return @enumFromInt(read(u8, ip));
}

/// 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 {
    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];

            vm.storeContext(.{ .fuel = fuel });
            const result = try system_fn(.{
                .vm = vm,
                .allocator = allocator,
                .args = vm.stack[vm.stack_top - arg_count ..],
            });
            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});
        },
    }
}