wasm3x 0.1.0

//! Callable Wasm functions: the dynamically typed [`Func`] and the statically
//! typed [`TypedFunc`], plus the traits that drive host-function wrapping.

use core::marker::PhantomData;
use std::ffi::CString;
use std::os::raw::c_void;

use wasm3x_sys as ffi;

use crate::caller::Caller;
use crate::error::{Error, Result};
use crate::store::{AsContext, Store, StoreId, assert_same_store};
use crate::trampoline::HostRawFn;
use crate::value::{FuncType, Val, ValType};

/// The maximum arity supported by the statically typed call path.
const MAX_ARITY: usize = 16;

/// A callable Wasm function belonging to a [`Store`].
///
/// [`Func`] is a small copyable handle; the function itself is owned by the
/// store. Using a [`Func`] with a different store than the one it came from
/// panics.
#[derive(Debug, Clone, Copy)]
pub struct Func {
    raw: ffi::IM3Function,
    store_id: StoreId,
}

impl Func {
    /// Looks up an exported function by name in the given runtime.
    pub(crate) fn find_raw(
        runtime: ffi::IM3Runtime,
        store_id: StoreId,
        name: &str,
    ) -> Option<Func> {
        let cname = CString::new(name).ok()?;
        // SAFETY: valid runtime and nul-terminated name.
        unsafe {
            let mut raw: ffi::IM3Function = core::ptr::null_mut();
            let result = ffi::m3_FindFunction(&mut raw, runtime, cname.as_ptr());
            if !result.is_null() || raw.is_null() {
                return None;
            }
            Some(Func { raw, store_id })
        }
    }

    /// Reads the function signature directly from Wasm3 (needs no store access).
    ///
    /// Fails if any parameter or result type is not representable through this
    /// wrapper (e.g. `v128`), mirroring [`Global::ty`](crate::Global::ty) rather
    /// than panicking.
    fn signature(&self) -> Result<FuncType> {
        // SAFETY: `raw` is a valid function handle for the lifetime tied to its
        // store, which the caller keeps borrowed.
        unsafe {
            let num_args = ffi::m3_GetArgCount(self.raw);
            let num_results = ffi::m3_GetRetCount(self.raw);
            let mut params = Vec::with_capacity(num_args as usize);
            for i in 0..num_args {
                params.push(
                    ValType::from_ffi(ffi::m3_GetArgType(self.raw, i)).ok_or_else(|| {
                        Error::mismatch("function has an unsupported parameter type")
                    })?,
                );
            }
            let mut results = Vec::with_capacity(num_results as usize);
            for i in 0..num_results {
                results.push(
                    ValType::from_ffi(ffi::m3_GetRetType(self.raw, i)).ok_or_else(|| {
                        Error::mismatch("function has an unsupported result type")
                    })?,
                );
            }
            Ok(FuncType::new(params, results))
        }
    }

    /// Returns the [`FuncType`] of this function.
    ///
    /// Fails if the function's signature uses a type not representable through
    /// this wrapper (e.g. `v128`).
    pub fn ty(&self, store: impl AsContext) -> Result<FuncType> {
        assert_same_store(store.as_context().store_id(), self.store_id);
        self.signature()
    }

    /// Calls this function with the given `params`, writing the outputs into
    /// `results`.
    ///
    /// The number and types of `params` must match the function signature, and
    /// `results` must have exactly as many elements as the function has results.
    pub fn call<T>(&self, store: &mut Store<T>, params: &[Val], results: &mut [Val]) -> Result<()> {
        store.ensure_owns(self.store_id);
        let ty = self.signature()?;
        check_params(&ty, params)?;
        if results.len() != ty.results().len() {
            return Err(Error::mismatch(format!(
                "expected {} results, got a buffer for {}",
                ty.results().len(),
                results.len()
            )));
        }

        let arg_slots: Vec<u64> = params.iter().map(|v| v.to_slot()).collect();
        let mut arg_ptrs: Vec<*const c_void> = arg_slots
            .iter()
            .map(|slot| (slot as *const u64).cast())
            .collect();
        // SAFETY: `arg_slots`/`arg_ptrs` live across the call; argc matches.
        unsafe { raw_call(store, self.raw, params.len(), arg_ptrs.as_mut_ptr())? };

        let num_results = ty.results().len();
        let mut ret_slots = vec![0u64; num_results];
        // Pointers must carry mutable provenance: Wasm3 writes results through
        // them (it casts away `const` internally).
        let mut ret_ptrs: Vec<*const c_void> = ret_slots
            .iter_mut()
            .map(|slot| (slot as *mut u64) as *const c_void)
            .collect();
        // SAFETY: `ret_slots` has `num_results` writable 64-bit slots.
        unsafe {
            Error::from_ffi(ffi::m3_GetResults(
                self.raw,
                num_results as u32,
                ret_ptrs.as_mut_ptr(),
            ))?;
        }
        for (i, result_ty) in ty.results().iter().enumerate() {
            results[i] = Val::from_slot(*result_ty, ret_slots[i]);
        }
        Ok(())
    }

