rkgk/crates/haku2/src/lib.rs

use std::{
    alloc::{self, Layout},
    error::Error,
    fmt::{self, Display},
    marker::{PhantomData, PhantomPinned},
    ptr::{self, NonNull},
    slice,
};

use log::trace;

pub static WASM: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/zig-out/bin/haku2.wasm"));

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_alloc(size: usize, align: usize) -> *mut u8 {
    if let Ok(layout) = Layout::from_size_align(size, align) {
        alloc::alloc(layout)
    } else {
        ptr::null_mut()
    }
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_realloc(
    ptr: *mut u8,
    size: usize,
    align: usize,
    new_size: usize,
) -> *mut u8 {
    if let Ok(layout) = Layout::from_size_align(size, align) {
        alloc::realloc(ptr, layout, new_size)
    } else {
        ptr::null_mut()
    }
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_dealloc(ptr: *mut u8, size: usize, align: usize) {
    match Layout::from_size_align(size, align) {
        Ok(layout) => alloc::dealloc(ptr, layout),
        Err(_) => {
            log::error!("__haku2_dealloc: invalid layout size={size} align={align} ptr={ptr:?}")
        }
    }
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_log_err(
    scope: *const u8,
    scope_len: usize,
    msg: *const u8,
    len: usize,
) {
    let scope = String::from_utf8_lossy(slice::from_raw_parts(scope, scope_len));
    let msg = String::from_utf8_lossy(slice::from_raw_parts(msg, len));
    log::error!("{scope}: {msg}");
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_log_warn(
    scope: *const u8,
    scope_len: usize,
    msg: *const u8,
    len: usize,
) {
    let scope = String::from_utf8_lossy(slice::from_raw_parts(scope, scope_len));
    let msg = String::from_utf8_lossy(slice::from_raw_parts(msg, len));
    log::warn!("{scope}: {msg}");
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_log_info(
    scope: *const u8,
    scope_len: usize,
    msg: *const u8,
    len: usize,
) {
    let scope = String::from_utf8_lossy(slice::from_raw_parts(scope, scope_len));
    let msg = String::from_utf8_lossy(slice::from_raw_parts(msg, len));
    log::info!("{scope}: {msg}");
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_log_debug(
    scope: *const u8,
    scope_len: usize,
    msg: *const u8,
    len: usize,
) {
    let scope = String::from_utf8_lossy(slice::from_raw_parts(scope, scope_len));
    let msg = String::from_utf8_lossy(slice::from_raw_parts(msg, len));
    log::debug!("{scope}: {msg}");
}

#[repr(C)]
struct ScratchC {
    _data: (),
    _marker: PhantomData<(*mut u8, PhantomPinned)>,
}

#[repr(C)]
struct LimitsC {
    _data: (),
    _marker: PhantomData<(*mut u8, PhantomPinned)>,
}

#[repr(C)]
struct DefsC {
    _data: (),
    _marker: PhantomData<(*mut u8, PhantomPinned)>,
}

#[repr(C)]
struct VmC {
    _data: (),
    _marker: PhantomData<(*mut u8, PhantomPinned)>,
}

extern "C" {
    fn haku2_scratch_new(max: usize) -> *mut ScratchC;
    fn haku2_scratch_destroy(scratch: *mut ScratchC);
    fn haku2_scratch_reset(scratch: *mut ScratchC);

    fn haku2_limits_new() -> *mut LimitsC;
    fn haku2_limits_destroy(limits: *mut LimitsC);
    fn haku2_limits_set_stack_capacity(limits: *mut LimitsC, new: usize);
    fn haku2_limits_set_call_stack_capacity(limits: *mut LimitsC, new: usize);

    fn haku2_defs_parse(
        defs_string: *const u8,
        defs_len: usize,
        tags_string: *const u8,
        tags_len: usize,
    ) -> *mut DefsC;
    fn haku2_defs_destroy(defs: *mut DefsC);

    fn haku2_vm_new() -> *mut VmC;
    fn haku2_vm_destroy(vm: *mut VmC);
    fn haku2_vm_reset(
        vm: *mut VmC,
        s: *mut ScratchC,
        defs: *const DefsC,
        limits: *const LimitsC,
        fuel: u32,
    );
    fn haku2_vm_run_main(
        vm: *mut VmC,
        scratch: *mut ScratchC,
        code: *const u8,
        code_len: usize,
        local_count: u8,
    ) -> bool;
    fn haku2_vm_has_cont(vm: *const VmC) -> bool;
    fn haku2_vm_is_dotter(vm: *const VmC) -> bool;
    fn haku2_vm_run_dotter(
        vm: *mut VmC,
        scratch: *mut ScratchC,
        from_x: f32,
        from_y: f32,
        to_x: f32,
        to_y: f32,
        num: f32,
    ) -> bool;
    fn haku2_vm_exception_len(vm: *const VmC) -> usize;
    fn haku2_vm_exception_render(vm: *const VmC, buffer: *mut u8);

    // improper_ctypes is emitted for `*mut CanvasC`, which is an opaque {} on the Zig side and
    // therefore FFI-safe.
    #[expect(improper_ctypes)]
    fn haku2_render(vm: *mut VmC, canvas: *mut CanvasC, max_depth: usize) -> bool;
}

#[derive(Debug)]
pub struct Scratch {
    raw: NonNull<ScratchC>,
}

impl Scratch {
    pub fn new(max: usize) -> Scratch {
        // SAFETY: haku2_scratch_new does not have any safety invariants.
        let raw = NonNull::new(unsafe { haku2_scratch_new(max) }).expect("out of memory");
        trace!("Scratch::new -> {raw:?}");
        Scratch { raw }
    }

    pub fn reset(&mut self) {
        trace!("Scratch::reset({:?})", self.raw);
        // SAFETY: The pointer passed is non-null.
        unsafe {
            haku2_scratch_reset(self.raw.as_ptr());
        }
    }
}

impl Drop for Scratch {
    fn drop(&mut self) {
        trace!("Scratch::drop({:?})", self.raw);
        // SAFETY: The pointer passed is non-null.
        unsafe {
            haku2_scratch_destroy(self.raw.as_ptr());
        }
    }
}

#[derive(Debug, Clone, Copy)]
pub struct LimitsSpec {
    pub stack_capacity: usize,
    pub call_stack_capacity: usize,
}

#[derive(Debug)]
pub struct Limits {
    raw: NonNull<LimitsC>,
}

// SAFETY: Limits's backing storage is only modified on creation.
// Changing the limits requires creating a new instance.
unsafe impl Send for Limits {}
unsafe impl Sync for Limits {}

impl Limits {
    pub fn new(spec: LimitsSpec) -> Self {
        // SAFETY: haku2_limits_new has no safety invariants.
        let limits = NonNull::new(unsafe { haku2_limits_new() }).expect("out of memory");

        // SAFETY: The following functions are called on a valid pointer.
        unsafe {
            haku2_limits_set_stack_capacity(limits.as_ptr(), spec.stack_capacity);
            haku2_limits_set_call_stack_capacity(limits.as_ptr(), spec.call_stack_capacity);
        }

        Self { raw: limits }
    }
}

impl Drop for Limits {
    fn drop(&mut self) {
        // SAFETY: The pointer passed is non-null.
        unsafe {
            haku2_limits_destroy(self.raw.as_ptr());
        }
    }
}

#[derive(Debug)]
pub struct Defs {
    raw: NonNull<DefsC>,
}

// SAFETY: Defs' backing storage is not modified after creation.
unsafe impl Send for Defs {}
unsafe impl Sync for Defs {}

impl Defs {
    pub fn parse(defs: &str, tags: &str) -> Self {
        Self {
            raw: NonNull::new(unsafe {
                haku2_defs_parse(defs.as_ptr(), defs.len(), tags.as_ptr(), tags.len())
            })
            .expect("out of memory"),
        }
    }
}

impl Drop for Defs {
    fn drop(&mut self) {
        // SAFETY: The pointer passed is non-null.
        unsafe {
            haku2_defs_destroy(self.raw.as_ptr());
        }
    }
}

#[derive(Debug)]
pub struct Code {
    defs: Defs,
    main_chunk: Vec<u8>,
    main_local_count: u8,
}

impl Code {
    /// Creates a new instance of `Code` from a valid vector of bytes.
    ///
    /// # Safety
    ///
    /// This does not perform any validation, and there is no way to perform such
    /// validation before constructing this. The bytecode must simply be valid, which is the case
    /// for bytecode emitted directly by the compiler.
    ///
    /// Untrusted bytecode should never ever be loaded under any circumstances.
    pub unsafe fn new(defs: Defs, main_chunk: Vec<u8>, main_local_count: u8) -> Self {
        Self {
            defs,
            main_chunk,
            main_local_count,
        }
    }
}

/// A VM that is ready to run and loaded with valid bytecode.
#[derive(Debug)]
pub struct Vm {
    scratch: Scratch,
    code: Code,
    limits: Limits,
    inner: VmInner,
}

#[derive(Debug)]
pub enum Cont<'vm> {
    None,
    Dotter(ContDotter<'vm>),
}

#[derive(Debug)]
pub struct ContDotter<'vm> {
    vm: &'vm mut Vm,
}

#[derive(Debug, Clone, Copy)]
pub struct Dotter {
    pub from: (f32, f32),
    pub to: (f32, f32),
    pub num: f32,
}

impl Vm {
    pub fn new(scratch: Scratch, code: Code, limits: Limits) -> Self {
        // SAFETY: haku2_vm_new cannot fail.
        // Do note that this returns an uninitialized VM, which must be reset before use.
        let raw = NonNull::new(unsafe { haku2_vm_new() }).expect("out of memory");
        trace!("Vm::new({scratch:?}, {code:?}, {limits:?}) -> {raw:?}");
        Self {
            // SAFETY:
            // - Ownership of scratch is passed to the VM, so the VM cannot outlive the scratch space.
            //   - The VM never gives you any references back, so this is safe to do.
            // - The other arguments are only borrowed immutably for construction.
            inner: VmInner { raw },
            scratch,
            code,
            limits,
        }
    }

    /// Begin running code. This makes the VM enter a "trampoline" state: after this call, you may
    /// proceed to call `cont` as many times as it returns a value other than [`Cont::None`].
    ///
    /// Calling `begin` again during this process will work correctly, and result in another
    /// continuation being stack on top of the old one---at the expense of a stack slot.
    pub fn begin(&mut self, fuel: u32) -> Result<(), Exception> {
        trace!("Vm::begin({self:?}, {fuel})");
        self.scratch.reset();
        let ok = unsafe {
            haku2_vm_reset(
                self.inner.raw.as_ptr(),
                self.scratch.raw.as_ptr(),
                self.code.defs.raw.as_ptr(),
                self.limits.raw.as_ptr(),
                fuel,
            );
            haku2_vm_run_main(
                self.inner.raw.as_ptr(),
                self.scratch.raw.as_ptr(),
                self.code.main_chunk.as_ptr(),
                self.code.main_chunk.len(),
                self.code.main_local_count,
            )
        };
        if ok {
            Ok(())
        } else {
            Err(self.exception().expect("missing exception after !ok"))
        }
    }

    /// Returns whether `cont()` can be called to run the next continuation.
    pub fn has_cont(&self) -> bool {
        unsafe { haku2_vm_has_cont(self.inner.raw.as_ptr()) }
    }

    fn is_dotter(&self) -> bool {
        // SAFETY: The pointer is valid.
        unsafe { haku2_vm_is_dotter(self.inner.raw.as_ptr()) }
    }

    /// Returns how the VM should continue executing after the previous execution.
    pub fn cont(&mut self) -> Cont<'_> {
        match () {
            _ if self.is_dotter() => Cont::Dotter(ContDotter { vm: self }),
            _ => Cont::None,
        }
    }

    /// Renders the current scribble on top of the stack.
    /// If the value on top is not a scribble, throws an exception.
    ///
    /// The rendering is performed by calling into the [`Canvas`] trait.
    pub fn render(&mut self, canvas: &mut dyn Canvas, max_depth: usize) -> Result<(), Exception> {
        let mut wrapped = CanvasC { inner: canvas };
        let ok = unsafe { haku2_render(self.inner.raw.as_ptr(), &mut wrapped, max_depth) };
        if ok {
            Ok(())
        } else {
            Err(self.exception().expect("missing exception after !ok"))
        }
    }

    /// Render the current exception out to a string.
    /// Returns `None` if there's no exception.
    pub fn exception(&self) -> Option<Exception> {
        // SAFETY: The pointer passed to this function is valid.
        let len = unsafe { haku2_vm_exception_len(self.inner.raw.as_ptr()) };
        if len == 0 {
            return None;
        }

        let mut buffer = vec![0; len];
        // SAFETY: The length of the buffer is as indicated by haku2_vm_exception_len.
        unsafe {
            haku2_vm_exception_render(self.inner.raw.as_ptr(), buffer.as_mut_ptr());
        }
        Some(Exception {
            message: String::from_utf8_lossy(&buffer).into_owned(),
        })
    }

    /// Take the `Scratch` out of the VM for reuse in another one.
    /// The scratch memory will be reset (no bytes will be consumed.)
    pub fn into_scratch(self) -> Scratch {
        trace!("Vm::into_scratch({self:?})");
        let Vm {
            mut scratch,
            code: _,
            inner: _,
            limits: _,
        } = self;
        scratch.reset();
        scratch
    }
}

impl ContDotter<'_> {
    pub fn run(self, dotter: &Dotter) -> Result<(), Exception> {
        trace!("ContDotter::run({self:?}, {dotter:?})");

        let Dotter {
            from: (from_x, from_y),
            to: (to_x, to_y),
            num,
        } = *dotter;

        let ok = unsafe {
            haku2_vm_run_dotter(
                self.vm.inner.raw.as_ptr(),
                self.vm.scratch.raw.as_ptr(),
                from_x,
                from_y,
                to_x,
                to_y,
                num,
            )
        };
        if ok {
            Ok(())
        } else {
            Err(self.vm.exception().expect("missing exception after !ok"))
        }
    }
}

#[derive(Debug)]
struct VmInner {
    raw: NonNull<VmC>,
}

impl Drop for VmInner {
    fn drop(&mut self) {
        trace!("VmInner::drop({:?})", self.raw);
        // SAFETY: The pointer passed is non-null.
        unsafe {
            haku2_vm_destroy(self.raw.as_ptr());
        }
    }
}

#[derive(Debug)]
pub struct Exception {
    message: String,
}

impl Display for Exception {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.message)
    }
}

impl Error for Exception {}

/// Marker for the VM to indicate that the rendering did not go down correctly.
/// If this is encountered, it throws an exception and aborts rendering.
#[derive(Debug)]
pub struct RenderError;

pub trait Canvas {
    fn begin(&mut self) -> Result<(), RenderError>;
    fn line(&mut self, x1: f32, y1: f32, x2: f32, y2: f32) -> Result<(), RenderError>;
    fn rectangle(&mut self, x: f32, y: f32, width: f32, height: f32) -> Result<(), RenderError>;
    fn circle(&mut self, x: f32, y: f32, r: f32) -> Result<(), RenderError>;
    fn fill(&mut self, r: u8, g: u8, b: u8, a: u8) -> Result<(), RenderError>;
    fn stroke(&mut self, r: u8, g: u8, b: u8, a: u8, thickness: f32) -> Result<(), RenderError>;
}

// SAFETY NOTE: I'm not sure the ownership model for this is quite correct.
// Given how the &mut's ownership flows through the Zig side of the code, it _should_ be fine,
// but I'm not an unsafe code expert to say this is the case for sure.

#[repr(C)]
struct CanvasC<'a> {
    inner: &'a mut dyn Canvas,
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_canvas_begin(c: *mut CanvasC) -> bool {
    let c = &mut *c;
    c.inner.begin().is_ok()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_canvas_line(
    c: *mut CanvasC,
    x1: f32,
    y1: f32,
    x2: f32,
    y2: f32,
) -> bool {
    let c = &mut *c;
    c.inner.line(x1, y1, x2, y2).is_ok()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_canvas_rectangle(
    c: *mut CanvasC,
    x: f32,
    y: f32,
    width: f32,
    height: f32,
) -> bool {
    let c = &mut *c;
    c.inner.rectangle(x, y, width, height).is_ok()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_canvas_circle(c: *mut CanvasC, x: f32, y: f32, r: f32) -> bool {
    let c = &mut *c;
    c.inner.circle(x, y, r).is_ok()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_canvas_fill(c: *mut CanvasC, r: u8, g: u8, b: u8, a: u8) -> bool {
    let c = &mut *c;
    c.inner.fill(r, g, b, a).is_ok()
}

#[unsafe(no_mangle)]
unsafe extern "C" fn __haku2_canvas_stroke(
    c: *mut CanvasC,
    r: u8,
    g: u8,
    b: u8,
    a: u8,
    thickness: f32,
) -> bool {
    let c = &mut *c;
    c.inner.stroke(r, g, b, a, thickness).is_ok()
}