submilli-wasm 0.1.11

//! `StoreInner` — the non-generic entity storage the runtime operates on, plus a
//! simple index arena. Public handles (`Memory`/`Global`/… ) are indices into these.

use std::num::NonZeroU64;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;

use crate::canon::CanonicalTypeId;
use crate::engine::Engine;
use crate::extern_::{Global, Memory, Table, Tag};
use crate::func::Func;
use crate::instance::Instance;
use crate::value::{ExnRef, Rooted};

use super::arena::Arena;
use super::entity::ExternRefs;
use super::gc::GcHeap;
use super::reclaim::ReclaimArena;
use super::{
    ExnEntity, FuncEntity, GlobalEntity, InstanceEntity, MemoryEntity, TableEntity, TagEntity,
};

#[cfg(feature = "async")]
#[derive(Debug)]
pub(crate) struct AsyncCallBoundary {
    pub(crate) value_base: usize,
    pub(crate) stop_depth: usize,
}

#[cfg(feature = "async")]
#[derive(Debug)]
struct AsyncCallCleanup {
    cancelled: Arc<AtomicBool>,
    roots_mark: Option<usize>,
    pending: Option<Rooted<ExnRef>>,
    pending_generation: u64,
    boundary: Option<AsyncCallBoundary>,
}

/// Outcome of charging one unit of fuel (see [`StoreInner::consume_fuel_step`]).
pub(crate) enum FuelStep {
    /// Fuel was available and charged; keep running.
    Ran,
    /// Active fuel is exhausted but reserve remains: the async driver should yield
    /// to the executor, then [`refuel_from_reserve`](StoreInner::refuel_from_reserve).
    NeedYield,
    /// No fuel left at all → `Trap::OutOfFuel`.
    Exhausted,
}

/// Process-unique store id generator. `0` is reserved as the **unchecked** sentinel (funcref *values*
/// rebuilt from a bare operand-cell index carry it — they are same-store by construction), so ids
/// start at `1`. See [`StoreInner::check_handle`] (#34).
static NEXT_STORE_ID: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);

/// Non-generic storage for all of a store's runtime entities.
#[derive(Debug)]
pub(crate) struct StoreInner {
    /// Process-unique id stamped onto every handle this store mints (#34). A handle whose `store`
    /// doesn't match this is from another store — using it here is an embedder bug, caught by
    /// [`check_handle`](Self::check_handle) rather than silently touching the wrong entity.
    store_id: u64,
    engine: Engine,
    pub(super) funcs: Arena<FuncEntity>,
    pub(super) memories: Arena<MemoryEntity>,
    pub(super) tables: Arena<TableEntity>,
    pub(super) globals: Arena<GlobalEntity>,
    pub(super) tags: Arena<TagEntity>,
    /// Reused host-call arg/result buffers: taken per call and returned after, so the
    /// steady-state boundary is allocation-free (`exec::host::invoke_host`).
    host_scratch: (Vec<crate::value::Val>, Vec<crate::value::Val>),
    /// Exception instances backing `exnref` values; `Rooted<ExnRef>` indexes here. Grow-only
    /// (freed on store drop); its `args` are GC roots the tracing collector (#27g) enumerates
    /// transitively (the exn arena itself is not reclaimed yet — see `gc_collect`).
    pub(super) exns: ReclaimArena<ExnEntity>,
    pub(super) instances: Arena<InstanceEntity>,
    /// Host payloads backing `externref` values; `Rooted<ExternRef>` indexes here. Traced for GC
    /// reachability; box reclamation is a follow-up (see `gc_collect`).
    pub(super) externrefs: ExternRefs,
    /// Managed `struct`/`array` objects; `Rooted<AnyRef>` slot handles index here. Reclaimed by the
    /// tracing collector (`Collector::Auto`/`MarkSweep`); allocate-only under `Collector::Null`.
    pub(crate) gc: GcHeap,
    /// Live host-held GC roots (`Rooted` handed to the embedder via the GC host API, scoped by
    /// `RootScope`). Enumerated at collection so a host reference held across a guest collection
    /// keeps its object alive (#27g). Each is `(handle, hierarchy)`.
    pub(super) gc_roots: std::sync::RwLock<Vec<(u32, crate::canon::RefKind)>>,
    /// Canonical type ids of host-allocated GC objects, each pinned with one type-registration
    /// (decref'd on store drop) so a host object outliving its `StructType` handle keeps its type
    /// (#27i; mirrors wasmtime's `gc_host_alloc_types`).
    pub(super) gc_host_alloc_types: std::collections::HashSet<CanonicalTypeId>,
    /// Active fuel — charged per op (meaningful only when `engine.consume_fuel()`).
    /// With an async yield interval this is the current slice; total = `fuel + fuel_reserve`.
    fuel: u64,
    /// Fuel held back from the active slice, released on each async yield.
    fuel_reserve: u64,
    /// Async fuel-yield granularity: yield to the executor every this-many units.
    fuel_yield_interval: Option<NonZeroU64>,
    /// Absolute epoch value at/after which execution interrupts (`u64::MAX` = none).
    epoch_deadline: u64,
    /// Exception surfaced from the last call / a host `throw`; taken via `take_pending_exception`.
    pub(super) pending_exception: Option<Rooted<ExnRef>>,
    /// Monotonic identity of the last exception installed in `pending_exception`. Host-call
    /// boundaries snapshot this so a stale pending exception cannot be paired with a later,
    /// unrelated `ThrownException` error.
    pub(super) pending_exception_generation: u64,
    /// First execution invariant failure observed by an infallible public operation.
    internal_error: Option<crate::error::InternalError>,
    /// Cleanup records for async host calls currently suspended at `.await`. A drop token marks a
    /// record cancelled; the next mutable store access restores it before exposing the store.
    #[cfg(feature = "async")]
    async_call_cleanups: Vec<AsyncCallCleanup>,
    /// The shared interpreter execution, parked for the duration of a host call so a host fn
    /// that re-enters wasm (`Func::call`) runs on the *same* operand/frame stacks (separated
    /// by a [`Delimiter`](crate::exec::Delimiter)); empty (default) outside a host call.
    exec_slot: crate::exec::Execution,
    /// This store's GC-request mailbox (the engine holds a `Weak`). The engine posts to it under
    /// engine-wide GC pressure; the run loop reads-and-clears it at a back-edge and self-collects.
    pub(super) gc_request: Arc<AtomicBool>,
}

