diff --git a/crates/wiggle/generate/src/wasmtime.rs b/crates/wiggle/generate/src/wasmtime.rs index be3603754645..aaffe8989c0a 100644 --- a/crates/wiggle/generate/src/wasmtime.rs +++ b/crates/wiggle/generate/src/wasmtime.rs @@ -112,13 +112,23 @@ fn generate_func( }; let body = quote! { - let mem = match caller.get_export("memory") { - Some(#rt::wasmtime_crate::Extern::Memory(m)) => m, + let (mem, ctx, shared) = match caller.get_export("memory") { + Some(#rt::wasmtime_crate::Extern::Memory(m)) => { + let (mem, ctx) = m.data_and_store_mut(&mut caller); + let ctx = get_cx(ctx); + (mem, ctx, false) + } + Some(#rt::wasmtime_crate::Extern::SharedMemory(m)) => { + // SAFETY: `WasmtimeGuestMemory` never hands out the mutable + // slice and the `unsafe`-ness is propagated by the `shared` + // flag. + let mem = unsafe { std::mem::transmute::<&[std::cell::UnsafeCell], &mut [u8]>(m.data()) }; + let ctx = get_cx(caller.data_mut()); + (mem, ctx, true) + } _ => #rt::anyhow::bail!("missing required memory export"), }; - let (mem , ctx) = mem.data_and_store_mut(&mut caller); - let ctx = get_cx(ctx); - let mem = #rt::wasmtime::WasmtimeGuestMemory::new(mem, false); + let mem = #rt::wasmtime::WasmtimeGuestMemory::new(mem, shared); Ok(<#ret_ty>::from(#abi_func(ctx, &mem #(, #arg_names)*) #await_ ?)) };