wasm3x 0.1.0

//! The [`Store`] owns all runtime state of a set of instantiated modules.

use std::os::raw::c_void;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use wasm3x_sys as ffi;

use crate::engine::Engine;
use crate::error::Error;
use crate::trampoline::HostFuncEntry;

/// A unique identifier used to ensure items are used with their owning store.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct StoreId(u64);

impl StoreId {
    /// A sentinel that matches no real store (ids start at 1). Used only as a
    /// defensive fallback when a store's call state is unexpectedly missing.
    pub(crate) const PLACEHOLDER: StoreId = StoreId(0);
}

fn next_store_id() -> StoreId {
    static NEXT: AtomicU64 = AtomicU64::new(1);
    StoreId(NEXT.fetch_add(1, Ordering::Relaxed))
}

/// Per-store state reachable from host-function trampolines via the runtime's
/// user data pointer. Boxed so that its address is stable.
pub(crate) struct CallState {
    /// The last error raised by a host function during a call, if any.
    pub(crate) last_host_error: Option<Error>,
    /// The owning store's identity, so host trampolines can build a [`Caller`]
    /// that validates handles.
    ///
    /// [`Caller`]: crate::Caller
    pub(crate) store_id: StoreId,
    /// Raw pointer to the store's host `data`, valid only while a call is in
    /// flight. Set by [`Store::set_call_data`] before entering Wasm and cleared
    /// afterwards; read by the trampoline to hand host functions a `&mut T`.
    pub(crate) data_ptr: *mut c_void,
}

/// Owns the runtime state (memory, globals, tables, compiled code) of the
/// modules instantiated into it, together with arbitrary host `data`.
///
/// A [`Store`] wraps a Wasm3 `IM3Runtime`. It is single-threaded and thus
/// neither [`Send`] nor [`Sync`].
///
/// The host `data` of type `T` is accessible directly via [`Store::data`]/
/// [`Store::data_mut`], and is also threaded into host functions: during a call
/// each host function receives a [`Caller<'_, T>`](crate::Caller) granting
/// `&T`/`&mut T` access to this `data` and to the runtime's linear memory. Host
/// functions still return at most one value (a Wasm3 signature limitation).
pub struct Store<T> {
    runtime: ffi::IM3Runtime,
    engine: Engine,
    id: StoreId,
    /// Boxed so the runtime's user-data pointer stays valid across moves.
    call_state: Box<CallState>,
    /// Backing byte buffers of instantiated modules. Wasm3 does not copy the
    /// Wasm bytes on parse, so they must outlive the runtime.
    module_bytes: Vec<Arc<[u8]>>,
    /// Host-function closures linked into this store, kept alive for the
    /// lifetime of the runtime that may call them.
    host_entries: Vec<Arc<HostFuncEntry>>,
    data: T,
}

impl<T> Store<T> {
    /// Creates a new [`Store`] with the given host `data`.
    pub fn new(engine: &Engine, data: T) -> Self {
        let id = next_store_id();
        let mut call_state = Box::new(CallState {
            last_host_error: None,
            store_id: id,
            data_ptr: core::ptr::null_mut(),
        });
        let user_data = (&mut *call_state) as *mut CallState as *mut c_void;
        // SAFETY: `engine.raw()` is a valid environment; user_data outlives the
        // runtime (both owned by this store, runtime freed first on drop).
        let runtime = unsafe {
            ffi::m3_NewRuntime(engine.raw(), engine.config().get_stack_size(), user_data)
        };
        assert!(!runtime.is_null(), "wasm3: failed to allocate runtime");
        Self {
            runtime,
            engine: engine.clone(),
            id,
            call_state,
            module_bytes: Vec::new(),
            host_entries: Vec::new(),
            data,
        }
    }

    /// Returns a shared reference to the host data.
    pub fn data(&self) -> &T {
        &self.data
    }

    /// Returns an exclusive reference to the host data.
    pub fn data_mut(&mut self) -> &mut T {
        &mut self.data
    }

    /// Returns the [`Engine`] this store belongs to.
    pub fn engine(&self) -> &Engine {
        &self.engine
    }

    pub(crate) fn raw(&self) -> ffi::IM3Runtime {
        self.runtime
    }

    pub(crate) fn id(&self) -> StoreId {
        self.id
    }

    pub(crate) fn register_bytes(&mut self, bytes: Arc<[u8]>) {
        self.module_bytes.push(bytes);
    }

