feat: add custom store state

Signed-off-by: Henry <mail@henrygressmann.de>

Author: explodingcamera

CHANGELOG.md

@@ -41,6 +41,7 @@ This release adds garbage collection and typed function reference support, singl
 
 ### Added
 
+- Optional store-owned host state through `Store::with_state`, `Store::state`, and `Store::state_mut`.
 - Support for the typed function references, garbage collection, exception handling, and compact import section proposals.
 - `ResourceLimiter` callbacks for memory and table allocation or growth.
 - A default-enabled `validate` Cargo feature and a parser option to skip WebAssembly validation.

README.md

@@ -56,6 +56,7 @@ See the [examples](./examples) directory and [documentation](https://docs.rs/tin
 - **`archive`:** Enables serialization and deserialization of the internal `twasm` format. Enabled by default.
 - **`log`:** Enables integration with the `log` crate. Enabled by default.
 - **`send`:** Makes stores and store-local handles movable across threads.
+- **`state`:** Adds store-owned host state for access from host callbacks.
 - **`portable-atomic`:** Supports targets without native atomic compare-and-swap.
 - **`canonicalize-nans`:** Uses a [canonical NaN](https://en.wikipedia.org/wiki/NaN#Canonical_NaN) for normalized NaN results. Enabled by default.
 - **`debug`:** Derives `Debug` for runtime types. Enabled by default.

crates/tinywasm/Cargo.toml

@@ -11,7 +11,7 @@ keywords.workspace = true
 categories.workspace = true
 
 [package.metadata.docs.rs]
-features = ["std", "validate", "parser", "archive", "log", "canonicalize-nans", "debug", "guest-debug"]
+features = ["std", "validate", "parser", "archive", "log", "canonicalize-nans", "state", "debug", "guest-debug"]
 rustdoc-args = ["--cfg", "docsrs"]
 
 [lib]
@@ -90,14 +90,17 @@ wat.workspace = true
 
 [features]
 default = ["full", "std"]
-full = ["archive", "debug", "parallel-parser", "parser", "validate"]
+full = ["archive", "debug", "parallel-parser", "parser", "state", "validate"]
 
 log = ["dep:log", "tinywasm-parser?/log"]
 std = ["tinywasm-parser?/std", "tinywasm-types/std"]
 
 # make stores, instances, references, and host callbacks movable across threads
 send = []
 
+# store host state for access from host callbacks
+state = []
+
 # support for parsing WebAssembly
 parser = ["dep:tinywasm-parser"]
 

crates/tinywasm/src/store/mod.rs

@@ -1,4 +1,8 @@
+#[cfg(feature = "state")]
+use alloc::collections::BTreeMap;
 use alloc::{boxed::Box, format, vec::Vec};
+#[cfg(feature = "state")]
+use core::any::{Any, TypeId};
 use core::hint::cold_path;
 use tinywasm_types::*;
 
@@ -110,6 +114,25 @@ pub trait ResourceLimiter: Send + Sync {
     }
 }
 
+/// A value stored as host state.
+#[doc(hidden)]
+#[cfg(all(feature = "state", not(feature = "send")))]
+pub trait StoreState: Any {}
+#[cfg(all(feature = "state", not(feature = "send")))]
+impl<T: Any> StoreState for T {}
+
+/// A value stored as host state.
+#[doc(hidden)]
+#[cfg(all(feature = "state", feature = "send"))]
+pub trait StoreState: Any + Send {}
+#[cfg(all(feature = "state", feature = "send"))]
+impl<T: Any + Send> StoreState for T {}
+
+#[cfg(all(feature = "state", not(feature = "send")))]
+type StoredState = dyn Any;
+#[cfg(all(feature = "state", feature = "send"))]
+type StoredState = dyn Any + Send;
+
 /// Runtime state used by WebAssembly instances and host functions.
 ///
 /// ## Example
@@ -133,6 +156,8 @@ pub struct Store {
     pub(crate) call_stack: CallStack,
     pub(crate) value_stack: ValueStack,
     value_scratch: ValueScratch,
+    #[cfg(feature = "state")]
+    host_state: BTreeMap<TypeId, Box<StoredState>>,
 }
 
 #[derive(Default)]
@@ -238,9 +263,47 @@ impl Store {
             engine,
             execution_fuel: 0,
             execution_active: false,
+            #[cfg(feature = "state")]
+            host_state: BTreeMap::new(),
         }
     }
 