    /// Converts this function into a statically typed [`TypedFunc`], validating
    /// that `Params` and `Results` match the function signature.
    pub fn typed<Params, Results>(
        &self,
        store: impl AsContext,
    ) -> Result<TypedFunc<Params, Results>>
    where
        Params: WasmParams,
        Results: WasmResults,
    {
        self.typed_checked(store.as_context().store_id())
    }

    /// Validates `Params`/`Results` against the signature and produces a typed
    /// handle. Shared by [`Func::typed`] and [`Instance::get_typed_func`].
    ///
    /// [`Instance::get_typed_func`]: crate::Instance::get_typed_func
    pub(crate) fn typed_checked<Params, Results>(
        &self,
        store_id: StoreId,
    ) -> Result<TypedFunc<Params, Results>>
    where
        Params: WasmParams,
        Results: WasmResults,
    {
        assert!(
            self.store_id == store_id,
            "wasm3: attempted to use a function with a different `Store`"
        );
        let ty = self.signature()?;
        let mut expected_params = Vec::new();
        Params::types(&mut expected_params);
        let mut expected_results = Vec::new();
        Results::types(&mut expected_results);
        if expected_params.as_slice() != ty.params() {
            return Err(Error::mismatch(format!(
                "function parameters {:?} do not match requested {:?}",
                ty.params(),
                expected_params
            )));
        }
        if expected_results.as_slice() != ty.results() {
            return Err(Error::mismatch(format!(
                "function results {:?} do not match requested {:?}",
                ty.results(),
                expected_results
            )));
        }
        Ok(TypedFunc {
            func: *self,
            _marker: PhantomData,
        })
    }
}

/// A statically typed view of a [`Func`] with compile-time known parameter and
/// result types. Type checks happen once, at construction.
#[derive(Debug, Clone, Copy)]
pub struct TypedFunc<Params, Results> {
    func: Func,
    _marker: PhantomData<fn(Params) -> Results>,
}

impl<Params, Results> TypedFunc<Params, Results>
where
    Params: WasmParams,
    Results: WasmResults,
{
    /// Calls this function with the given typed parameters.
    ///
    /// Takes a `&mut Store<T>` rather than an `impl AsContextMut` on purpose: a
    /// host function only ever holds a [`Caller`], never a
    /// `&mut Store`, so this signature makes re-entering the same store from a
    /// host function a compile error. Wasm3's `m3_Call` runs from the base of
    /// the runtime stack and would corrupt the suspended outer frame.
    pub fn call<T>(&self, store: &mut Store<T>, params: Params) -> Result<Results> {
        store.ensure_owns(self.func.store_id);
        let raw = self.func.raw;

        let argc = Params::LEN;
        let mut arg_slots = [0u64; MAX_ARITY];
        params.write_slots(&mut arg_slots[..argc]);
        let mut arg_ptrs = [core::ptr::null::<c_void>(); MAX_ARITY];
        for i in 0..argc {
            arg_ptrs[i] = (&arg_slots[i] as *const u64).cast();
        }
        // SAFETY: scratch buffers live across the call; argc matches signature.
        unsafe { raw_call(store, raw, argc, arg_ptrs.as_mut_ptr())? };

        let retc = Results::LEN;
        let mut ret_slots = [0u64; MAX_ARITY];
        let mut ret_ptrs = [core::ptr::null::<c_void>(); MAX_ARITY];
        for i in 0..retc {
            // Mutable provenance: Wasm3 writes results through these pointers.
            ret_ptrs[i] = (&mut ret_slots[i] as *mut u64) as *const c_void;
        }
        // SAFETY: `ret_slots` has enough writable slots for `retc` results.
        unsafe {
            Error::from_ffi(ffi::m3_GetResults(raw, retc as u32, ret_ptrs.as_mut_ptr()))?;
        }
        Ok(Results::read_slots(&ret_slots[..retc]))
    }

    /// Returns the underlying dynamically typed [`Func`].
    pub fn func(&self) -> &Func {
        &self.func
    }
}