impl StoreInner {
    pub(crate) fn new(engine: Engine) -> Self {
        let gc = GcHeap::new(engine.is_collecting(), engine.gc_heap_reservation());
        let gc_request = engine.register_gc_request();
        StoreInner {
            store_id: NEXT_STORE_ID.fetch_add(1, std::sync::atomic::Ordering::Relaxed),
            engine,
            funcs: Arena::default(),
            memories: Arena::default(),
            tables: Arena::default(),
            globals: Arena::default(),
            tags: Arena::default(),
            host_scratch: (Vec::new(), Vec::new()),
            exns: ReclaimArena::default(),
            instances: Arena::default(),
            externrefs: ExternRefs::default(),
            gc,
            gc_roots: std::sync::RwLock::new(Vec::new()),
            gc_host_alloc_types: std::collections::HashSet::new(),
            fuel: 0,
            fuel_reserve: 0,
            fuel_yield_interval: None,
            epoch_deadline: u64::MAX,
            pending_exception: None,
            pending_exception_generation: 0,
            internal_error: None,
            #[cfg(feature = "async")]
            async_call_cleanups: Vec::new(),
            exec_slot: crate::exec::Execution::default(),
            gc_request,
        }
    }

    // The GC-request mailbox (`take_gc_request`) and the pending-exception accessors
    // (`set`/`take_pending_exception`) live in `managed` with the rest of the GC machinery; their
    // fields (`gc_request`, `pending_exception`) stay on the struct above.

    /// Takes the parked execution (empty at a top-level entry = fresh stacks).
    pub(crate) fn take_exec(&mut self) -> crate::exec::Execution {
        std::mem::take(&mut self.exec_slot)
    }

    /// Parks the shared execution back in its slot (see `exec_slot`).
    pub(crate) fn park_exec(&mut self, exec: crate::exec::Execution) {
        self.exec_slot = exec;
    }

    /// Swap-based park/unpark for the per-host-call crossing (no placeholder writes/drops).
    /// Used symmetrically; a re-entrant `Func::call` takes/re-parks through the same slot,
    /// so the second swap always reclaims the (possibly updated) shared execution.
    #[inline]
    pub(crate) fn swap_exec(&mut self, exec: &mut crate::exec::Execution) {
        std::mem::swap(&mut self.exec_slot, exec);
    }

    /// The parked execution (frames back `WasmBacktrace::capture`); empty outside a host call.
    pub(crate) fn parked_exec(&self) -> &crate::exec::Execution {
        &self.exec_slot
    }

    /// The operand/local GC roots of the parked execution (empty when none is parked). Seeds a
    /// collection triggered from a host call (`gc_reserve_host`/`Store::gc`), where the guest's
    /// operands live on the parked execution rather than a running one.
    pub(crate) fn exec_roots(&self) -> crate::Result<Vec<(u32, crate::canon::RefKind)>> {
        self.exec_slot.operand_roots()
    }

    pub(crate) fn has_parked_execution(&self) -> bool {
        !self.exec_slot.is_empty()
    }

    pub(crate) fn latch_internal_error(&mut self, error: crate::error::InternalError) {
        if self.internal_error.is_none() {
            self.internal_error = Some(error);
        }
    }