+    /// Adds host state to this store.
+    ///
+    /// Host callbacks can access the value through [`Store::state`] and
+    /// [`Store::state_mut`]. If state of the same type already exists, this
+    /// method replaces it.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use tinywasm::Store;
+    ///
+    /// struct AppState {
+    ///     requests: usize,
+    /// }
+    ///
+    /// let store = Store::default().with_state(AppState { requests: 0 });
+    /// assert_eq!(store.state::<AppState>().unwrap().requests, 0);
+    /// ```
+    #[cfg(feature = "state")]
+    pub fn with_state<T: StoreState>(mut self, state: T) -> Self {
+        self.host_state.insert(TypeId::of::<T>(), Box::new(state));
+        self
+    }
+
+    /// Get host state.
+    #[cfg(feature = "state")]
+    pub fn state<T: StoreState>(&self) -> Option<&T> {
+        self.host_state.get(&TypeId::of::<T>()).and_then(|state| state.downcast_ref())
+    }
+
+    /// Get mutable access to host state.
+    #[cfg(feature = "state")]
+    pub fn state_mut<T: StoreState>(&mut self) -> Option<&mut T> {
+        self.host_state.get_mut(&TypeId::of::<T>()).and_then(|state| state.downcast_mut())
+    }
+
     pub(crate) const fn id(&self) -> StoreId {
         self.id
     }

crates/tinywasm/tests/store_state.rs

@@ -0,0 +1,36 @@
+#![cfg(feature = "state")]
+
+use tinywasm::types::WasmValue;
+use tinywasm::{HostFunction, Store};
+
+struct Counter(i32);
+struct Label(&'static str);
+
+#[test]
+fn stores_state_by_type() {
+    let mut store = Store::default().with_state(Counter(1)).with_state(Label("first")).with_state(Counter(2));
+
+    assert_eq!(store.state::<Counter>().unwrap().0, 2);
+    assert_eq!(store.state::<Label>().unwrap().0, "first");
+
+    store.state_mut::<Counter>().unwrap().0 = 3;
+    assert_eq!(store.state::<Counter>().unwrap().0, 3);
+}
+
+#[test]
+fn host_function_mutates_store_state() -> tinywasm::Result<()> {
+    let mut store = Store::default().with_state(Counter(40));
+    let function = HostFunction::from(|mut ctx, value: i32| {
+        let counter = ctx.state_mut::<Counter>().expect("counter state");
+        counter.0 += value;
+        Ok(counter.0)
+    })
+    .instantiate(&mut store)?;
+
+    let mut results = [WasmValue::I32(0)];
+    function.call(&mut store, &[WasmValue::I32(2)], &mut results)?;
+
+    assert_eq!(results, [WasmValue::I32(42)]);
+    assert_eq!(store.state::<Counter>().unwrap().0, 42);
+    Ok(())
+}

examples/doom/Cargo.toml

@@ -14,6 +14,7 @@ softbuffer = "0.4"
 tinywasm = { path = "../../crates/tinywasm", features = [
   "archive",
   "canonicalize-nans",
+  "state",
   "debug",
   "log",
   "parallel-parser",

examples/doom/src/main.rs

@@ -129,7 +129,7 @@ impl ApplicationHandler for DoomApp {
             window.request_redraw();
         }
 
-        if self.runtime.host_state.lock().unwrap().exit_code.is_some() {
+        if self.runtime.exited() {
             event_loop.exit();
         }
     }