/// Validates that `params` match the parameter types of `ty`.
fn check_params(ty: &FuncType, params: &[Val]) -> Result<()> {
    if params.len() != ty.params().len() {
        return Err(Error::mismatch(format!(
            "expected {} parameters, got {}",
            ty.params().len(),
            params.len()
        )));
    }
    for (i, (param, expected)) in params.iter().zip(ty.params()).enumerate() {
        if param.ty() != *expected {
            return Err(Error::mismatch(format!(
                "parameter {i} has type {:?}, expected {:?}",
                param.ty(),
                expected
            )));
        }
    }
    Ok(())
}

/// Performs `m3_Call` and maps failures to rich errors (host trap or Wasm trap).
///
/// # Safety
///
/// `argptrs` must point to `argc` valid argument-value pointers that live for
/// the duration of the call.
unsafe fn raw_call<T>(
    store: &mut Store<T>,
    raw: ffi::IM3Function,
    argc: usize,
    argptrs: *mut *const c_void,
) -> Result<()> {
    unsafe {
        store.clear_host_error();
        // Publish the host-data pointer so host functions invoked during this
        // call can reach it through their `Caller`.
        store.set_call_data();
        let result = ffi::m3_Call(raw, argc as u32, argptrs);
        store.clear_call_data();
        if !result.is_null() {
            if let Some(error) = store.take_host_error() {
                return Err(error);
            }
            return Err(Error::from_trap(store.raw(), result));
        }
        Ok(())
    }
}

// ===========================================================================
// Typed value machinery
// ===========================================================================

/// A primitive Wasm type usable as a typed function parameter or result.
///
/// This trait is effectively sealed: it is implemented only for the primitive
/// Wasm types (`i32`, `i64`, `f32`, `f64`, and the unsigned integer aliases).
pub trait WasmTy: Copy + Send + Sync + 'static {
    #[doc(hidden)]
    const TYPE: ValType;
    #[doc(hidden)]
    fn into_slot(self) -> u64;
    #[doc(hidden)]
    fn from_slot(slot: u64) -> Self;
    #[doc(hidden)]
    fn into_val(self) -> Val;
}

macro_rules! impl_wasm_ty {
    ($ty:ty, $vt:expr, $into_slot:expr, $from_slot:expr, $into_val:expr) => {
        impl WasmTy for $ty {
            const TYPE: ValType = $vt;
            #[inline]
            fn into_slot(self) -> u64 {
                let f: fn($ty) -> u64 = $into_slot;
                f(self)
            }
            #[inline]
            fn from_slot(slot: u64) -> Self {
                let f: fn(u64) -> $ty = $from_slot;
                f(slot)
            }
            #[inline]
            fn into_val(self) -> Val {
                let f: fn($ty) -> Val = $into_val;
                f(self)
            }
        }
    };
}

impl_wasm_ty!(
    i32,
    ValType::I32,
    |v| v as u32 as u64,
    |s| s as u32 as i32,
    Val::I32
);
impl_wasm_ty!(u32, ValType::I32, |v| v as u64, |s| s as u32, |v| Val::I32(
    v as i32
));
impl_wasm_ty!(i64, ValType::I64, |v| v as u64, |s| s as i64, Val::I64);
impl_wasm_ty!(u64, ValType::I64, |v| v, |s| s, |v| Val::I64(v as i64));
impl_wasm_ty!(
    f32,
    ValType::F32,
    |v| v.to_bits() as u64,
    |s| f32::from_bits(s as u32),
    Val::F32
);
impl_wasm_ty!(f64, ValType::F64, |v| v.to_bits(), f64::from_bits, Val::F64);

/// Types usable as the parameter list of a [`TypedFunc`] or host function.
pub trait WasmParams {
    #[doc(hidden)]
    const LEN: usize;
    #[doc(hidden)]
    fn write_slots(self, slots: &mut [u64]);
    #[doc(hidden)]
    fn types(out: &mut Vec<ValType>);
}

/// Types usable as the result list of a [`TypedFunc`] or host function.
pub trait WasmResults: Sized {
    #[doc(hidden)]
    const LEN: usize;
    #[doc(hidden)]
    fn read_slots(slots: &[u64]) -> Self;
    #[doc(hidden)]
    fn write_slots(self, slots: &mut [u64]);
    #[doc(hidden)]
    fn write_vals(self, out: &mut [Val]);
    #[doc(hidden)]
    fn types(out: &mut Vec<ValType>);
}