    pub(crate) fn register_host_entry(&mut self, entry: Arc<HostFuncEntry>) {
        self.host_entries.push(entry);
    }

    /// The number of host-function closures retained by this store. Used by
    /// tests to assert that non-imported definitions are not accumulated.
    #[cfg(test)]
    pub(crate) fn host_entry_count(&self) -> usize {
        self.host_entries.len()
    }

    pub(crate) fn clear_host_error(&mut self) {
        self.call_state.last_host_error = None;
    }

    /// Publishes a pointer to the host `data` for the duration of a call so that
    /// host-function trampolines can materialize a `&mut T` through a [`Caller`].
    ///
    /// Only a raw pointer is stashed; no live borrow is held across the FFI
    /// call, and `self.data`'s address is stable while `&mut self` is held.
    ///
    /// [`Caller`]: crate::Caller
    pub(crate) fn set_call_data(&mut self) {
        let ptr = (&mut self.data) as *mut T as *mut c_void;
        self.call_state.data_ptr = ptr;
    }

    /// Clears the host-`data` pointer once a call has returned.
    pub(crate) fn clear_call_data(&mut self) {
        self.call_state.data_ptr = core::ptr::null_mut();
    }

    pub(crate) fn take_host_error(&mut self) -> Option<Error> {
        self.call_state.last_host_error.take()
    }

    /// Panics if `id` does not match this store, guarding against using an item
    /// (function, memory, ...) with the wrong store.
    pub(crate) fn ensure_owns(&self, id: StoreId) {
        assert!(
            self.id == id,
            "wasm3: attempted to use an item with a different `Store` than the one that owns it"
        );
    }
}

impl<T: Default> Default for Store<T> {
    fn default() -> Self {
        Self::new(&Engine::default(), T::default())
    }
}

/// Panics if a handle's `store_id` does not match the `store_id` of the context
/// it is being used with. The shared ownership guard used by every handle type.
pub(crate) fn assert_same_store(ctx: StoreId, handle: StoreId) {
    assert!(
        ctx == handle,
        "wasm3: attempted to use an item with a different `Store` than the one that owns it"
    );
}

/// A shared view of a [`Store`]'s runtime, produced by [`AsContext::as_context`].
///
/// Handle methods that only read take an `impl AsContext<Data = T>`, so they
/// accept a `&Store<T>`, a [`StoreContext`]/[`StoreContextMut`], or a
/// [`Caller`](crate::Caller) uniformly.
pub struct StoreContext<'a, T> {
    runtime: ffi::IM3Runtime,
    store_id: StoreId,
    data: &'a T,
}

/// An exclusive view of a [`Store`]'s runtime, produced by
/// [`AsContextMut::as_context_mut`]. Used by handle methods that mutate. See
/// [`AsContext`] for why this uniform view exists.
pub struct StoreContextMut<'a, T> {
    runtime: ffi::IM3Runtime,
    store_id: StoreId,
    data: &'a mut T,
}

impl<'a, T> StoreContext<'a, T> {
    pub(crate) fn new(runtime: ffi::IM3Runtime, store_id: StoreId, data: &'a T) -> Self {
        Self {
            runtime,
            store_id,
            data,
        }
    }

    /// Returns a shared reference to the host data.
    pub fn data(&self) -> &T {
        self.data
    }

    pub(crate) fn runtime(&self) -> ffi::IM3Runtime {
        self.runtime
    }

    pub(crate) fn store_id(&self) -> StoreId {
        self.store_id
    }
}

impl<'a, T> StoreContextMut<'a, T> {
    pub(crate) fn new(runtime: ffi::IM3Runtime, store_id: StoreId, data: &'a mut T) -> Self {
        Self {
            runtime,
            store_id,
            data,
        }
    }

    /// Returns a shared reference to the host data.
    pub fn data(&self) -> &T {
        self.data
    }

    /// Returns an exclusive reference to the host data.
    pub fn data_mut(&mut self) -> &mut T {
        self.data
    }

    pub(crate) fn runtime(&self) -> ffi::IM3Runtime {
        self.runtime
    }

    pub(crate) fn store_id(&self) -> StoreId {
        self.store_id
    }
}

