wasm3x 0.1.0

//! The C-ABI trampoline that dispatches Wasm3 imports to Rust host functions.

use std::ffi::CStr;
use std::os::raw::c_void;
use std::panic::{self, AssertUnwindSafe};

use wasm3x_sys as ffi;

use crate::error::{Error, Result};
use crate::store::{CallState, StoreId};
use crate::value::{FuncType, Val};

/// The type-erased context the trampoline hands to a host closure. The linking
/// closure (monomorphized by `Linker<T>`) turns this back into a typed
/// [`Caller<'_, T>`](crate::Caller) via [`Caller::from_raw`](crate::Caller).
///
/// This is nominally `pub` (with private fields) only so it can appear in the
/// `#[doc(hidden)]` return type of [`IntoFunc::into_host_func`](crate::IntoFunc);
/// it is opaque and not re-exported.
#[doc(hidden)]
pub struct RawCaller {
    pub(crate) runtime: ffi::IM3Runtime,
    pub(crate) store_id: StoreId,
    pub(crate) data: *mut std::os::raw::c_void,
}

/// The dynamic host implementation used by [`func_new`](crate::Linker::func_new):
/// reads arguments and writes results through `&[Val]` slices.
pub(crate) type HostValFn =
    Box<dyn Fn(RawCaller, &[Val], &mut [Val]) -> Result<()> + Send + Sync + 'static>;

/// The slot-native host implementation used by
/// [`func_wrap`](crate::Linker::func_wrap): reads args and writes results
/// directly on the raw Wasm3 stack, using its compile-time-known arity. Args
/// live at `sp[num_results..]`, results at `sp[..num_results]`.
///
/// This is nominally `pub` (opaque) only so it can appear in the
/// `#[doc(hidden)]` return type of [`IntoFunc::into_host_func`](crate::IntoFunc),
/// like [`RawCaller`]; it is not re-exported.
#[doc(hidden)]
pub type HostRawFn = Box<dyn Fn(RawCaller, *mut u64) -> Result<()> + Send + Sync + 'static>;

/// The type-erased implementation of a host function. `func_wrap` produces the
/// allocation-free [`Raw`](HostImpl::Raw) variant; `func_new` needs the dynamic
/// [`Val`](HostImpl::Val) view.
pub(crate) enum HostImpl {
    Val(HostValFn),
    Raw(HostRawFn),
}

/// A host function definition kept alive by the [`Store`](crate::Store) it is
/// linked into. Its address (via `Arc`) is passed as Wasm3 user data.
pub(crate) struct HostFuncEntry {
    pub(crate) ty: FuncType,
    pub(crate) func: HostImpl,
}

/// The static trap message returned to Wasm3 when a host function fails. The
/// rich error is stored in the store's [`CallState`] and recovered afterwards.
static HOST_TRAP_MESSAGE: &CStr = c"wasm3: host function trapped";

/// Inline `Val` scratch capacity for the dynamic `func_new` path before falling
/// back to the heap. Sized to cover essentially all real host signatures; Wasm3
/// permits up to `d_m3MaxSaneFunctionArgRetCount` (1000) args+rets, which must
/// still work via the heap fallback. Deliberately independent of `func.rs`'s
/// `MAX_ARITY`, which sizes the unrelated typed-path slot buffers.
const INLINE_VAL_SLOTS: usize = 16;

/// The `M3RawCall` trampoline linked for every host function.
///
/// # ABI
///
/// Wasm3 lays out the raw stack as `[results.., args..]`, one 64-bit slot each
/// (64-bit slots are enabled in the sys build). We read the arguments after the
/// result slots and write the results into the leading slots.
///
/// # Safety
///
/// Invoked by Wasm3 with a valid import context whose `userdata` points to a
/// live [`HostFuncEntry`], a stack pointer with enough slots for the function's
/// signature, and a runtime whose user data is a valid [`CallState`].
pub(crate) unsafe extern "C" fn host_trampoline(
    runtime: ffi::IM3Runtime,
    ctx: ffi::IM3ImportContext,
    sp: *mut u64,
    _mem: *mut c_void,
) -> *const c_void {
    unsafe {
        // Recover the per-store call state (store identity + host-data pointer)
        // so the host function can be given a `Caller`.
        let call_state = ffi::m3_GetUserData(runtime) as *const CallState;
        let (store_id, data_ptr) = if call_state.is_null() {
            (StoreId::PLACEHOLDER, core::ptr::null_mut())
        } else {
            ((*call_state).store_id, (*call_state).data_ptr)
        };

        // Guard against unwinding across the FFI boundary, which would be UB.
        let outcome = panic::catch_unwind(AssertUnwindSafe(|| {
            let entry = &*((*ctx).userdata as *const HostFuncEntry);
            let caller = RawCaller {
                runtime,
                store_id,
                data: data_ptr,
            };
            match &entry.func {
                // Slot-native fast path: no `Vec`, no `Val`. The closure reads
                // args and writes results directly on `sp` using its own arity.
                HostImpl::Raw(func) => func(caller, sp),
                // Dynamic path: marshal the raw stack through `&[Val]`.
                HostImpl::Val(func) => {
                    let num_results = entry.ty.results().len();
                    let total = num_results + entry.ty.params().len();

                    // One combined buffer laid out like the raw stack
                    // (`[results.., args..]`), split into the two slices the
                    // closure expects. Zero-alloc for the common case; heap
                    // fallback for large dynamic sigs (Wasm3 permits up to
                    // `d_m3MaxSaneFunctionArgRetCount` = 1000 args+rets). `heap`
                    // is only touched on the fallback path, so the inline path
                    // stays alloc-free.
                    let mut inline = [Val::I32(0); INLINE_VAL_SLOTS];
                    let mut heap: Vec<Val>;
                    let buf: &mut [Val] = if total <= INLINE_VAL_SLOTS {
                        &mut inline[..total]
                    } else {
                        heap = vec![Val::I32(0); total];
                        heap.as_mut_slice()
                    };

                    let (results, args) = buf.split_at_mut(num_results);
                    // Read every arg before writing any result: args
                    // (`sp[num_results..]`) and results (`sp[..num_results]`)
                    // alias the same region on the raw stack.
                    for (i, (slot, ty)) in args.iter_mut().zip(entry.ty.params()).enumerate() {
                        *slot = Val::from_slot(*ty, *sp.add(num_results + i));
                    }
                    for (slot, ty) in results.iter_mut().zip(entry.ty.results()) {
                        *slot = Val::default_for_ty(*ty);
                    }

                    func(caller, args, results)?;

                    for (i, result) in results.iter().enumerate() {
                        *sp.add(i) = result.to_slot();
                    }
                    Ok(())
                }
            }
        }));

        match outcome {
            Ok(Ok(())) => core::ptr::null(),
            Ok(Err(error)) => stash_error(runtime, error),
            Err(_) => stash_error(runtime, Error::new("host function panicked")),
        }
    }
}

/// Records `error` in the store's call state and returns the static trap
/// pointer that signals a trap to Wasm3.
unsafe fn stash_error(runtime: ffi::IM3Runtime, error: Error) -> *const c_void {
    unsafe {
        let call_state = ffi::m3_GetUserData(runtime) as *mut CallState;
        if !call_state.is_null() {
            (*call_state).last_host_error = Some(error);
        }
        HOST_TRAP_MESSAGE.as_ptr() as *const c_void
    }
}