// A single primitive acts as a 1-element parameter / result list.
macro_rules! impl_single {
    ($ty:ty) => {
        impl WasmParams for $ty {
            const LEN: usize = 1;
            fn write_slots(self, slots: &mut [u64]) {
                slots[0] = WasmTy::into_slot(self);
            }
            fn types(out: &mut Vec<ValType>) {
                out.push(<$ty as WasmTy>::TYPE);
            }
        }
        impl WasmResults for $ty {
            const LEN: usize = 1;
            fn read_slots(slots: &[u64]) -> Self {
                <$ty as WasmTy>::from_slot(slots[0])
            }
            fn write_slots(self, slots: &mut [u64]) {
                slots[0] = WasmTy::into_slot(self);
            }
            fn write_vals(self, out: &mut [Val]) {
                out[0] = WasmTy::into_val(self);
            }
            fn types(out: &mut Vec<ValType>) {
                out.push(<$ty as WasmTy>::TYPE);
            }
        }
    };
}
impl_single!(i32);
impl_single!(u32);
impl_single!(i64);
impl_single!(u64);
impl_single!(f32);
impl_single!(f64);

impl WasmParams for () {
    const LEN: usize = 0;
    fn write_slots(self, _slots: &mut [u64]) {}
    fn types(_out: &mut Vec<ValType>) {}
}
impl WasmResults for () {
    const LEN: usize = 0;
    fn read_slots(_slots: &[u64]) -> Self {}
    fn write_slots(self, _slots: &mut [u64]) {}
    fn write_vals(self, _out: &mut [Val]) {}
    fn types(_out: &mut Vec<ValType>) {}
}

macro_rules! impl_tuple {
    ($n:expr; $($T:ident => $idx:tt),+) => {
        impl<$($T: WasmTy),+> WasmParams for ($($T,)+) {
            const LEN: usize = $n;
            fn write_slots(self, slots: &mut [u64]) {
                $( slots[$idx] = WasmTy::into_slot(self.$idx); )+
            }
            fn types(out: &mut Vec<ValType>) {
                $( out.push(<$T as WasmTy>::TYPE); )+
            }
        }
        impl<$($T: WasmTy),+> WasmResults for ($($T,)+) {
            const LEN: usize = $n;
            fn read_slots(slots: &[u64]) -> Self {
                ( $( <$T as WasmTy>::from_slot(slots[$idx]), )+ )
            }
            fn write_slots(self, slots: &mut [u64]) {
                $( slots[$idx] = WasmTy::into_slot(self.$idx); )+
            }
            fn write_vals(self, out: &mut [Val]) {
                $( out[$idx] = WasmTy::into_val(self.$idx); )+
            }
            fn types(out: &mut Vec<ValType>) {
                $( out.push(<$T as WasmTy>::TYPE); )+
            }
        }
    };
}
impl_tuple!(1; T0 => 0);
impl_tuple!(2; T0 => 0, T1 => 1);
impl_tuple!(3; T0 => 0, T1 => 1, T2 => 2);
impl_tuple!(4; T0 => 0, T1 => 1, T2 => 2, T3 => 3);
impl_tuple!(5; T0 => 0, T1 => 1, T2 => 2, T3 => 3, T4 => 4);
impl_tuple!(6; T0 => 0, T1 => 1, T2 => 2, T3 => 3, T4 => 4, T5 => 5);
impl_tuple!(7; T0 => 0, T1 => 1, T2 => 2, T3 => 3, T4 => 4, T5 => 5, T6 => 6);
impl_tuple!(8; T0 => 0, T1 => 1, T2 => 2, T3 => 3, T4 => 4, T5 => 5, T6 => 6, T7 => 7);

/// Types usable as the return value of a host function: a value, a tuple of
/// values, or a [`Result`] thereof.
pub trait WasmRet {
    #[doc(hidden)]
    type Output: WasmResults;
    #[doc(hidden)]
    fn into_result(self) -> Result<Self::Output>;
}

impl<T: WasmResults> WasmRet for T {
    type Output = T;
    fn into_result(self) -> Result<Self::Output> {
        Ok(self)
    }
}

impl<T: WasmResults> WasmRet for Result<T> {
    type Output = T;
    fn into_result(self) -> Result<Self::Output> {
        self
    }
}

/// Rust closures and functions usable as Wasm3 host functions.
///
/// The callable's remaining parameters and its return type map to Wasm value
/// types. It may optionally take a leading [`Caller<'_, T>`](crate::Caller),
/// giving access to the store's host data and linear memory; host functions that
/// do not need the caller can simply omit it (e.g. `|x: i32| x * 2` or a plain
/// `fn() -> u64`).
pub trait IntoFunc<T, Params, Results>: Send + Sync + 'static {
    #[doc(hidden)]
    fn into_host_func(self) -> (FuncType, HostRawFn);
}

