diff --git a/crates/pecos-wasm/src/wasmtime_foreign_object.rs b/crates/pecos-wasm/src/wasmtime_foreign_object.rs index 5de2049c7..bf44ef953 100644 --- a/crates/pecos-wasm/src/wasmtime_foreign_object.rs +++ b/crates/pecos-wasm/src/wasmtime_foreign_object.rs @@ -17,11 +17,11 @@ use crate::foreign_object::ForeignObject; use log::{debug, warn}; -use parking_lot::{Mutex, RwLock}; +use parking_lot::RwLock; use pecos_core::errors::PecosError; use std::any::Any; use std::path::Path; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use std::thread; use std::time::Duration; use wasmtime::{ @@ -82,9 +82,8 @@ impl StoreContext { #[derive(Debug)] pub struct WasmForeignObject { /// WebAssembly binary - #[allow(dead_code)] wasm_bytes: Vec, - /// Wasmtime engine + /// Wasmtime engine (must be kept alive for module/store to remain valid) #[allow(dead_code)] engine: Engine, /// Wasmtime module @@ -93,8 +92,8 @@ pub struct WasmForeignObject { store: RwLock>, /// Wasmtime instance (thread-safe) instance: RwLock>, - /// Cached function names (thread-safe) - func_names: Mutex>>, + /// Cached function names (initialized once, lock-free reads) + func_names: OnceLock>, /// Stop flag for epoch increment thread stop_flag: Arc>, /// Last function call results @@ -281,7 +280,7 @@ impl WasmForeignObject { module, store: RwLock::new(store), instance: RwLock::new(None), - func_names: Mutex::new(None), + func_names: OnceLock::new(), stop_flag, last_results: Vec::new(), timeout_seconds, @@ -332,13 +331,30 @@ impl WasmForeignObject { /// /// Returns an error if the `shot_reinit` function exists but execution fails pub fn shot_reinit(&mut self) -> Result<(), PecosError> { - let funcs = self.get_funcs(); - if funcs.contains(&"shot_reinit".to_string()) { + if self.has_func("shot_reinit") { self.exec("shot_reinit", &[])?; } Ok(()) } + /// Check if a function exists in the WebAssembly module (zero-allocation) + /// + /// This is more efficient than `get_funcs().contains()` for existence checks + /// as it doesn't allocate or clone the function list. + /// + /// # Parameters + /// + /// * `func_name` - Name of the function to check + /// + /// # Returns + /// + /// `true` if the function exists, `false` otherwise + fn has_func(&self, func_name: &str) -> bool { + self.module + .exports() + .any(|e| e.name() == func_name && e.ty().func().is_some()) + } + /// Get the WebAssembly binary bytes /// /// This is useful for serialization and cloning. @@ -392,8 +408,7 @@ impl ForeignObject for WasmForeignObject { self.new_instance()?; // Check if the init function exists - let funcs = self.get_funcs(); - if !funcs.contains(&"init".to_string()) { + if !self.has_func("init") { return Err(PecosError::Input( "WebAssembly module must contain an 'init' function".to_string(), )); @@ -420,23 +435,27 @@ impl ForeignObject for WasmForeignObject { } fn get_funcs(&self) -> Vec { - // Check if we've already cached the function names - if let Some(ref funcs) = *self.func_names.lock() { - return funcs.clone(); - } - - // Get the function names - let mut funcs = Vec::new(); - for export in self.module.exports() { - if export.ty().func().is_some() { - funcs.push(export.name().to_string()); - } - } - - // Cache the function names - *self.func_names.lock() = Some(funcs.clone()); + // Use OnceLock::get_or_init for lock-free cached access + // The Arc is cloned (cheap), then converted to Vec for trait compatibility + let arc_funcs = self.func_names.get_or_init(|| { + // Collect function names and convert to Arc<[String]> + let funcs: Vec = self + .module + .exports() + .filter_map(|export| { + if export.ty().func().is_some() { + Some(export.name().to_string()) + } else { + None + } + }) + .collect(); + Arc::from(funcs) + }); - funcs + // Convert Arc<[String]> to Vec for trait compatibility + // This clones the strings but Arc::clone is cheap + arc_funcs.to_vec() } fn exec(&mut self, func_name: &str, args: &[i64]) -> Result, PecosError> { diff --git a/python/pecos-rslib/rust/src/wasm_foreign_object_bindings.rs b/python/pecos-rslib/rust/src/wasm_foreign_object_bindings.rs index 5846d0d80..965b3fd73 100644 --- a/python/pecos-rslib/rust/src/wasm_foreign_object_bindings.rs +++ b/python/pecos-rslib/rust/src/wasm_foreign_object_bindings.rs @@ -180,6 +180,15 @@ impl PyWasmForeignObject { } } + /// Get the WebAssembly binary bytes + /// + /// Returns: + /// The WASM binary as bytes + #[getter] + fn wasm_bytes<'py>(&self, py: Python<'py>) -> Bound<'py, PyBytes> { + PyBytes::new(py, self.inner.wasm_bytes()) + } + /// Cleanup resources /// /// Stops the epoch increment thread. This is called automatically diff --git a/python/quantum-pecos/src/pecos/foreign_objects/wasmtime.py b/python/quantum-pecos/src/pecos/foreign_objects/wasmtime.py index 647407db5..cd3827107 100644 --- a/python/quantum-pecos/src/pecos/foreign_objects/wasmtime.py +++ b/python/quantum-pecos/src/pecos/foreign_objects/wasmtime.py @@ -65,8 +65,8 @@ def __init__( memory_size=memory_size, ) - # Get WASM bytes for compatibility with serialization - self.wasm_bytes = self._rust_obj.to_dict()["wasm_bytes"] + # Get WASM bytes directly from property (no dict allocation) + self.wasm_bytes = self._rust_obj.wasm_bytes def init(self) -> None: """Initialize object before running a series of experiments.