perf: keep shared-memory locking out of line

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

@@ -1432,7 +1432,7 @@ impl<'store> Executor<'store> {
         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 addr = crate::store::with_memory!(@lock_inline 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 {
@@ -1484,7 +1484,7 @@ impl<'store> Executor<'store> {
 
         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| {
+        crate::store::with_memory!(@lock_inline 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::<N>(base, memory.offset()))?;
             if addr % N != 0 {

crates/tinywasm/src/store/mod.rs

@@ -34,6 +34,8 @@ pub(crate) use memory::{MemValue, MemoryInstance};
 pub use memory::{MemoryShared, MemorySharedGuard};
 pub(crate) use state::State;
 pub(crate) use state::with_memory;
+#[cfg(feature = "std")]
+pub(crate) use state::with_shared_memory;
 pub(crate) use types::{canonicalize_ref_type, canonicalize_value_type};
 pub(crate) use {data::*, element::*, function::*, global::*, table::*, tag::*};
 
@@ -822,9 +824,8 @@ impl Store {
                     };
                     let offset = usize::try_from(offset).unwrap_or(usize::MAX);
                     with_memory!(self.state, *mem_addr, |mem, kind| {
-                        mem.write_all(offset, &data.data)
-                            .ok_or_else(|| memory::memory_oob(offset, data.data.len(), mem.len()))?;
-                    });
+                        mem.write_all(offset, &data.data).ok_or_else(|| memory::memory_oob(offset, data.data.len(), mem.len()))
+                    })?;
                     self.state.data[data_addrs[i] as usize].drop();
                 }
                 tinywasm_types::DataKind::Passive => {}

crates/tinywasm/src/store/state.rs

@@ -26,16 +26,19 @@ pub(crate) struct State {
     pub(crate) roots: gc::Roots,
 }
 
-// Dispatch once per operation, keeping ordinary memory on the direct-access path
+// Dispatch once per operation, keeping ordinary memory on the direct-access path.
+// A shared memory runs the body out of line, under its lock, so the handlers of
+// ordinary memory accesses contain no lock/unlock calls and spill no registers for them.
+// Atomic operations, whose memory is usually shared, keep the lock inline with
+// `with_memory!(@lock_inline ...)`.
 macro_rules! with_memory {
-    ($state:expr, $addr:expr, |$memory:ident, $kind:ident| $body:block) => {{
+    (@lock_inline $state:expr, $addr:expr, |$memory:ident, $kind:ident| $body:block) => {{
         let state = &mut $state;
         let addr = $addr;
         #[cfg(feature = "std")]
         let mut guard;
         #[cfg(feature = "std")]
         let (kind, bytes) = if addr & $crate::store::SHARED_MEM_BIT != 0 {
-            core::hint::cold_path();
             guard = state.shared_memories[(addr & !$crate::store::SHARED_MEM_BIT) as usize].lock();
             (guard.kind, &mut *guard.inner)
         } else {
@@ -52,9 +55,50 @@ macro_rules! with_memory {
         let $memory = bytes;
         $body
     }};
+    ($state:expr, $addr:expr, |$memory:ident, $kind:ident| $body:block) => {{
+        let state = &mut $state;
+        let addr = $addr;
+        #[cfg(feature = "std")]
+        let result = if addr & $crate::store::SHARED_MEM_BIT != 0 {
+            $crate::store::with_shared_memory(state, addr, |$memory, kind| {
+                #[allow(unused_variables)]
+                let $kind = kind;
+                $body
+            })
+        } else {
+            let ordinary = &mut state.memories[addr as usize];
+            #[allow(unused_variables)]
+            let $kind = ordinary.kind;
+            let $memory = &mut ordinary.inner;
+            $body
+        };
+        #[cfg(not(feature = "std"))]
+        let result = {
+            let ordinary = state.get_mem_mut(addr);
+            #[allow(unused_variables)]
+            let $kind = ordinary.kind;
+            let $memory = &mut ordinary.inner;
+            $body
+        };
+        result
+    }};
 }
 pub(crate) use with_memory;
 
+/// Runs `body` on a shared memory's bytes while holding its lock.
+#[cfg(feature = "std")]
+#[cold]
+#[inline(never)]
+pub(crate) fn with_shared_memory<R>(
+    state: &mut State,
+    addr: MemAddr,
+    body: impl FnOnce(&mut super::memory::MemoryStorage, MemoryType) -> R,
+) -> R {
+    let mut guard = state.shared_memories[(addr & !SHARED_MEM_BIT) as usize].lock();
+    let kind = guard.kind;
+    body(&mut guard.inner, kind)
+}
+
 impl State {
     /// Returns the immutable memory type without taking a shared-memory lock.
     pub(crate) fn memory_type(&self, addr: MemAddr) -> MemoryType {
@@ -120,11 +164,11 @@ impl State {
         // Never hold two backing locks, since another store may copy in the opposite direction.
         // Growth cannot invalidate these ranges. Check both before changing the destination.
         with_memory!(*self, src_addr, |source, kind| {
-            source.checked_range(src, size).ok_or_else(|| memory_oob(src, size, source.len()))?;
-        });
+            source.checked_range(src, size).ok_or_else(|| memory_oob(src, size, source.len()))
+        })?;
         with_memory!(*self, dst_addr, |destination, kind| {
-            destination.checked_range(dst, size).ok_or_else(|| memory_oob(dst, size, destination.len()))?;
-        });
+            destination.checked_range(dst, size).ok_or_else(|| memory_oob(dst, size, destination.len()))
+        })?;
         let mut bytes = [0u8; 4096];
         for offset in (0..size).step_by(bytes.len()) {
             let chunk = bytes.len().min(size - offset);