    pub(crate) fn take_internal_error(&mut self) -> Option<crate::Error> {
        self.internal_error.take().map(Into::into)
    }

    pub(crate) fn internal_error(&self) -> Option<crate::Error> {
        self.internal_error.map(Into::into)
    }

    pub(crate) fn clear_internal_error(&mut self) {
        self.internal_error = None;
    }

    #[cfg(feature = "async")]
    pub(crate) fn begin_async_call_cleanup(
        &mut self,
        roots_mark: Option<usize>,
        pending: Option<Rooted<ExnRef>>,
        pending_generation: u64,
        boundary: Option<AsyncCallBoundary>,
    ) -> Arc<AtomicBool> {
        let cancelled = Arc::new(AtomicBool::new(false));
        self.async_call_cleanups.push(AsyncCallCleanup {
            cancelled: cancelled.clone(),
            roots_mark,
            pending,
            pending_generation,
            boundary,
        });
        cancelled
    }

    #[cfg(feature = "async")]
    pub(crate) fn complete_async_call_cleanup(
        &mut self,
        token: &Arc<AtomicBool>,
    ) -> crate::Result<()> {
        let record =
            self.async_call_cleanups
                .last()
                .ok_or(crate::error::InternalError::FrameStack(
                    "missing async call cleanup record",
                ))?;
        if !Arc::ptr_eq(&record.cancelled, token) {
            return Err(crate::error::InternalError::FrameStack(
                "async call cleanup order mismatch",
            )
            .into());
        }
        self.async_call_cleanups.pop();
        Ok(())
    }

    #[cfg(feature = "async")]
    pub(crate) fn recover_cancelled_async_calls(&mut self) {
        use std::sync::atomic::Ordering;

        while self
            .async_call_cleanups
            .last()
            .is_some_and(|record| record.cancelled.load(Ordering::Acquire))
        {
            let Some(record) = self.async_call_cleanups.pop() else {
                break;
            };
            if let Some(boundary) = record.boundary {
                let mut exec = self.take_exec();
                if let Err(error) = exec.discard_to(boundary.value_base, boundary.stop_depth) {
                    if let Some(internal) = error.downcast_ref::<crate::error::InternalError>() {
                        self.latch_internal_error(*internal);
                    }
                }
                if boundary.stop_depth != 0 {
                    self.park_exec(exec);
                }
            }
            if let Some(mark) = record.roots_mark {
                self.gc_roots_truncate(mark);
            }
            self.restore_pending_exception(record.pending, record.pending_generation);
        }
    }

    #[cfg(feature = "async")]
    pub(crate) fn abandon_async_call_boundary(&mut self, value_base: usize, stop_depth: usize) {
        if let Some(index) = self.async_call_cleanups.iter().rposition(|record| {
            record.boundary.as_ref().is_some_and(|boundary| {
                boundary.value_base == value_base && boundary.stop_depth == stop_depth
            })
        }) {
            self.async_call_cleanups.truncate(index);
        }
    }

    #[inline]
    pub(crate) fn engine(&self) -> &Engine {
        &self.engine
    }

    pub(crate) fn fuel(&self) -> u64 {
        self.fuel + self.fuel_reserve
    }

    /// Sets total fuel, split into active + reserve per the yield interval (none ⇒ all active).
    pub(crate) fn set_fuel(&mut self, total: u64) {
        let active = self
            .fuel_yield_interval
            .map_or(total, |i| total.min(i.get()));
        self.fuel = active;
        self.fuel_reserve = total - active;
    }

    /// Sets the async fuel-yield interval, then re-splits the current total fuel.
    pub(crate) fn set_fuel_yield_interval(&mut self, interval: Option<u64>) {
        self.fuel_yield_interval = interval.and_then(NonZeroU64::new);
        self.set_fuel(self.fuel());
    }

    /// Charges one unit; `NeedYield` if only reserve remains (yield point), `Exhausted` if none.
    pub(crate) fn consume_fuel_step(&mut self) -> FuelStep {
        if self.fuel > 0 {
            self.fuel -= 1;
            FuelStep::Ran
        } else if self.fuel_reserve > 0 {
            FuelStep::NeedYield
        } else {
            FuelStep::Exhausted
        }
    }

    /// Moves one interval's fuel from reserve into the active slice (after a yield); no-op if unset.
    pub(crate) fn refuel_from_reserve(&mut self) {
        let take = self
            .fuel_yield_interval
            .map_or(0, NonZeroU64::get)
            .min(self.fuel_reserve);
        self.fuel += take;
        self.fuel_reserve -= take;
    }

    pub(crate) fn set_epoch_deadline(&mut self, deadline: u64) {
        self.epoch_deadline = deadline;
    }

