feat: add support for atomic waiters / notify

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

Author: explodingcamera

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added WebAssembly threads support, including atomic wait/notify and cross-store shared memory with `MemoryShared`.
 - Added new `Memory::data` and `Memory::data_mut` methods for direct access to linear memory.
 - Added several internal methods to support a future c-api crate
 

README.md

@@ -46,7 +46,7 @@ See the [examples](./examples) directory and [documentation](https://docs.rs/tin
 ## Cargo Features
 
 - **`full`:** Enables `archive`, `debug`, `parallel-parser`, `parser`, `state`, and `validate`. Enabled by default.
-- **`std`:** Enables `std` and parsing from files and streams. Enabled by default.
+- **`std`:** Enables parsing from files and streams, shared memory handles, and blocking atomic wait/notify across stores. Enabled by default.
 - **`parser`:** Enables `tinywasm-parser` and top-level parse helpers. Enabled by default.
 - **`validate`:** Enables WebAssembly validation while parsing. Enabled by default and configurable through [`ParserOptions`](https://docs.rs/tinywasm/latest/tinywasm/parser/struct.ParserOptions.html).
 - **`parallel-parser`:** Parallelizes function parsing when `std` is enabled. Enabled by default.
@@ -104,8 +104,8 @@ TinyWasm generally implements non-JavaScript core proposals at [phase 4](https:/
 | [**Garbage Collection**](https://github.com/WebAssembly/gc)                                                      | 🟢     | 0.11.0             |
 | [**Exception Handling**](https://github.com/WebAssembly/exception-handling)                                      | 🟢     | 0.11.0             |
 | [**Compact Import Section**](https://github.com/WebAssembly/compact-import-section)                              | 🟢     | 0.11.0             |
+| [**Threads**](https://github.com/WebAssembly/threads)                                                            | 🚧     | next               |
 | [**Stack Switching**](https://github.com/WebAssembly/stack-switching)                                            | 🌑     | -                  |
-| [**Threads**](https://github.com/WebAssembly/threads)                                                            | 🌑     | -                  |
 
 **Legend**\
 🌑 -- not available\

crates/parser/src/visit.rs

@@ -11,9 +11,9 @@ use alloc::{
     vec::Vec,
 };
 use tinywasm_types::{
-    AtomicArg, AtomicOp, AtomicWidth, BinOp, BinOp128, CmpOp, ExceptionHandler, Global, Import, ImportKind,
-    Instruction, MemoryType, Operand64, Operand64Idx, Operand128, Operand128Idx, StorageType, TableDefinition, TagType,
-    TypeSection, ValueCounts, ValueLane, WasmFunctionData,
+    AtomicArg, AtomicOp, AtomicWaitOp, AtomicWidth, BinOp, BinOp128, CmpOp, ExceptionHandler, Global, Import,
+    ImportKind, Instruction, MemoryType, Operand64, Operand64Idx, Operand128, Operand128Idx, StorageType,
+    TableDefinition, TagType, TypeSection, ValueCounts, ValueLane, WasmFunctionData,
 };
 use wasmparser::{FunctionBody, OperatorsReader, OperatorsReaderAllocations, VisitSimdOperator};
 
@@ -231,6 +231,21 @@ impl<'a> FunctionBuilder<'a> {
         self.uses_local_memory |= memory >= self.metadata.imported_memories;
     }
 
+    fn atomic_wait(&mut self, memarg: wasmparser::MemArg, width: u8, op: AtomicWaitOp) -> Result<()> {
+        if memarg.align != width.trailing_zeros() as u8 {
+            return Err(crate::ParseError::Other("invalid atomic alignment".into()));
+        }
+        let address = self.metadata.memory_size(memarg.memory)?;
+        self.mark_memory(memarg.memory);
+        let memory = self.push128(Operand128::<tinywasm_types::MemoryOperand>::new(memarg.offset, memarg.memory))?;
+        let inputs: &[ValueLane] = match op {
+            AtomicWaitOp::Notify => &[address, ValueLane::S32],
+            AtomicWaitOp::Wait32 => &[address, ValueLane::S32, ValueLane::S64],
+            AtomicWaitOp::Wait64 => &[address, ValueLane::S64, ValueLane::S64],
+        };
+        self.emit(inputs, &[ValueLane::S32], Instruction::AtomicWait(memory, op))
+    }
+
     fn visit_struct_get_impl(
         &mut self,
         type_index: u32,
@@ -539,14 +554,14 @@ impl<'a> wasmparser::VisitOperator<'a> for FunctionBuilder<'_> {
         self.emit(&[], &[], Instruction::AtomicFence)
     }
 
-    fn visit_memory_atomic_notify(&mut self, _: wasmparser::MemArg) -> Self::Output {
-        Err(crate::ParseError::UnsupportedOperator("memory.atomic.notify".into()))
+    fn visit_memory_atomic_notify(&mut self, memarg: wasmparser::MemArg) -> Self::Output {
+        self.atomic_wait(memarg, 4, AtomicWaitOp::Notify)
     }
-    fn visit_memory_atomic_wait32(&mut self, _: wasmparser::MemArg) -> Self::Output {
-        Err(crate::ParseError::UnsupportedOperator("memory.atomic.wait32".into()))
+    fn visit_memory_atomic_wait32(&mut self, memarg: wasmparser::MemArg) -> Self::Output {
+        self.atomic_wait(memarg, 4, AtomicWaitOp::Wait32)
     }
-    fn visit_memory_atomic_wait64(&mut self, _: wasmparser::MemArg) -> Self::Output {
-        Err(crate::ParseError::UnsupportedOperator("memory.atomic.wait64".into()))
+    fn visit_memory_atomic_wait64(&mut self, memarg: wasmparser::MemArg) -> Self::Output {
+        self.atomic_wait(memarg, 8, AtomicWaitOp::Wait64)
     }
 
     atomic_visitors! {

crates/tinywasm/Cargo.toml

@@ -42,7 +42,6 @@ name = "test-wasm-multi-memory"
 [[test]]
 harness = false
 name = "test-wasm-threads"
-test = false
 
 [[test]]
 harness = false

crates/tinywasm/src/interpreter/executor/instructions.rs

@@ -324,6 +324,7 @@ macro_rules! instruction_handlers {
             MemorySize(addr) => executor.exec_memory_size(*addr)?,
             MemoryGrow(addr) => executor.exec_memory_grow(*addr)?,
             Atomic(arg) => executor.exec_atomic(*arg)?,
+            AtomicWait(memory, op) => executor.exec_atomic_wait(*memory, *op)?,
             AtomicFence => crate::shared::fence(crate::shared::Ordering::SeqCst),
 
             // Bulk memory operations

crates/tinywasm/src/interpreter/executor/mod.rs

@@ -561,10 +561,7 @@ impl<'store> Executor<'store> {
     fn throw_exception(&mut self, exception: ValueRef, protected_ip: usize) -> ExecResult<ExecFlow> {
         match self.dispatch_exception(exception, protected_ip)? {
             Some(landing_pad) => Ok(ExecFlow::next(landing_pad)),
-            None => match self.store.root_exception(exception) {
-                Ok(exception) => Err(Error::Exception(exception).into()),
-                Err(error) => Err(error.into()),
-            },
+            None => Err(Error::Exception(self.store.root_exception(exception)?).into()),
         }
     }
 
@@ -734,6 +731,7 @@ impl<'store> Executor<'store> {
         let type_addr = operand.a();
         let table_addr = operand.b();
         self.charge_call_fuel(FUEL_COST_CALL_TOTAL);
+
         // verify that the table is of the right type, this should be validated by the parser already
         let table_addr = self.module.resolve_table_addr(table_addr);
         let table_idx = self.pop_table_operand(self.store.state.get_table(table_addr).kind.arch())?;
@@ -1100,14 +1098,13 @@ impl<'store> Executor<'store> {
         let type_index = operand.a();
         let len = operand.b();
         let type_addr = self.module.resolve_type_addr(type_index);
-        let storage =
-            self.store.state.get_type(type_addr).as_array().expect("validated array.new_fixed type").field.storage;
+        let arr_type = self.store.state.get_type(type_addr).as_array().expect("validated array.new_fixed type");
         let len = len as usize;
         self.store.state.gc.check_allocation(len, self.store.state.gc_type_has_references(type_addr))?;
         let mut values = Vec::new();
         cold_err!(values.try_reserve_exact(len)).map_err(|_| Trap::OutOfMemory)?;
         for _ in 0..len {
-            values.push(pop_value(&mut self.store.value_stack, storage));
+            values.push(pop_value(&mut self.store.value_stack, arr_type.field.storage));
         }
         values.reverse();
         self.push_gc_object(type_addr, values)
@@ -1257,8 +1254,8 @@ impl<'store> Executor<'store> {
         let object = self.store.state.gc_object(reference, type_addr)?;
         let object_len = self.store.state.gc.get_handle(object).expect("validated array").values.len();
         dst.checked_add(len).filter(|end| *end <= object_len).ok_or(Trap::ArrayOutOfBounds)?;
-        let items =
-            self.store.state.elements[self.module.resolve_elem_addr(elem_index) as usize].items_range(src, len)?;
+        let elem_addr = self.module.resolve_elem_addr(elem_index);
+        let items = self.store.state.elements[elem_addr as usize].items_range(src, len)?;
         let mut values = Vec::new();
         cold_err!(values.try_reserve_exact(len)).map_err(|_| Trap::OutOfMemory)?;
         values.extend(items.iter().copied().map(RuntimeValue::ValueRef));
@@ -1267,11 +1264,9 @@ impl<'store> Executor<'store> {
     }
 
     fn exec_memory_size(&mut self, addr: u32) -> Result<(), Trap> {
-        let mem_addr = self.mem_addr(addr);
-        let (arch, pages) = self.store.state.memory_size(mem_addr);
-        match arch {
-            MemoryArch::I64 => i64::stack_push(&mut self.store.value_stack, pages as i64),
-            MemoryArch::I32 => i32::stack_push(&mut self.store.value_stack, pages as i32),
+        match self.store.state.memory_size(self.mem_addr(addr)) {
+            (MemoryArch::I64, pages) => i64::stack_push(&mut self.store.value_stack, pages as i64),
+            (MemoryArch::I32, pages) => i32::stack_push(&mut self.store.value_stack, pages as i32),
         }
     }
 
@@ -1441,24 +1436,65 @@ impl<'store> Executor<'store> {
         }
     }
 
+    fn exec_atomic_wait(&mut self, arg: Operand128Idx<MemoryOperand>, op: AtomicWaitOp) -> Result<(), Trap> {
+        let timeout = if op == AtomicWaitOp::Notify { 0 } else { i64::stack_pop(&mut self.store.value_stack) };
+        let value = match op {
+            AtomicWaitOp::Wait64 => u64::stack_pop(&mut self.store.value_stack),
+            _ => u32::stack_pop(&mut self.store.value_stack) as u64,
+        };
+        let memory = arg.resolve(&self.func.data);
+        let mem_addr = self.mem_addr(memory.memory());
+        let width = if op == AtomicWaitOp::Wait64 { 8 } else { 4 };
+        let addr = crate::store::with_memory!(self.store.state, mem_addr, |mem, kind| {
+            let base = self.store.value_stack.pop_memory_operand(kind.arch())?;
+            let addr = cold_err!(mem.effective_addr::<1>(base, memory.offset()))?;
+            if addr % width != 0 {
+                return cold!(Err(Trap::UnalignedAtomic));
+            }
+
+            if op != AtomicWaitOp::Notify && !kind.shared() {
+                return cold!(Err(Trap::Other("atomic wait requires shared memory")));
+            }
+
+            // Check the entire access even for notify and even when no waiter exists.
+            if mem.checked_range(addr, width).is_none() {
+                return cold!(Err(Trap::MemoryOutOfBounds { offset: addr, len: width, max: mem.len() }));
+            }
+
+            Ok(addr)
+        })?;
+
+        #[cfg(not(feature = "std"))]
+        let _ = (timeout, value, addr);
+
+        #[cfg(feature = "std")]
+        if mem_addr & crate::store::SHARED_MEM_BIT != 0 {
+            let shared = &self.store.state.shared_memories[(mem_addr & !crate::store::SHARED_MEM_BIT) as usize];
+            let result = match op {
+                AtomicWaitOp::Notify => shared.notify(addr, value as u32),
+                AtomicWaitOp::Wait32 => shared.wait::<4>(addr, value, timeout)?,
+                AtomicWaitOp::Wait64 => shared.wait::<8>(addr, value, timeout)?,
+            };
+            return u32::stack_push(&mut self.store.value_stack, result);
+        }
+
+        u32::stack_push(&mut self.store.value_stack, 0)
+    }
+
     fn exec_atomic_width<const N: usize>(&mut self, arg: AtomicArg) -> Result<(), Trap> {
         let op = arg.op();
-        let value = if op == AtomicOp::Load {
-            0
-        } else if arg.is_64() {
-            u64::stack_pop(&mut self.store.value_stack)
-        } else {
-            u32::stack_pop(&mut self.store.value_stack) as u64
+        let value = match op {
+            AtomicOp::Load => 0,
+            _ if arg.is_64() => u64::stack_pop(&mut self.store.value_stack),
+            _ => u32::stack_pop(&mut self.store.value_stack) as u64,
         };
-        let expected = if op == AtomicOp::Cmpxchg {
-            if arg.is_64() {
-                u64::stack_pop(&mut self.store.value_stack)
-            } else {
-                u32::stack_pop(&mut self.store.value_stack) as u64
-            }
-        } else {
-            0
+
+        let expected = match op {
+            AtomicOp::Cmpxchg if arg.is_64() => u64::stack_pop(&mut self.store.value_stack),
+            AtomicOp::Cmpxchg => u32::stack_pop(&mut self.store.value_stack) as u64,
+            _ => 0,
         };
+
         let memory = arg.memory.resolve(&self.func.data);
         let mem_addr = self.mem_addr(memory.memory());
         crate::store::with_memory!(self.store.state, mem_addr, |mem, kind| {
@@ -1493,10 +1529,10 @@ impl<'store> Executor<'store> {
                 bytes.copy_from_slice(&next.to_le_bytes()[..N]);
                 cold_err!(mem.write_fixed::<N>(addr, &bytes))?;
             }
-            if arg.is_64() {
-                u64::stack_push(&mut self.store.value_stack, old)?;
-            } else {
-                u32::stack_push(&mut self.store.value_stack, old as u32)?;
+
+            match arg.is_64() {
+                true => u64::stack_push(&mut self.store.value_stack, old)?,
+                false => u32::stack_push(&mut self.store.value_stack, old as u32)?,
             }
             Ok(())
         })
@@ -1580,19 +1616,16 @@ impl<'store> Executor<'store> {
 
     fn exec_table_init(&mut self, index: Operand64Idx<(u32, u32)>) -> Result<(), Trap> {
         let operand = index.resolve(&self.func.data);
-        let elem_index = operand.a();
-        let table_index = operand.b();
         let size = self.pop_table_operand(MemoryArch::I32)?; // n
         let offset = self.pop_table_operand(MemoryArch::I32)?; // s
-        let table_addr = self.module.resolve_table_addr(table_index);
+        let table_addr = self.module.resolve_table_addr(operand.b());
         let dst = self.pop_table_operand(self.store.state.get_table(table_addr).kind.arch())?; // d
-        let elem_addr = self.module.resolve_elem_addr(elem_index) as usize;
+        let elem_addr = self.module.resolve_elem_addr(operand.a()) as usize;
         let elem = self.store.state.elements.get(elem_addr).ok_or_else(|| Trap::Other("element not found"))?;
         let items = elem.items_range(offset, size)?;
 
-        let table =
-            self.store.state.tables.get_mut(table_addr as usize).ok_or_else(|| Trap::Other("table not found"))?;
-        table.init(dst, items)
+        let table = self.store.state.tables.get_mut(table_addr as usize);
+        table.ok_or_else(|| Trap::Other("table not found"))?.init(dst, items)
     }
 
     fn exec_table_grow(&mut self, table_index: u32) -> Result<(), Trap> {

crates/tinywasm/src/store/memory/shared.rs

@@ -1,5 +1,10 @@
-use crate::std::sync::{Mutex, MutexGuard};
-use alloc::{sync::Arc, vec::Vec};
+use crate::std::sync::{Condvar, Mutex, MutexGuard};
+use crate::std::time::Duration;
+use alloc::{
+    collections::{BTreeMap, VecDeque},
+    sync::Arc,
+    vec::Vec,
+};
 use core::sync::atomic::{AtomicUsize, Ordering};
 use tinywasm_types::MemoryType;
 
@@ -15,6 +20,13 @@ struct MemorySharedInstance {
     kind: MemoryType,
     pages: AtomicUsize,
     bytes: Mutex<MemoryStorage>,
+    // Lock order: bytes, waiters, then an individual waiter's notified flag.
+    waiters: Mutex<BTreeMap<usize, VecDeque<Arc<Waiter>>>>,
+}
+
+struct Waiter {
+    notified: Mutex<bool>,
+    wake: Condvar,
 }
 
 #[cfg(feature = "debug")]
@@ -35,7 +47,7 @@ impl MemoryShared {
         let kind = instance.kind;
         let pages = AtomicUsize::new(instance.page_count);
         let bytes = Mutex::new(instance.inner);
-        Self(Arc::new(MemorySharedInstance { kind, pages, bytes }))
+        Self(Arc::new(MemorySharedInstance { kind, pages, bytes, waiters: Mutex::new(BTreeMap::new()) }))
     }
 
     /// Whether two handles reference the same instance, even across store slots.
@@ -115,6 +127,76 @@ impl MemoryShared {
         guard.inner.fill(offset, len, val).ok_or_else(|| memory_oob(offset, len, guard.inner.len()).into())
     }
 
+    /// Waits until notified or the signed nanosecond timeout expires.
+    pub(crate) fn wait<const N: usize>(
+        &self,
+        addr: usize,
+        expected: u64,
+        timeout: i64,
+    ) -> core::result::Result<u32, Trap> {
+        let guard = self.lock();
+        let bytes = guard.inner.read_fixed::<N>(addr)?;
+        let mut value = [0u8; 8];
+        value[..N].copy_from_slice(&bytes);
+        if u64::from_le_bytes(value) != expected {
+            return Ok(1);
+        }
+        if timeout == 0 {
+            return Ok(2);
+        }
+        let waiter = Arc::new(Waiter { notified: Mutex::new(false), wake: Condvar::new() });
+        self.0
+            .waiters
+            .lock()
+            .unwrap_or_else(|poison| poison.into_inner())
+            .entry(addr)
+            .or_default()
+            .push_back(waiter.clone());
+        drop(guard);
+
+        let notified = waiter.notified.lock().unwrap_or_else(|poison| poison.into_inner());
+        if timeout < 0 {
+            drop(waiter.wake.wait_while(notified, |ready| !*ready).unwrap_or_else(|poison| poison.into_inner()));
+        } else {
+            drop(
+                waiter
+                    .wake
+                    .wait_timeout_while(notified, Duration::from_nanos(timeout as u64), |ready| !*ready)
+                    .unwrap_or_else(|poison| poison.into_inner())
+                    .0,
+            );
+        }
+        // The condvar guard was dropped before taking waiters, preserving lock order.
+        let mut waiters = self.0.waiters.lock().unwrap_or_else(|poison| poison.into_inner());
+        let notified = *waiter.notified.lock().unwrap_or_else(|poison| poison.into_inner());
+        if let Some(queue) = waiters.get_mut(&addr) {
+            queue.retain(|entry| !Arc::ptr_eq(entry, &waiter));
+            if queue.is_empty() {
+                waiters.remove(&addr);
+            }
+        }
+        Ok(if notified { 0 } else { 2 })
+    }
+
+    /// Wakes at most `count` waiters at the given byte address.
+    pub(crate) fn notify(&self, addr: usize, count: u32) -> u32 {
+        let _guard = self.lock();
+        let mut waiters = self.0.waiters.lock().unwrap_or_else(|poison| poison.into_inner());
+        let mut woken = 0;
+        if let Some(queue) = waiters.get_mut(&addr) {
+            while woken < count {
+                let Some(waiter) = queue.pop_front() else { break };
+                *waiter.notified.lock().unwrap_or_else(|poison| poison.into_inner()) = true;
+                waiter.wake.notify_one();
+                woken += 1;
+            }
+            if queue.is_empty() {
+                waiters.remove(&addr);
+            }
+        }
+        woken
+    }
+
     /// Grows the memory by the given number of pages, returning its previous size.
     pub fn grow(&self, pages: i64) -> Result<Option<i64>> {
         Ok(self.grow_with_limiter(pages, None)?)