macro_rules! impl_into_func {
    ($($P:ident),*) => {
        // Closures/functions that take a leading `Caller<'_, T>`. The `Caller` is
        // encoded into the `Params` marker tuple so this impl does not overlap the
        // no-`Caller` impl below (which uses `($($P,)*)`).
        impl<'a, T, F, $($P,)* R> IntoFunc<T, (Caller<'a, T>, $($P,)*), R> for F
        where
            F: Fn(Caller<'_, T>, $($P),*) -> R + Send + Sync + 'static,
            $($P: WasmTy,)*
            R: WasmRet,
        {
            #[allow(non_snake_case, unused_mut, unused_variables, unused_assignments)]
            fn into_host_func(self) -> (FuncType, HostRawFn) {
                let ty = impl_into_func!(@ty ($($P,)*) R);
                // Slot-native closure: read/write the raw Wasm3 stack directly.
                // The stack is laid out `[results.., args..]`; results occupy the
                // leading `num_results` slots, args follow.
                let host: HostRawFn = Box::new(move |raw, sp: *mut u64| -> Result<()> {
                    // `T` is baked in here because `Linker<T>` monomorphized this
                    // closure; rebuild the typed `Caller` from the erased bundle.
                    let caller: Caller<'_, T> = Caller::from_raw(raw);
                    let num_results = <R::Output as WasmResults>::LEN;
                    impl_into_func!(@read_args num_results, sp, $($P,)*);
                    let ret = (self)(caller, $($P),*);
                    impl_into_func!(@write_results num_results, sp, ret)
                });
                (ty, host)
            }
        }

        // Closures/functions that ignore the caller: identical marshalling, but
        // the `RawCaller` is dropped and no `Caller` is built.
        impl<T, F, $($P,)* R> IntoFunc<T, ($($P,)*), R> for F
        where
            F: Fn($($P),*) -> R + Send + Sync + 'static,
            $($P: WasmTy,)*
            R: WasmRet,
        {
            #[allow(non_snake_case, unused_mut, unused_variables, unused_assignments)]
            fn into_host_func(self) -> (FuncType, HostRawFn) {
                let ty = impl_into_func!(@ty ($($P,)*) R);
                let host: HostRawFn = Box::new(move |_raw, sp: *mut u64| -> Result<()> {
                    let num_results = <R::Output as WasmResults>::LEN;
                    impl_into_func!(@read_args num_results, sp, $($P,)*);
                    let ret = (self)($($P),*);
                    impl_into_func!(@write_results num_results, sp, ret)
                });
                (ty, host)
            }
        }
    };

    // Builds the `FuncType` from the parameter types and the return type.
    (@ty ($($P:ident,)*) $R:ident) => {{
        let params: Vec<ValType> = vec![$(<$P as WasmTy>::TYPE,)*];
        let mut results: Vec<ValType> = Vec::new();
        <$R::Output as WasmResults>::types(&mut results);
        FuncType::new(params, results)
    }};

    // Reads every arg into a local named after its type param before any result is
    // written: args (`sp[num_results..]`) and results (`sp[..num_results]`) alias
    // the same region on the raw stack.
    (@read_args $num_results:ident, $sp:ident, $($P:ident,)*) => {
        let mut idx = 0usize;
        $(
            // SAFETY: Wasm3 provides a stack with a slot for each signature entry;
            // args start after the result slots.
            let $P = unsafe { <$P as WasmTy>::from_slot(*$sp.add($num_results + idx)) };
            idx += 1;
        )*
    };

    // Writes the closure's result back into the leading `num_results` slots.
    (@write_results $num_results:ident, $sp:ident, $ret:ident) => {{
        let ok = $ret.into_result()?;
        // SAFETY: the leading `num_results` slots are writable.
        let out = unsafe { core::slice::from_raw_parts_mut($sp, $num_results) };
        ok.write_slots(out);
        Ok(())
    }};
}
impl_into_func!();
impl_into_func!(T0);
impl_into_func!(T0, T1);
impl_into_func!(T0, T1, T2);
impl_into_func!(T0, T1, T2, T3);
impl_into_func!(T0, T1, T2, T3, T4);
impl_into_func!(T0, T1, T2, T3, T4, T5);
impl_into_func!(T0, T1, T2, T3, T4, T5, T6);
impl_into_func!(T0, T1, T2, T3, T4, T5, T6, T7);