    pub(crate) fn epoch_deadline_reached(&self) -> bool {
        self.engine.current_epoch() >= self.epoch_deadline
    }

    /// This store's process-unique id (#34), stamped onto handles it mints.
    pub(crate) fn store_id(&self) -> u64 {
        self.store_id
    }

    /// Guards an entity access against **cross-store handle misuse** (#34): a handle minted by another
    /// store must not silently resolve to the wrong entity here. The `0` sentinel (an internal,
    /// same-store-by-construction funcref value rebuilt from a bare cell index) is accepted. A real
    /// mismatch is an embedder bug — not guest-reachable (a guest can't forge a foreign handle) — so
    /// it panics, matching wasmtime's "wrong store" behavior.
    #[inline]
    pub(crate) fn check_handle(&self, store: u64) {
        assert!(
            store == 0 || store == self.store_id,
            "handle used with the wrong store (cross-store misuse)"
        );
    }

    pub(crate) fn alloc_func(&mut self, entity: FuncEntity) -> Func {
        Func {
            index: self.funcs.alloc(entity),
            store: self.store_id,
        }
    }

    pub(crate) fn alloc_memory(&mut self, entity: MemoryEntity) -> Memory {
        Memory {
            index: self.memories.alloc(entity),
            store: self.store_id,
        }
    }

    /// Takes the reused host-call buffers (a nested call sees empty defaults).
    #[inline]
    pub(crate) fn take_host_scratch(&mut self) -> (Vec<crate::value::Val>, Vec<crate::value::Val>) {
        std::mem::take(&mut self.host_scratch)
    }

    #[inline]
    pub(crate) fn put_host_scratch(
        &mut self,
        mut params: Vec<crate::value::Val>,
        mut results: Vec<crate::value::Val>,
    ) {
        params.clear();
        results.clear();
        self.host_scratch = (params, results);
    }

    pub(crate) fn alloc_table(&mut self, entity: TableEntity) -> Table {
        Table {
            index: self.tables.alloc(entity),
            store: self.store_id,
        }
    }

    pub(crate) fn alloc_global(&mut self, entity: GlobalEntity) -> Global {
        Global {
            index: self.globals.alloc(entity),
            store: self.store_id,
        }
    }

    pub(crate) fn alloc_tag(&mut self, entity: TagEntity) -> Tag {
        Tag {
            index: self.tags.alloc(entity),
            store: self.store_id,
        }
    }

    // Exception-arena accessors (`alloc_exn`/`exn`/`exn_mut`/`exn_checked`/`exn_generation`) live in
    // `managed` with the other GC-managed reference arenas (the `exns` field stays here).

    /// The handle the *next* [`alloc_instance`](Self::alloc_instance) will return,
    /// without allocating. Lets instantiation build `FuncEntity`s that point back
    /// at the instance before it exists (see `instance::init`). Only valid while no
    /// other instance is allocated in between (we hold `&mut StoreInner` throughout).
    pub(crate) fn reserve_instance(&self) -> Instance {
        Instance {
            index: self.instances.len(),
            store: self.store_id,
        }
    }

    pub(crate) fn alloc_instance(&mut self, entity: InstanceEntity) -> Instance {
        Instance {
            index: self.instances.alloc(entity),
            store: self.store_id,
        }
    }

    /// Fallible instance lookup — `None` for an unregistered handle (synthetic test executions).
    pub(crate) fn try_instance(&self, handle: Instance) -> Option<&InstanceEntity> {
        self.instances.get_opt(handle.index)
    }

    pub(crate) fn memory_count(&self) -> usize {
        self.memories.len() as usize
    }

    pub(crate) fn table_count(&self) -> usize {
        self.tables.len() as usize
    }

    pub(crate) fn instance_count(&self) -> usize {
        self.instances.len() as usize
    }
}

impl Drop for StoreInner {
    fn drop(&mut self) {
        // Release this store's GC reservation from the engine-wide committed total.
        self.engine.sub_gc_committed(self.gc.reserved());
        // Release the type-registrations pinned by host GC allocations.
        for &id in &self.gc_host_alloc_types {
            self.engine.decref_type(id);
        }
    }
}

#[cfg(test)]
#[path = "exn_tests.rs"]
mod exn_tests;

#[cfg(test)]
mod tests {
    use super::*;
    use crate::value::{GlobalType, Mutability, Val, ValType};

    #[test]
    fn global_arena_round_trip() {
        let mut inner = StoreInner::new(Engine::default());
        let handle = inner.alloc_global(GlobalEntity {
            value: Val::I32(7),
            ty: GlobalType::new(ValType::I32, Mutability::Var),
        });
        assert_eq!(inner.global(handle).value.unwrap_i32(), 7);
        inner.global_mut(handle).value = Val::I32(9);
        assert_eq!(inner.global(handle).value.unwrap_i32(), 9);
    }
}