/// A type that can yield a shared [`StoreContext`]: `Store<T>`, a context, or a
/// [`Caller`](crate::Caller), plus `&`/`&mut` references to any of them.
///
/// This trait is the uniform entry point for read-only handle methods
/// ([`Func::ty`], [`Memory::read`], [`Global::get`], [`Instance::get_func`], …):
/// taking `impl AsContext` lets each of them accept a `&Store<T>`, a
/// [`StoreContext`]/[`StoreContextMut`], or a [`Caller<'_, T>`] interchangeably.
///
/// [`Caller`] is what makes that uniformity load-bearing rather than cosmetic: a
/// host function holds a `Caller`, not a `&Store`, yet must still be able to read
/// linear memory and globals mid-call. Routing every handle method through
/// `AsContext` is what lets it do so with the same code the outside world uses.
///
/// The borrowed `data` rides along for the few methods that expose store data
/// (e.g. [`StoreContext::data`]). Pure read/write handle operations only need the
/// underlying runtime and store id and never touch `data`; it is carried for API
/// symmetry with `wasmtime`, not because the operation requires it.
///
/// [`Caller<'_, T>`]: crate::Caller
/// [`Caller`]: crate::Caller
/// [`Func::ty`]: crate::Func::ty
/// [`Memory::read`]: crate::Memory::read
/// [`Global::get`]: crate::Global::get
/// [`Instance::get_func`]: crate::Instance::get_func
pub trait AsContext {
    /// The host data type of the underlying store.
    type Data;
    /// Returns a shared view of the underlying store.
    fn as_context(&self) -> StoreContext<'_, Self::Data>;
}

/// A type that can yield an exclusive [`StoreContextMut`].
///
/// The mutating counterpart to [`AsContext`], taken by handle methods that write
/// ([`Memory::write`], [`Global::set`], …) so they too accept a `&mut Store<T>`,
/// a [`StoreContextMut`], or a [`Caller<'_, T>`](crate::Caller) interchangeably.
///
/// [`Memory::write`]: crate::Memory::write
/// [`Global::set`]: crate::Global::set
pub trait AsContextMut: AsContext {
    /// Returns an exclusive view of the underlying store.
    fn as_context_mut(&mut self) -> StoreContextMut<'_, Self::Data>;
}

impl<T> AsContext for Store<T> {
    type Data = T;
    fn as_context(&self) -> StoreContext<'_, T> {
        StoreContext {
            runtime: self.runtime,
            store_id: self.id,
            data: &self.data,
        }
    }
}

impl<T> AsContextMut for Store<T> {
    fn as_context_mut(&mut self) -> StoreContextMut<'_, T> {
        StoreContextMut {
            runtime: self.runtime,
            store_id: self.id,
            data: &mut self.data,
        }
    }
}

impl<C: AsContext + ?Sized> AsContext for &C {
    type Data = C::Data;
    fn as_context(&self) -> StoreContext<'_, C::Data> {
        C::as_context(self)
    }
}

impl<C: AsContext + ?Sized> AsContext for &mut C {
    type Data = C::Data;
    fn as_context(&self) -> StoreContext<'_, C::Data> {
        C::as_context(self)
    }
}

impl<C: AsContextMut + ?Sized> AsContextMut for &mut C {
    fn as_context_mut(&mut self) -> StoreContextMut<'_, C::Data> {
        C::as_context_mut(self)
    }
}

impl<'a, T> AsContext for StoreContext<'a, T> {
    type Data = T;
    fn as_context(&self) -> StoreContext<'_, T> {
        StoreContext {
            runtime: self.runtime,
            store_id: self.store_id,
            data: self.data,
        }
    }
}

impl<'a, T> AsContext for StoreContextMut<'a, T> {
    type Data = T;
    fn as_context(&self) -> StoreContext<'_, T> {
        StoreContext {
            runtime: self.runtime,
            store_id: self.store_id,
            data: &*self.data,
        }
    }
}

impl<'a, T> AsContextMut for StoreContextMut<'a, T> {
    fn as_context_mut(&mut self) -> StoreContextMut<'_, T> {
        StoreContextMut {
            runtime: self.runtime,
            store_id: self.store_id,
            data: &mut *self.data,
        }
    }
}

impl<T> Drop for Store<T> {
    fn drop(&mut self) {
        // Free the runtime first: this releases all loaded modules and compiled
        // code that reference `module_bytes` and `host_entries`, which are
        // dropped afterwards when the remaining fields drop.
        //
        // SAFETY: `runtime` was created by `m3_NewRuntime` and freed once.
        unsafe { ffi::m3_FreeRuntime(self.runtime) }
    }
}