Repository: bytecodealliance/wasmtime
Pull request #5307 (2022-11-21 07:06:31)

Refactor Sigdata methods with sigset
Changed stack_ret_arg to become a boolean on SigData, as the stack_return_arg is always the last argument.
Refactored all SigSet methods on SigData to be in SigSet instead.


tnachen: @jameysharp changing the Option<u16> to bool made the SigData size from 40 -> 32.
However, I'm trying to figure out how to remove more fields based on the SigSet suggestion. 
Based on your earlier comment, seems like the motivation is to be able to index over SigDatas since they're all lined up in the abi args list.
But from the function calls it seems like we're parsing a function sig one at a time and not sure what the change should look like here.
bjorn3: > I'll defer to @fitzgen and @cfallin on whether it's okay to make stack_ret_arg a bool. You've shown that every ABI we currently support places that argument at the end, which I think is a pretty good argument for doing this. But we want to support more targets in the future, so maybe we don't want to over-simplify this detail? On the other hand, we can always change this again if we ever find a target that doesn't put the stack ret arg at the end.

The position shouldn't matter I think. Only the register assignment.
jameysharp: I talked with @fitzgen and @cfallin about this offline. We aren't comfortable hard-coding the assumption that `stack_return_arg` is always last right now. (@bjorn3, you're right that the register assignment is all that matters, but as I understand it the register assignment gets computed from this list, so order still matters.)

So @tnachen, could you remove the `stack_return_arg` changes from this PR? I want to merge the rest of the work you did here. Let me know whether you want to try de-duplicating the `args_start` fields in this PR, or if I should merge this and you can open another PR. I would recommend opening another PR for that unless you feel you can do it within a day or two, just so we can get the work you've already done merged promptly.
tnachen: @jameysharp ah ok! I just reverted those changes and kept the Sigdata refactoring. 
We can merge this as-is and I can submit another one for the return field changes
tnachen: @jameysharp addressed comments!

Patch:
diff --git a/RELEASES.md b/RELEASES.md
index 6cc80f22c..f99a433a6 100644
--- a/RELEASES.md
+++ b/RELEASES.md
@@ -12,7 +12,7 @@ Unreleased.
 
 ## 3.0.0
 
-Unreleased.
+Released 2022-11-21
 
 ### Added
 
diff --git a/cranelift/codegen/src/isa/s390x/lower/isle.rs b/cranelift/codegen/src/isa/s390x/lower/isle.rs
index 44ec3994b..1f6f5ce9a 100644
--- a/cranelift/codegen/src/isa/s390x/lower/isle.rs
+++ b/cranelift/codegen/src/isa/s390x/lower/isle.rs
@@ -129,7 +129,7 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Flags, IsaFlags, 6>
         for i in 0..self.lower_ctx.sigs()[*abi].num_rets() {
             if let &ABIArg::Slots {
                 ref slots, purpose, ..
-            } = &self.lower_ctx.sigs()[*abi].get_ret(self.lower_ctx.sigs(), i)
+            } = &self.lower_ctx.sigs().get_ret(*abi, i)
             {
                 if purpose == ArgumentPurpose::StructReturn {
                     continue;
@@ -192,8 +192,10 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Flags, IsaFlags, 6>
         defs: &CallRetList,
         opcode: &Opcode,
     ) -> BoxCallInfo {
-        let clobbers =
-            self.lower_ctx.sigs()[*abi].call_clobbers::<S390xMachineDeps>(self.lower_ctx.sigs());
+        let clobbers = self
+            .lower_ctx
+            .sigs()
+            .call_clobbers::<S390xMachineDeps>(*abi);
         Box::new(CallInfo {
             dest: name.clone(),
             uses: uses.clone(),
@@ -214,8 +216,10 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Flags, IsaFlags, 6>
         defs: &CallRetList,
         opcode: &Opcode,
     ) -> BoxCallIndInfo {
-        let clobbers =
-            self.lower_ctx.sigs()[*abi].call_clobbers::<S390xMachineDeps>(self.lower_ctx.sigs());
+        let clobbers = self
+            .lower_ctx
+            .sigs()
+            .call_clobbers::<S390xMachineDeps>(*abi);
         Box::new(CallIndInfo {
             rn: target,
             uses: uses.clone(),
diff --git a/cranelift/codegen/src/machinst/abi.rs b/cranelift/codegen/src/machinst/abi.rs
index f83b59ea8..84b4a7864 100644
--- a/cranelift/codegen/src/machinst/abi.rs
+++ b/cranelift/codegen/src/machinst/abi.rs
@@ -643,114 +643,6 @@ pub struct SigData {
 }
 
 impl SigData {
-    pub fn from_func_sig<M: ABIMachineSpec>(
-        sigs: &mut SigSet,
-        sig: &ir::Signature,
-        flags: &settings::Flags,
-    ) -> CodegenResult<SigData> {
-        let sret = missing_struct_return(sig);
-        let returns = sret.as_ref().into_iter().chain(&sig.returns);
-
-        // Compute args and retvals from signature. Handle retvals first,
-        // because we may need to add a return-area arg to the args.
-        let rets_start = u32::try_from(sigs.abi_args.len()).unwrap();
-        let (sized_stack_ret_space, _) = M::compute_arg_locs(
-            sig.call_conv,
-            flags,
-            returns,
-            ArgsOrRets::Rets,
-            /* extra ret-area ptr = */ false,
-            ArgsAccumulator::new(&mut sigs.abi_args),
-        )?;
-        let rets_end = u32::try_from(sigs.abi_args.len()).unwrap();
-
-        let need_stack_return_area = sized_stack_ret_space > 0;
-        let args_start = u32::try_from(sigs.abi_args.len()).unwrap();
-        let (sized_stack_arg_space, stack_ret_arg) = M::compute_arg_locs(
-            sig.call_conv,
-            flags,
-            &sig.params,
-            ArgsOrRets::Args,
-            need_stack_return_area,
-            ArgsAccumulator::new(&mut sigs.abi_args),
-        )?;
-        let args_end = u32::try_from(sigs.abi_args.len()).unwrap();
-
-        let args_length = u16::try_from(args_end - args_start).unwrap();
-        let rets_length = u16::try_from(rets_end - rets_start).unwrap();
-
-        trace!(
-            "ABISig: sig {:?} => args start = {} args len = {} rets start = {} rets len = {}
-             arg stack = {} ret stack = {} stack_ret_arg = {:?}",
-            sig,
-            args_start,
-            args_length,
-            rets_start,
-            rets_length,
-            sized_stack_arg_space,
-            sized_stack_ret_space,
-            stack_ret_arg,
-        );
-
-        let stack_ret_arg = stack_ret_arg.map(|s| u16::try_from(s).unwrap());
-        Ok(SigData {
-            args_start,
-            args_length,
-            rets_start,
-            rets_length,
-            sized_stack_arg_space,
-            sized_stack_ret_space,
-            stack_ret_arg,
-            call_conv: sig.call_conv,
-        })
-    }
-
-    /// Get this signature's ABI arguments.
-    pub fn args<'a>(&self, sigs: &'a SigSet) -> &'a [ABIArg] {
-        let start = usize::try_from(self.args_start).unwrap();
-        let end = usize::try_from(self.args_start + u32::from(self.args_length)).unwrap();
-        &sigs.abi_args[start..end]
-    }
-
-    /// Get this signature's ABI returns.
-    pub fn rets<'a>(&self, sigs: &'a SigSet) -> &'a [ABIArg] {
-        let start = usize::try_from(self.rets_start).unwrap();
-        let end = usize::try_from(self.rets_start + u32::from(self.rets_length)).unwrap();
-        &sigs.abi_args[start..end]
-    }
-
-    /// Return all clobbers for the callsite.
-    pub fn call_clobbers<M: ABIMachineSpec>(&self, sigs: &SigSet) -> PRegSet {
-        // Get clobbers: all caller-saves. These may include return value
-        // regs, which we will remove from the clobber set below.
-        let mut clobbers = M::get_regs_clobbered_by_call(self.call_conv);
-
-        // Remove retval regs from clobbers. Skip StructRets: these
-        // are not, semantically, returns at the CLIF level, so we
-        // treat such a value as a clobber instead.
-        for ret in self.rets(sigs) {
-            if let &ABIArg::Slots {
-                ref slots, purpose, ..
-            } = ret
-            {
-                if purpose == ir::ArgumentPurpose::StructReturn {
-                    continue;
-                }
-                for slot in slots {
-                    match slot {
-                        &ABIArgSlot::Reg { reg, .. } => {
-                            log::trace!("call_clobbers: retval reg {:?}", reg);
-                            clobbers.remove(PReg::from(reg));
-                        }
-                        _ => {}
-                    }
-                }
-            }
-        }
-
-        clobbers
-    }
-
     /// Get the number of arguments expected.
     pub fn num_args(&self) -> usize {
         let num = if self.stack_ret_arg.is_some() {
@@ -762,11 +654,6 @@ impl SigData {
         usize::from(num)
     }
 
-    /// Get information specifying how to pass one argument.
-    pub fn get_arg(&self, sigs: &SigSet, idx: usize) -> ABIArg {
-        self.args(sigs)[idx].clone()
-    }
-
     /// Get total stack space required for arguments.
     pub fn sized_stack_arg_space(&self) -> i64 {
         self.sized_stack_arg_space
@@ -777,23 +664,11 @@ impl SigData {
         usize::from(self.rets_length)
     }
 
-    /// Get information specifying how to pass one return value.
-    pub fn get_ret(&self, sigs: &SigSet, idx: usize) -> ABIArg {
-        self.rets(sigs)[idx].clone()
-    }
-
     /// Get total stack space required for return values.
     pub fn sized_stack_ret_space(&self) -> i64 {
         self.sized_stack_ret_space
     }
 
-    /// Get information specifying how to pass the implicit pointer
-    /// to the return-value area on the stack, if required.
-    pub fn get_ret_arg(&self, sigs: &SigSet) -> Option<ABIArg> {
-        let ret_arg = usize::from(self.stack_ret_arg?);
-        Some(self.args(sigs)[ret_arg].clone())
-    }
-
     /// Get calling convention used.
     pub fn call_conv(&self) -> isa::CallConv {
         self.call_conv
@@ -878,7 +753,7 @@ impl SigSet {
         // `ir::Signature`.
         debug_assert!(!self.have_abi_sig_for_signature(&signature));
 
-        let sig_data = SigData::from_func_sig::<M>(self, &signature, flags)?;
+        let sig_data = self.from_func_sig::<M>(&signature, flags)?;
         let sig = self.sigs.push(sig_data);
         self.ir_signature_to_abi_sig.insert(signature, sig);
         Ok(sig)
@@ -897,7 +772,7 @@ impl SigSet {
             return Ok(sig);
         }
         let signature = &dfg.signatures[sig_ref];
-        let sig_data = SigData::from_func_sig::<M>(self, signature, flags)?;
+        let sig_data = self.from_func_sig::<M>(signature, flags)?;
         let sig = self.sigs.push(sig_data);
         self.ir_sig_ref_to_abi_sig[sig_ref] = Some(sig);
         Ok(sig)
@@ -920,6 +795,138 @@ impl SigSet {
             .copied()
             .expect("must call `make_abi_sig_from_ir_signature` before `get_abi_sig_for_signature`")
     }
+
+    pub fn from_func_sig<M: ABIMachineSpec>(
+        &mut self,
+        sig: &ir::Signature,
+        flags: &settings::Flags,
+    ) -> CodegenResult<SigData> {
+        let sret = missing_struct_return(sig);
+        let returns = sret.as_ref().into_iter().chain(&sig.returns);
+
+        // Compute args and retvals from signature. Handle retvals first,
+        // because we may need to add a return-area arg to the args.
+        let rets_start = u32::try_from(self.abi_args.len()).unwrap();
+        let (sized_stack_ret_space, _) = M::compute_arg_locs(
+            sig.call_conv,
+            flags,
+            returns,
+            ArgsOrRets::Rets,
+            /* extra ret-area ptr = */ false,
+            ArgsAccumulator::new(&mut self.abi_args),
+        )?;
+        let rets_end = u32::try_from(self.abi_args.len()).unwrap();
+
+        let need_stack_return_area = sized_stack_ret_space > 0;
+        let args_start = u32::try_from(self.abi_args.len()).unwrap();
+        let (sized_stack_arg_space, stack_ret_arg) = M::compute_arg_locs(
+            sig.call_conv,
+            flags,
+            &sig.params,
+            ArgsOrRets::Args,
+            need_stack_return_area,
+            ArgsAccumulator::new(&mut self.abi_args),
+        )?;
+        let args_end = u32::try_from(self.abi_args.len()).unwrap();
+
+        let args_length = u16::try_from(args_end - args_start).unwrap();
+        let rets_length = u16::try_from(rets_end - rets_start).unwrap();
+
+        trace!(
+            "ABISig: sig {:?} => args start = {} args len = {} rets start = {} rets len = {}
+             arg stack = {} ret stack = {} stack_ret_arg = {:?}",
+            sig,
+            args_start,
+            args_length,
+            rets_start,
+            rets_length,
+            sized_stack_arg_space,
+            sized_stack_ret_space,
+            need_stack_return_area,
+        );
+
+        let stack_ret_arg = stack_ret_arg.map(|s| u16::try_from(s).unwrap());
+        Ok(SigData {
+            args_start,
+            args_length,
+            rets_start,
+            rets_length,
+            sized_stack_arg_space,
+            sized_stack_ret_space,
+            stack_ret_arg,
+            call_conv: sig.call_conv,
+        })
+    }
+
+    /// Get this signature's ABI arguments.
+    pub fn args(&self, sig: Sig) -> &[ABIArg] {
+        let sig_data = &self.sigs[sig];
+        let start = usize::try_from(sig_data.args_start).unwrap();
+        let end = usize::try_from(sig_data.args_start + u32::from(sig_data.args_length)).unwrap();
+        &self.abi_args[start..end]
+    }
+
+    /// Get information specifying how to pass the implicit pointer
+    /// to the return-value area on the stack, if required.
+    pub fn get_ret_arg(&self, sig: Sig) -> Option<ABIArg> {
+        let sig_data = &self.sigs[sig];
+        if let Some(i) = sig_data.stack_ret_arg {
+            Some(self.args(sig)[usize::from(i)].clone())
+        } else {
+            None
+        }
+    }
+
+    /// Get information specifying how to pass one argument.
+    pub fn get_arg(&self, sig: Sig, idx: usize) -> ABIArg {
+        self.args(sig)[idx].clone()
+    }
+
+    /// Get this signature's ABI returns.
+    pub fn rets(&self, sig: Sig) -> &[ABIArg] {
+        let sig_data = &self.sigs[sig];
+        let start = usize::try_from(sig_data.rets_start).unwrap();
+        let end = usize::try_from(sig_data.rets_start + u32::from(sig_data.rets_length)).unwrap();
+        &self.abi_args[start..end]
+    }
+
+    /// Get information specifying how to pass one return value.
+    pub fn get_ret(&self, sig: Sig, idx: usize) -> ABIArg {
+        self.rets(sig)[idx].clone()
+    }
+
+    /// Return all clobbers for the callsite.
+    pub fn call_clobbers<M: ABIMachineSpec>(&self, sig: Sig) -> PRegSet {
+        let sig_data = &self.sigs[sig];
+        // Get clobbers: all caller-saves. These may include return value
+        // regs, which we will remove from the clobber set below.
+        let mut clobbers = M::get_regs_clobbered_by_call(sig_data.call_conv);
+
+        // Remove retval regs from clobbers. Skip StructRets: these
+        // are not, semantically, returns at the CLIF level, so we
+        // treat such a value as a clobber instead.
+        for ret in self.rets(sig) {
+            if let &ABIArg::Slots {
+                ref slots, purpose, ..
+            } = ret
+            {
+                if purpose == ir::ArgumentPurpose::StructReturn {
+                    continue;
+                }
+                for slot in slots {
+                    match slot {
+                        &ABIArgSlot::Reg { reg, .. } => {
+                            log::trace!("call_clobbers: retval reg {:?}", reg);
+                            clobbers.remove(PReg::from(reg));
+                        }
+                        _ => {}
+                    }
+                }
+            }
+        }
+
+        clobbers
+    }
 }
 
 // NB: we do _not_ implement `IndexMut` because these signatures are
@@ -999,11 +1006,11 @@ pub struct Callee<M: ABIMachineSpec> {
 fn get_special_purpose_param_register(
     f: &ir::Function,
     sigs: &SigSet,
-    abi: &SigData,
+    sig: &Sig,
     purpose: ir::ArgumentPurpose,
 ) -> Option<Reg> {
     let idx = f.signature.special_param_index(purpose)?;
-    match &abi.args(sigs)[idx] {
+    match &sigs.args(*sig)[idx] {
         &ABIArg::Slots { ref slots, .. } => match &slots[0] {
             &ABIArgSlot::Reg { reg, .. } => Some(reg.into()),
             _ => None,
@@ -1080,17 +1087,13 @@ impl<M: ABIMachineSpec> Callee<M> {
         // stack limit. This can either be specified as a special-purpose
         // argument or as a global value which often calculates the stack limit
         // from the arguments.
-        let stack_limit = get_special_purpose_param_register(
-            f,
-            sigs,
-            &sigs[sig],
-            ir::ArgumentPurpose::StackLimit,
-        )
-        .map(|reg| (reg, smallvec![]))
-        .or_else(|| {
-            f.stack_limit
-                .map(|gv| gen_stack_limit::<M>(f, sigs, &sigs[sig], gv))
-        });
+        let stack_limit =
+            get_special_purpose_param_register(f, sigs, &sig, ir::ArgumentPurpose::StackLimit)
+                .map(|reg| (reg, smallvec![]))
+                .or_else(|| {
+                    f.stack_limit
+                        .map(|gv| gen_stack_limit::<M>(f, sigs, &sig, gv))
+                });
 
         // Determine whether a probestack call is required for large enough
         // frames (and the minimum frame size if so).
@@ -1210,25 +1213,25 @@ impl<M: ABIMachineSpec> Callee<M> {
 fn gen_stack_limit<M: ABIMachineSpec>(
     f: &ir::Function,
     sigs: &SigSet,
-    abi: &SigData,
+    sig: &Sig,
     gv: ir::GlobalValue,
 ) -> (Reg, SmallInstVec<M::I>) {
     let mut insts = smallvec![];
-    let reg = generate_gv::<M>(f, sigs, abi, gv, &mut insts);
+    let reg = generate_gv::<M>(f, sigs, sig, gv, &mut insts);
     return (reg, insts);
 }
 
 fn generate_gv<M: ABIMachineSpec>(
     f: &ir::Function,
     sigs: &SigSet,
-    abi: &SigData,
+    sig: &Sig,
     gv: ir::GlobalValue,
     insts: &mut SmallInstVec<M::I>,
 ) -> Reg {
     match f.global_values[gv] {
         // Return the direct register the vmcontext is in
         ir::GlobalValueData::VMContext => {
-            get_special_purpose_param_register(f, sigs, abi, ir::ArgumentPurpose::VMContext)
+            get_special_purpose_param_register(f, sigs, sig, ir::ArgumentPurpose::VMContext)
                 .expect("no vmcontext parameter found")
         }
         // Load our base value into a register, then load from that register
@@ -1239,7 +1242,7 @@ fn generate_gv<M: ABIMachineSpec>(
             global_type: _,
             readonly: _,
         } => {
-            let base = generate_gv::<M>(f, sigs, abi, base, insts);
+            let base = generate_gv::<M>(f, sigs, sig, base, insts);
             let into_reg = Writable::from_reg(M::get_stacklimit_reg());
             insts.push(M::gen_load_base_offset(
                 into_reg,
@@ -1321,7 +1324,7 @@ impl<M: ABIMachineSpec> Callee<M> {
     /// They will be provided to `init()` as the `temps` arg if so.
     pub fn temps_needed(&self, sigs: &SigSet) -> Vec<Type> {
         let mut temp_tys = vec![];
-        for arg in sigs[self.sig].args(sigs) {
+        for arg in sigs.args(self.sig) {
             match arg {
                 &ABIArg::ImplicitPtrArg { pointer, .. } => match &pointer {
                     &ABIArgSlot::Reg { .. } => {}
@@ -1343,7 +1346,7 @@ impl<M: ABIMachineSpec> Callee<M> {
     /// once the lowering context exists.
     pub fn init(&mut self, sigs: &SigSet, temps: Vec<Writable<Reg>>) {
         let mut temps_iter = temps.into_iter();
-        for arg in sigs[self.sig].args(sigs) {
+        for arg in sigs.args(self.sig) {
             let temp = match arg {
                 &ABIArg::ImplicitPtrArg { pointer, .. } => match &pointer {
                     &ABIArgSlot::Reg { .. } => None,
@@ -1441,7 +1444,7 @@ impl<M: ABIMachineSpec> Callee<M> {
             }
         };
 
-        match &sigs[self.sig].args(sigs)[idx] {
+        match &sigs.args(self.sig)[idx] {
             &ABIArg::Slots { ref slots, .. } => {
                 assert_eq!(into_regs.len(), slots.len());
                 for (slot, into_reg) in slots.iter().zip(into_regs.regs().iter()) {
@@ -1518,7 +1521,7 @@ impl<M: ABIMachineSpec> Callee<M> {
         let mut reg_pairs = smallvec![];
         let mut ret = smallvec![];
         let word_bits = M::word_bits() as u8;
-        match &sigs[self.sig].rets(sigs)[idx] {
+        match &sigs.rets(self.sig)[idx] {
             &ABIArg::Slots { ref slots, .. } => {
                 assert_eq!(from_regs.len(), slots.len());
                 for (slot, &from_reg) in slots.iter().zip(from_regs.regs().iter()) {
@@ -2015,7 +2018,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         flags: settings::Flags,
     ) -> CodegenResult<Caller<M>> {
         let sig = sigs.abi_sig_for_sig_ref(sig_ref);
-        let clobbers = sigs[sig].call_clobbers::<M>(sigs);
+        let clobbers = sigs.call_clobbers::<M>(sig);
         Ok(Caller {
             sig,
             uses: smallvec![],
@@ -2040,7 +2043,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         flags: settings::Flags,
     ) -> CodegenResult<Caller<M>> {
         let sig = sigs.abi_sig_for_signature(sig);
-        let clobbers = sigs[sig].call_clobbers::<M>(sigs);
+        let clobbers = sigs.call_clobbers::<M>(sig);
         Ok(Caller {
             sig,
             uses: smallvec![],
@@ -2065,7 +2068,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         flags: settings::Flags,
     ) -> CodegenResult<Caller<M>> {
         let sig = sigs.abi_sig_for_sig_ref(sig_ref);
-        let clobbers = sigs[sig].call_clobbers::<M>(sigs);
+        let clobbers = sigs.call_clobbers::<M>(sig);
         Ok(Caller {
             sig,
             uses: smallvec![],
@@ -2094,9 +2097,8 @@ fn adjust_stack_and_nominal_sp<M: ABIMachineSpec>(ctx: &mut Lower<M::I>, off: i3
 impl<M: ABIMachineSpec> Caller<M> {
     /// Get the number of arguments expected.
     pub fn num_args(&self, sigs: &SigSet) -> usize {
-        let data = &sigs[self.sig];
-        let len = data.args(sigs).len();
-        if data.stack_ret_arg.is_some() {
+        let len = sigs.args(self.sig).len();
+        if sigs[self.sig].stack_ret_arg.is_some() {
             len - 1
         } else {
             len
@@ -2128,7 +2130,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         idx: usize,
         from_regs: ValueRegs<Reg>,
     ) {
-        match &ctx.sigs()[self.sig].args(ctx.sigs())[idx] {
+        match &ctx.sigs().args(self.sig)[idx] {
             &ABIArg::Slots { .. } => {}
             &ABIArg::StructArg { offset, size, .. } => {
                 let src_ptr = from_regs.only_reg().unwrap();
@@ -2180,7 +2182,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         // How many temps do we need for extends? Allocate them ahead
         // of time, since we can't do it while we're iterating over
         // the sig and immutably borrowing `ctx`.
-        let needed_tmps = match &ctx.sigs()[self.sig].args(ctx.sigs())[idx] {
+        let needed_tmps = match &ctx.sigs().args(self.sig)[idx] {
             &ABIArg::Slots { ref slots, .. } => slots
                 .iter()
                 .map(|slot| match slot {
@@ -2205,7 +2207,7 @@ impl<M: ABIMachineSpec> Caller<M> {
             .map(|_| ctx.alloc_tmp(M::word_type()).only_reg().unwrap())
             .collect();
 
-        match &ctx.sigs()[self.sig].args(ctx.sigs())[idx] {
+        match &ctx.sigs().args(self.sig)[idx] {
             &ABIArg::Slots { ref slots, .. } => {
                 assert_eq!(from_regs.len(), slots.len());
                 for (slot, from_reg) in slots.iter().zip(from_regs.regs().iter()) {
@@ -2307,7 +2309,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         into_regs: ValueRegs<Writable<Reg>>,
     ) -> SmallInstVec<M::I> {
         let mut insts = smallvec![];
-        match &ctx.sigs()[self.sig].rets(ctx.sigs())[idx] {
+        match &ctx.sigs().rets(self.sig)[idx] {
             &ABIArg::Slots { ref slots, .. } => {
                 assert_eq!(into_regs.len(), slots.len());
                 for (slot, into_reg) in slots.iter().zip(into_regs.regs().iter()) {
diff --git a/cranelift/codegen/src/machinst/isle.rs b/cranelift/codegen/src/machinst/isle.rs
index 441a37fa8..a7a22ca1a 100644
--- a/cranelift/codegen/src/machinst/isle.rs
+++ b/cranelift/codegen/src/machinst/isle.rs
@@ -436,7 +436,7 @@ macro_rules! isle_lower_prelude_methods {
         }
 
         fn abi_get_arg(&mut self, abi: &Sig, idx: usize) -> ABIArg {
-            self.lower_ctx.sigs()[*abi].get_arg(self.lower_ctx.sigs(), idx)
+            self.lower_ctx.sigs().get_arg(*abi, idx)
         }
 
         fn abi_num_rets(&mut self, abi: &Sig) -> usize {
@@ -444,15 +444,15 @@ macro_rules! isle_lower_prelude_methods {
         }
 
         fn abi_get_ret(&mut self, abi: &Sig, idx: usize) -> ABIArg {
-            self.lower_ctx.sigs()[*abi].get_ret(self.lower_ctx.sigs(), idx)
+            self.lower_ctx.sigs().get_ret(*abi, idx)
         }
 
         fn abi_ret_arg(&mut self, abi: &Sig) -> Option<ABIArg> {
-            self.lower_ctx.sigs()[*abi].get_ret_arg(self.lower_ctx.sigs())
+            self.lower_ctx.sigs().get_ret_arg(*abi)
         }
 
         fn abi_no_ret_arg(&mut self, abi: &Sig) -> Option<()> {
-            if let Some(_) = self.lower_ctx.sigs()[*abi].get_ret_arg(self.lower_ctx.sigs()) {
+            if let Some(_) = self.lower_ctx.sigs().get_ret_arg(*abi) {
                 None
             } else {
                 Some(())
@@ -716,8 +716,7 @@ macro_rules! isle_prelude_method_helpers {
                 // Borrow `sigdata` again so we don't hold a `self`
                 // borrow across the `&mut self` arg to
                 // `abi_arg_slot_regs()` below.
-                let sigdata = &self.lower_ctx.sigs()[abi];
-                let ret = sigdata.get_ret(self.lower_ctx.sigs(), i);
+                let ret = self.lower_ctx.sigs().get_ret(abi, i);
                 let retval_regs = self.abi_arg_slot_regs(&ret).unwrap();
                 retval_insts.extend(
                     caller
diff --git a/crates/runtime/src/instance.rs b/crates/runtime/src/instance.rs
index e8d9ed929..9e4f88ecb 100644
--- a/crates/runtime/src/instance.rs
+++ b/crates/runtime/src/instance.rs
@@ -966,8 +966,8 @@ impl Instance {
                 let def_ptr = self.memories[defined_memory_index]
                     .as_shared_memory()
                     .unwrap()
-                    .vmmemory_ptr_mut();
-                ptr::write(ptr, def_ptr);
+                    .vmmemory_ptr();
+                ptr::write(ptr, def_ptr.cast_mut());
             } else {
                 ptr::write(owned_ptr, self.memories[defined_memory_index].vmmemory());
                 ptr::write(ptr, owned_ptr);
diff --git a/crates/runtime/src/lib.rs b/crates/runtime/src/lib.rs
index 325e34740..2c1db36ad 100644
--- a/crates/runtime/src/lib.rs
+++ b/crates/runtime/src/lib.rs
@@ -231,3 +231,17 @@ pub fn page_size() -> usize {
         unsafe { libc::sysconf(libc::_SC_PAGESIZE) as usize }
     }
 }
+
+/// Result of [`Memory::atomic_wait32`] and [`Memory::atomic_wait64`]
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
+pub enum WaitResult {
+    /// Indicates that a `wait` completed by being awoken by a different thread.
+    /// This means the thread went to sleep and didn't time out.
+    Ok = 0,
+    /// Indicates that `wait` did not complete and instead returned due to the
+    /// value in memory not matching the expected value.
+    Mismatch = 1,
+    /// Indicates that `wait` completed with a timeout, meaning that the
+    /// original value matched as expected but nothing ever called `notify`.
+    TimedOut = 2,
+}
diff --git a/crates/runtime/src/libcalls.rs b/crates/runtime/src/libcalls.rs
index 384e7f190..e6bebc25d 100644
--- a/crates/runtime/src/libcalls.rs
+++ b/crates/runtime/src/libcalls.rs
@@ -57,7 +57,7 @@
 use crate::externref::VMExternRef;
 use crate::table::{Table, TableElementType};
 use crate::vmcontext::{VMCallerCheckedAnyfunc, VMContext};
-use crate::{SharedMemory, TrapReason};
+use crate::TrapReason;
 use anyhow::Result;
 use std::mem;
 use std::ptr::{self, NonNull};
@@ -436,25 +436,12 @@ unsafe fn memory_atomic_notify(
     memory_index: u32,
     addr_index: u64,
     count: u32,
-) -> Result<u32, TrapReason> {
+) -> Result<u32, Trap> {
     let memory = MemoryIndex::from_u32(memory_index);
     let instance = (*vmctx).instance_mut();
     instance
-        .get_memory(memory)
-        .validate_addr(addr_index, 4, 4)?;
-
-    let shared_mem = instance.get_runtime_memory(memory).as_shared_memory();
-
-    if count == 0 {
-        return Ok(0);
-    }
-
-    let unparked_threads = shared_mem.map_or(0, |shared_mem| {
-        // SAFETY: checked `addr_index` above
-        unsafe { shared_mem.unchecked_atomic_notify(addr_index, count) }
-    });
-
-    Ok(unparked_threads)
+        .get_runtime_memory(memory)
+        .atomic_notify(addr_index, count)
 }
 
 // Implementation of `memory.atomic.wait32` for locally defined memories.
@@ -464,24 +451,14 @@ unsafe fn memory_atomic_wait32(
     addr_index: u64,
     expected: u32,
     timeout: u64,
-) -> Result<u32, TrapReason> {
+) -> Result<u32, Trap> {
     // convert timeout to Instant, before any wait happens on locking
     let timeout = (timeout as i64 >= 0).then(|| Instant::now() + Duration::from_nanos(timeout));
-
     let memory = MemoryIndex::from_u32(memory_index);
     let instance = (*vmctx).instance_mut();
-    let addr = instance
-        .get_memory(memory)
-        .validate_addr(addr_index, 4, 4)?;
-
-    let shared_mem: SharedMemory = instance
+    Ok(instance
         .get_runtime_memory(memory)
-        .as_shared_memory()
-        .ok_or(Trap::AtomicWaitNonSharedMemory)?;
-
-    // SAFETY: checked `addr_index` above
-    let res = unsafe { shared_mem.unchecked_atomic_wait32(addr_index, addr, expected, timeout) };
-    Ok(res)
+        .atomic_wait32(addr_index, expected, timeout)? as u32)
 }
 
 // Implementation of `memory.atomic.wait64` for locally defined memories.
@@ -491,24 +468,14 @@ unsafe fn memory_atomic_wait64(
     addr_index: u64,
     expected: u64,
     timeout: u64,
-) -> Result<u32, TrapReason> {
+) -> Result<u32, Trap> {
     // convert timeout to Instant, before any wait happens on locking
     let timeout = (timeout as i64 >= 0).then(|| Instant::now() + Duration::from_nanos(timeout));
-
     let memory = MemoryIndex::from_u32(memory_index);
     let instance = (*vmctx).instance_mut();
-    let addr = instance
-        .get_memory(memory)
-        .validate_addr(addr_index, 8, 8)?;
-
-    let shared_mem: SharedMemory = instance
+    Ok(instance
         .get_runtime_memory(memory)
-        .as_shared_memory()
-        .ok_or(Trap::AtomicWaitNonSharedMemory)?;
-
-    // SAFETY: checked `addr_index` above
-    let res = unsafe { shared_mem.unchecked_atomic_wait64(addr_index, addr, expected, timeout) };
-    Ok(res)
+        .atomic_wait64(addr_index, expected, timeout)? as u32)
 }
 
 // Hook for when an instance runs out of fuel.
diff --git a/crates/runtime/src/memory.rs b/crates/runtime/src/memory.rs
index 5ba8bc52e..3cbcd125d 100644
--- a/crates/runtime/src/memory.rs
+++ b/crates/runtime/src/memory.rs
@@ -3,11 +3,9 @@
 //! `RuntimeLinearMemory` is to WebAssembly linear memories what `Table` is to WebAssembly tables.
 
 use crate::mmap::Mmap;
-use crate::parking_spot::{ParkResult, ParkingSpot};
+use crate::parking_spot::ParkingSpot;
 use crate::vmcontext::VMMemoryDefinition;
-use crate::MemoryImage;
-use crate::MemoryImageSlot;
-use crate::Store;
+use crate::{MemoryImage, MemoryImageSlot, Store, WaitResult};
 use anyhow::Error;
 use anyhow::{bail, format_err, Result};
 use std::convert::TryFrom;
@@ -152,8 +150,7 @@ pub trait RuntimeLinearMemory: Send + Sync {
     /// `RuntimeMemoryCreator::new_memory()`.
     fn needs_init(&self) -> bool;
 
-    /// For the pooling allocator, we must be able to downcast this trait to its
-    /// underlying structure.
+    /// Used for optional dynamic downcasting.
     fn as_any_mut(&mut self) -> &mut dyn std::any::Any;
 }
 
@@ -436,6 +433,13 @@ impl RuntimeLinearMemory for StaticMemory {
 #[derive(Clone)]
 pub struct SharedMemory(Arc<SharedMemoryInner>);
 
+struct SharedMemoryInner {
+    memory: RwLock<Box<dyn RuntimeLinearMemory>>,
+    spot: ParkingSpot,
+    ty: wasmtime_environ::Memory,
+    def: LongTermVMMemoryDefinition,
+}
+
 impl SharedMemory {
     /// Construct a new [`SharedMemory`].
     pub fn new(plan: MemoryPlan) -> Result<Self> {
@@ -460,15 +464,17 @@ impl SharedMemory {
             memory.as_any_mut().type_id() != std::any::TypeId::of::<SharedMemory>(),
             "cannot re-wrap a shared memory"
         );
-        let def = LongTermVMMemoryDefinition(memory.vmmemory());
-        let spot = ParkingSpot::default();
-        let inner = RwLock::new(SharedMemoryInnerRwLocked { memory, ty, def });
-        Ok(Self(Arc::new(SharedMemoryInner { spot, lock: inner })))
+        Ok(Self(Arc::new(SharedMemoryInner {
+            ty,
+            spot: ParkingSpot::default(),
+            def: LongTermVMMemoryDefinition(memory.vmmemory()),
+            memory: RwLock::new(memory),
+        })))
     }
 
     /// Return the memory type for this [`SharedMemory`].
     pub fn ty(&self) -> wasmtime_environ::Memory {
-        self.0.lock.read().unwrap().ty
+        self.0.ty
     }
 
     /// Convert this shared memory into a [`Memory`].
@@ -476,33 +482,52 @@ impl SharedMemory {
         Memory(Box::new(self))
     }
 
-    /// Return a mutable pointer to the shared memory's [VMMemoryDefinition].
-    pub fn vmmemory_ptr_mut(&mut self) -> *mut VMMemoryDefinition {
-        &self.0.lock.read().unwrap().def.0 as *const _ as *mut _
-    }
-
     /// Return a pointer to the shared memory's [VMMemoryDefinition].
     pub fn vmmemory_ptr(&self) -> *const VMMemoryDefinition {
-        &self.0.lock.read().unwrap().def.0 as *const _
+        &self.0.def.0
     }
 
-    /// Implementation of `memory.atomic.notify` for this shared memory.
-    pub fn atomic_notify(&self, addr_index: u64, count: u32) -> Result<u32, Trap> {
-        let definition = &self.0.lock.read().unwrap().def.0;
-        definition.validate_addr(addr_index, 4, 4)?;
-        // SAFETY: checked `addr_index` above
-        Ok(unsafe { self.unchecked_atomic_notify(addr_index, count) })
+    /// Same as `RuntimeLinearMemory::grow`, except with `&self`.
+    pub fn grow(
+        &self,
+        delta_pages: u64,
+        store: Option<&mut dyn Store>,
+    ) -> Result<Option<(usize, usize)>, Error> {
+        let mut memory = self.0.memory.write().unwrap();
+        let result = memory.grow(delta_pages, store)?;
+        if let Some((_old_size_in_bytes, new_size_in_bytes)) = result {
+            // Store the new size to the `VMMemoryDefinition` for JIT-generated
+            // code (and runtime functions) to access. No other code can be
+            // growing this memory due to the write lock, but code in other
+            // threads could have access to this shared memory and we want them
+            // to see the most consistent version of the `current_length`; a
+            // weaker consistency is possible if we accept them seeing an older,
+            // smaller memory size (assumption: memory only grows) but presently
+            // we are aiming for accuracy.
+            //
+            // Note that it could be possible to access a memory address that is
+            // now-valid due to changes to the page flags in `grow` above but
+            // beyond the `memory.size` that we are about to assign to. In these
+            // and similar cases, discussion in the thread proposal concluded
+            // that: "multiple accesses in one thread racing with another
+            // thread's `memory.grow` that are in-bounds only after the grow
+            // commits may independently succeed or trap" (see
+            // https://github.com/WebAssembly/threads/issues/26#issuecomment-433930711).
+            // In other words, some non-determinism is acceptable when using
+            // `memory.size` on work being done by `memory.grow`.
+            self.0
+                .def
+                .0
+                .current_length
+                .store(new_size_in_bytes, Ordering::SeqCst);
+        }
+        Ok(result)
     }
 
-    /// Unchecked implementation of `memory.atomic.notify` for this shared memory.
-    ///
-    /// # Safety
-    /// The caller must ensure that `addr_index` is a valid address in the memory.
-    pub unsafe fn unchecked_atomic_notify(&self, addr_index: u64, count: u32) -> u32 {
-        if count == 0 {
-            return 0;
-        }
-        self.0.spot.unpark(addr_index, count)
+    /// Implementation of `memory.atomic.notify` for this shared memory.
+    pub fn atomic_notify(&self, addr_index: u64, count: u32) -> Result<u32, Trap> {
+        validate_atomic_addr(&self.0.def.0, addr_index, 4, 4)?;
+        Ok(self.0.spot.unpark(addr_index, count))
     }
 
     /// Implementation of `memory.atomic.wait32` for this shared memory.
@@ -511,36 +536,18 @@ impl SharedMemory {
         addr_index: u64,
         expected: u32,
         timeout: Option<Instant>,
-    ) -> Result<u32, Trap> {
-        let definition = &self.0.lock.read().unwrap().def.0;
-        let addr = definition.validate_addr(addr_index, 4, 4)?;
-        // SAFETY: checked `addr` above
-        Ok(unsafe { self.unchecked_atomic_wait32(addr_index, addr, expected, timeout) })
-    }
-
-    /// Unchecked implementation of `memory.atomic.wait32` for this shared memory.
-    ///
-    /// # Safety
-    /// The caller must ensure that `addr` is a valid address in the memory.
-    pub unsafe fn unchecked_atomic_wait32(
-        &self,
-        addr_index: u64,
-        addr: *const u8,
-        expected: u32,
-        timeout: Option<Instant>,
-    ) -> u32 {
-        // SAFETY: `addr_index` was validated by `validate_addr` above.
+    ) -> Result<WaitResult, Trap> {
+        let addr = validate_atomic_addr(&self.0.def.0, addr_index, 4, 4)?;
+        // SAFETY: `addr_index` was validated by `validate_atomic_addr` above.
+        assert!(std::mem::size_of::<AtomicU32>() == 4);
+        assert!(std::mem::align_of::<AtomicU32>() <= 4);
         let atomic = unsafe { &*(addr as *const AtomicU32) };
 
         // We want the sequential consistency of `SeqCst` to ensure that the `load` sees the value that the `notify` will/would see.
         // All WASM atomic operations are also `SeqCst`.
         let validate = || atomic.load(Ordering::SeqCst) == expected;
 
-        match self.0.spot.park(addr_index, validate, timeout) {
-            ParkResult::Unparked => 0,
-            ParkResult::Invalid => 1,
-            ParkResult::TimedOut => 2,
-        }
+        Ok(self.0.spot.park(addr_index, validate, timeout))
     }
 
     /// Implementation of `memory.atomic.wait64` for this shared memory.
@@ -549,49 +556,21 @@ impl SharedMemory {
         addr_index: u64,
         expected: u64,
         timeout: Option<Instant>,
-    ) -> Result<u32, Trap> {
-        let definition = &self.0.lock.read().unwrap().def.0;
-        let addr = definition.validate_addr(addr_index, 8, 8)?;
-        // SAFETY: checked `addr` above
-        Ok(unsafe { self.unchecked_atomic_wait64(addr_index, addr, expected, timeout) })
-    }
-    /// Unchecked implementation of `memory.atomic.wait64` for this shared memory.
-    ///
-    /// # Safety
-    /// The caller must ensure that `addr` is a valid address in the memory.
-    pub unsafe fn unchecked_atomic_wait64(
-        &self,
-        addr_index: u64,
-        addr: *const u8,
-        expected: u64,
-        timeout: Option<Instant>,
-    ) -> u32 {
-        // SAFETY: `addr_index` was validated by `validate_addr` above.
+    ) -> Result<WaitResult, Trap> {
+        let addr = validate_atomic_addr(&self.0.def.0, addr_index, 8, 8)?;
+        // SAFETY: `addr_index` was validated by `validate_atomic_addr` above.
+        assert!(std::mem::size_of::<AtomicU64>() == 8);
+        assert!(std::mem::align_of::<AtomicU64>() <= 8);
         let atomic = unsafe { &*(addr as *const AtomicU64) };
 
         // We want the sequential consistency of `SeqCst` to ensure that the `load` sees the value that the `notify` will/would see.
         // All WASM atomic operations are also `SeqCst`.
         let validate = || atomic.load(Ordering::SeqCst) == expected;
 
-        match self.0.spot.park(addr_index, validate, timeout) {
-            ParkResult::Unparked => 0,
-            ParkResult::Invalid => 1,
-            ParkResult::TimedOut => 2,
-        }
+        Ok(self.0.spot.park(addr_index, validate, timeout))
     }
 }
 
-struct SharedMemoryInnerRwLocked {
-    memory: Box<dyn RuntimeLinearMemory>,
-    ty: wasmtime_environ::Memory,
-    def: LongTermVMMemoryDefinition,
-}
-
-struct SharedMemoryInner {
-    lock: RwLock<SharedMemoryInnerRwLocked>,
-    spot: ParkingSpot,
-}
-
 /// Shared memory needs some representation of a `VMMemoryDefinition` for
 /// JIT-generated code to access. This structure owns the base pointer and
 /// length to the actual memory and we share this definition across threads by:
@@ -607,11 +586,11 @@ unsafe impl Sync for LongTermVMMemoryDefinition {}
 /// Proxy all calls through the [`RwLock`].
 impl RuntimeLinearMemory for SharedMemory {
     fn byte_size(&self) -> usize {
-        self.0.lock.read().unwrap().memory.byte_size()
+        self.0.memory.read().unwrap().byte_size()
     }
 
     fn maximum_byte_size(&self) -> Option<usize> {
-        self.0.lock.read().unwrap().memory.maximum_byte_size()
+        self.0.memory.read().unwrap().maximum_byte_size()
     }
 
     fn grow(
@@ -619,39 +598,11 @@ impl RuntimeLinearMemory for SharedMemory {
         delta_pages: u64,
         store: Option<&mut dyn Store>,
     ) -> Result<Option<(usize, usize)>, Error> {
-        let mut inner = self.0.lock.write().unwrap();
-        let result = inner.memory.grow(delta_pages, store)?;
-        if let Some((_old_size_in_bytes, new_size_in_bytes)) = result {
-            // Store the new size to the `VMMemoryDefinition` for JIT-generated
-            // code (and runtime functions) to access. No other code can be
-            // growing this memory due to the write lock, but code in other
-            // threads could have access to this shared memory and we want them
-            // to see the most consistent version of the `current_length`; a
-            // weaker consistency is possible if we accept them seeing an older,
-            // smaller memory size (assumption: memory only grows) but presently
-            // we are aiming for accuracy.
-            //
-            // Note that it could be possible to access a memory address that is
-            // now-valid due to changes to the page flags in `grow` above but
-            // beyond the `memory.size` that we are about to assign to. In these
-            // and similar cases, discussion in the thread proposal concluded
-            // that: "multiple accesses in one thread racing with another
-            // thread's `memory.grow` that are in-bounds only after the grow
-            // commits may independently succeed or trap" (see
-            // https://github.com/WebAssembly/threads/issues/26#issuecomment-433930711).
-            // In other words, some non-determinism is acceptable when using
-            // `memory.size` on work being done by `memory.grow`.
-            inner
-                .def
-                .0
-                .current_length
-                .store(new_size_in_bytes, Ordering::SeqCst);
-        }
-        Ok(result)
+        SharedMemory::grow(self, delta_pages, store)
     }
 
     fn grow_to(&mut self, size: usize) -> Result<()> {
-        self.0.lock.write().unwrap().memory.grow_to(size)
+        self.0.memory.write().unwrap().grow_to(size)
     }
 
     fn vmmemory(&mut self) -> VMMemoryDefinition {
@@ -663,7 +614,7 @@ impl RuntimeLinearMemory for SharedMemory {
     }
 
     fn needs_init(&self) -> bool {
-        self.0.lock.read().unwrap().memory.needs_init()
+        self.0.memory.read().unwrap().needs_init()
     }
 
     fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
@@ -872,12 +823,78 @@ impl Memory {
 
     /// If the [Memory] is a [SharedMemory], unwrap it and return a clone to
     /// that shared memory.
-    pub fn as_shared_memory(&mut self) -> Option<SharedMemory> {
+    pub fn as_shared_memory(&mut self) -> Option<&mut SharedMemory> {
         let as_any = self.0.as_any_mut();
         if let Some(m) = as_any.downcast_mut::<SharedMemory>() {
-            Some(m.clone())
+            Some(m)
         } else {
             None
         }
     }
+
+    /// Implementation of `memory.atomic.notify` for all memories.
+    pub fn atomic_notify(&mut self, addr: u64, count: u32) -> Result<u32, Trap> {
+        match self.0.as_any_mut().downcast_mut::<SharedMemory>() {
+            Some(m) => m.atomic_notify(addr, count),
+            None => {
+                validate_atomic_addr(&self.vmmemory(), addr, 4, 4)?;
+                Ok(0)
+            }
+        }
+    }
+
+    /// Implementation of `memory.atomic.wait32` for all memories.
+    pub fn atomic_wait32(
+        &mut self,
+        addr: u64,
+        expected: u32,
+        deadline: Option<Instant>,
+    ) -> Result<WaitResult, Trap> {
+        match self.0.as_any_mut().downcast_mut::<SharedMemory>() {
+            Some(m) => m.atomic_wait32(addr, expected, deadline),
+            None => {
+                validate_atomic_addr(&self.vmmemory(), addr, 4, 4)?;
+                Err(Trap::AtomicWaitNonSharedMemory)
+            }
+        }
+    }
+
+    /// Implementation of `memory.atomic.wait64` for all memories.
+    pub fn atomic_wait64(
+        &mut self,
+        addr: u64,
+        expected: u64,
+        deadline: Option<Instant>,
+    ) -> Result<WaitResult, Trap> {
+        match self.0.as_any_mut().downcast_mut::<SharedMemory>() {
+            Some(m) => m.atomic_wait64(addr, expected, deadline),
+            None => {
+                validate_atomic_addr(&self.vmmemory(), addr, 8, 8)?;
+                Err(Trap::AtomicWaitNonSharedMemory)
+            }
+        }
+    }
+}
+
+/// In the configurations where bounds checks were elided in JIT code (because
+/// we are using static memories with virtual memory guard pages) this manual
+/// check is here so we don't segfault from Rust. For other configurations,
+/// these checks are required anyways.
+fn validate_atomic_addr(
+    def: &VMMemoryDefinition,
+    addr: u64,
+    access_size: u64,
+    access_alignment: u64,
+) -> Result<*mut u8, Trap> {
+    debug_assert!(access_alignment.is_power_of_two());
+    if !(addr % access_alignment == 0) {
+        return Err(Trap::HeapMisaligned);
+    }
+
+    let length = u64::try_from(def.current_length()).unwrap();
+    if !(addr.saturating_add(access_size) < length) {
+        return Err(Trap::MemoryOutOfBounds);
+    }
+
+    Ok(def.base.wrapping_add(addr as usize))
 }
diff --git a/crates/runtime/src/parking_spot.rs b/crates/runtime/src/parking_spot.rs
index a958d6d75..dee1019b2 100644
--- a/crates/runtime/src/parking_spot.rs
+++ b/crates/runtime/src/parking_spot.rs
@@ -14,23 +14,11 @@
 #![deny(missing_docs)]
 #![deny(unsafe_code)]
 
+use crate::WaitResult;
 use std::collections::BTreeMap;
 use std::sync::{Arc, Condvar, Mutex};
 use std::time::Instant;
 
-/// Result of a park operation.
-#[derive(Copy, Clone, Eq, PartialEq, Debug)]
-pub enum ParkResult {
-    /// Unparked by another thread.
-    Unparked,
-
-    /// The validation callback returned false.
-    Invalid,
-
-    /// The timeout expired.
-    TimedOut,
-}
-
 #[derive(Default, Debug)]
 struct Spot {
     /// The number of threads parked on this spot.
@@ -57,7 +45,7 @@ impl ParkingSpot {
     /// `unpark_all` or `unpark` with the same key, the current thread will be unparked.
     ///
     /// The `validate` callback is called before parking.
-    /// If it returns `false`, the thread is not parked and `ParkResult::Invalid` is returned.
+    /// If it returns `false`, the thread is not parked and `WaitResult::Mismatch` is returned.
     ///
     /// The `timeout` argument specifies the maximum amount of time the thread will be parked.
     pub fn park(
@@ -65,7 +53,7 @@ impl ParkingSpot {
         key: u64,
         validate: impl FnOnce() -> bool,
         timeout: impl Into<Option<Instant>>,
-    ) -> ParkResult {
+    ) -> WaitResult {
         self.park_inner(key, validate, timeout.into())
     }
 
@@ -74,7 +62,7 @@ impl ParkingSpot {
         key: u64,
         validate: impl FnOnce() -> bool,
         timeout: Option<Instant>,
-    ) -> ParkResult {
+    ) -> WaitResult {
         let mut inner = self
             .inner
             .lock()
@@ -82,7 +70,7 @@ impl ParkingSpot {
 
         // check validation with lock held
         if !validate() {
-            return ParkResult::Invalid;
+            return WaitResult::Mismatch;
         }
 
         // clone the condvar, so we can move the lock
@@ -115,7 +103,14 @@ impl ParkingSpot {
 
             let spot = inner.get_mut(&key).expect("failed to get spot");
 
-            if !timed_out {
+            if timed_out {
+                if let Some(timeout) = timeout {
+                    if Instant::now() < timeout {
+                        // Did not sleep long enough, try again.
+                        continue;
+                    }
+                }
+            } else {
                 if spot.to_unpark == 0 {
                     continue;
                 }
@@ -136,10 +131,10 @@ impl ParkingSpot {
             }
 
             if timed_out {
-                return ParkResult::TimedOut;
+                return WaitResult::TimedOut;
             }
 
-            return ParkResult::Unparked;
+            return WaitResult::Ok;
         }
     }
 
@@ -147,6 +142,9 @@ impl ParkingSpot {
     ///
     /// Returns the number of threads that were actually unparked.
     pub fn unpark(&self, key: u64, n: u32) -> u32 {
+        if n == 0 {
+            return 0;
+        }
         let mut num_unpark = 0;
 
         self.with_lot(key, |spot| {
diff --git a/crates/runtime/src/vmcontext.rs b/crates/runtime/src/vmcontext.rs
index f954d70e0..ae15342ea 100644
--- a/crates/runtime/src/vmcontext.rs
+++ b/crates/runtime/src/vmcontext.rs
@@ -12,7 +12,7 @@ use std::ptr::NonNull;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::u32;
 pub use vm_host_func_context::VMHostFuncContext;
-use wasmtime_environ::{DefinedMemoryIndex, Trap};
+use wasmtime_environ::DefinedMemoryIndex;
 
 pub const VMCONTEXT_MAGIC: u32 = u32::from_le_bytes(*b"core");
 
@@ -248,30 +248,6 @@ impl VMMemoryDefinition {
             current_length: other.current_length().into(),
         }
     }
-
-    /// In the configurations where bounds checks were elided in JIT code (because
-    /// we are using static memories with virtual memory guard pages) this manual
-    /// check is here so we don't segfault from Rust. For other configurations,
-    /// these checks are required anyways.
-    pub fn validate_addr(
-        &self,
-        addr: u64,
-        access_size: u64,
-        access_alignment: u64,
-    ) -> Result<*const u8, Trap> {
-        debug_assert!(access_alignment.is_power_of_two());
-        if !(addr % access_alignment == 0) {
-            return Err(Trap::HeapMisaligned);
-        }
-
-        let length = u64::try_from(self.current_length()).unwrap();
-        if !(addr.saturating_add(access_size) < length) {
-            return Err(Trap::MemoryOutOfBounds);
-        }
-
-        // SAFETY: checked above that the address is in bounds
-        Ok(unsafe { self.base.add(addr as usize) })
-    }
 }
 
 #[cfg(test)]
diff --git a/crates/wasmtime/src/memory.rs b/crates/wasmtime/src/memory.rs
index 4d5cc788b..61d5ef3a2 100644
--- a/crates/wasmtime/src/memory.rs
+++ b/crates/wasmtime/src/memory.rs
@@ -1,13 +1,17 @@
 use crate::store::{StoreData, StoreOpaque, Stored};
 use crate::trampoline::generate_memory_export;
+use crate::Trap;
 use crate::{AsContext, AsContextMut, Engine, MemoryType, StoreContext, StoreContextMut};
 use anyhow::{bail, Result};
 use std::cell::UnsafeCell;
 use std::convert::TryFrom;
 use std::slice;
+use std::time::Instant;
 use wasmtime_environ::MemoryPlan;
 use wasmtime_runtime::{RuntimeLinearMemory, VMMemoryImport};
 
+pub use wasmtime_runtime::WaitResult;
+
 /// Error for out of bounds [`Memory`] access.
 #[derive(Debug)]
 #[non_exhaustive]
@@ -780,7 +784,7 @@ impl SharedMemory {
     /// the maximum limits of this memory. A
     /// [`ResourceLimiter`](crate::ResourceLimiter) is another example of
     /// preventing a memory to grow.
-    pub fn grow(&mut self, delta: u64) -> Result<u64> {
+    pub fn grow(&self, delta: u64) -> Result<u64> {
         match self.0.grow(delta, None)? {
             Some((old_size, _new_size)) => {
                 // For shared memory, the `VMMemoryDefinition` is updated inside
@@ -791,6 +795,87 @@ impl SharedMemory {
         }
     }
 
+    /// Equivalent of the WebAssembly `memory.atomic.notify` instruction for
+    /// this shared memory.
+    ///
+    /// This method allows embedders to notify threads blocked on the specified
+    /// `addr`, an index into wasm linear memory. Threads could include
+    /// wasm threads blocked on a `memory.atomic.wait*` instruction or embedder
+    /// threads blocked on [`SharedMemory::atomic_wait32`], for example.
+    ///
+    /// The `count` argument is the number of threads to wake up.
+    ///
+    /// This function returns the number of threads awoken.
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if `addr` is not within bounds or
+    /// not aligned to a 4-byte boundary.
+    pub fn atomic_notify(&self, addr: u64, count: u32) -> Result<u32, Trap> {
+        self.0.atomic_notify(addr, count)
+    }
+
+    /// Equivalent of the WebAssembly `memory.atomic.wait32` instruction for
+    /// this shared memory.
+    ///
+    /// This method allows embedders to block the current thread until notified
+    /// via the `memory.atomic.notify` instruction or the
+    /// [`SharedMemory::atomic_notify`] method, enabling synchronization with
+    /// the wasm guest as desired.
+    ///
+    /// The `expected` argument is the expected 32-bit value to be stored at
+    /// the byte address `addr` specified. The `addr` specified is an index
+    /// into this linear memory.
+    ///
+    /// The optional `timeout` argument is the point in time after which the
+    /// calling thread is guaranteed to be woken up. Blocking will not occur
+    /// past this point.
+    ///
+    /// This function returns one of three possible values:
+    ///
+    /// * `WaitResult::Ok` - this function, loaded the value at `addr`, found
+    ///   it was equal to `expected`, and then blocked (all as one atomic
+    ///   operation). The thread was then awoken with a `memory.atomic.notify`
+    ///   instruction or the [`SharedMemory::atomic_notify`] method.
+    /// * `WaitResult::Mismatch` - the value at `addr` was loaded but was not
+    ///   equal to `expected` so the thread did not block and immediately
+    ///   returned.
+    /// * `WaitResult::TimedOut` - all the steps of `Ok` happened, except this
+    ///   thread was woken up due to a timeout.
+    ///
+    /// This function will not return due to spurious wakeups.
+    ///
+    /// # Errors
+    ///
+    /// This function will return an error if `addr` is not within bounds or
+    /// not aligned to a 4-byte boundary.
+    pub fn atomic_wait32(
+        &self,
+        addr: u64,
+        expected: u32,
+        timeout: Option<Instant>,
+    ) -> Result<WaitResult, Trap> {
+        self.0.atomic_wait32(addr, expected, timeout)
+    }
+
+    /// Equivalent of the WebAssembly `memory.atomic.wait64` instruction for
+    /// this shared memory.
+    ///
+    /// For more information see [`SharedMemory::atomic_wait32`].
+    ///
+    /// # Errors
+    ///
+    /// Returns the same error as [`SharedMemory::atomic_wait32`] except that
+    /// the specified address must be 8-byte aligned instead of 4-byte aligned.
+    pub fn atomic_wait64(
+        &self,
+        addr: u64,
+        expected: u64,
+        timeout: Option<Instant>,
+    ) -> Result<WaitResult, Trap> {
+        self.0.atomic_wait64(addr, expected, timeout)
+    }
+
     /// Return a reference to the [`Engine`] used to configure the shared
     /// memory.
     pub(crate) fn engine(&self) -> &Engine {
@@ -824,7 +909,8 @@ impl SharedMemory {
             .unwrap();
         let shared_memory = memory
             .as_shared_memory()
-            .expect("unable to convert from a shared memory");
+            .expect("unable to convert from a shared memory")
+            .clone();
         Self(shared_memory, store.engine().clone())
     }
 }


Test patch:
diff --git a/tests/all/memory.rs b/tests/all/memory.rs
index 281a2e7ab..2f7663db9 100644
--- a/tests/all/memory.rs
+++ b/tests/all/memory.rs
@@ -1,5 +1,7 @@
 use anyhow::Result;
 use rayon::prelude::*;
+use std::sync::atomic::{AtomicU32, Ordering::SeqCst};
+use std::time::{Duration, Instant};
 use wasmtime::*;
 
 fn module(engine: &Engine) -> Result<Module> {
@@ -462,3 +464,105 @@ fn memory64_maximum_minimum() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn shared_memory_basics() -> Result<()> {
+    let engine = Engine::default();
+    assert!(SharedMemory::new(&engine, MemoryType::new(1, None)).is_err());
+    assert!(SharedMemory::new(&engine, MemoryType::new(1, Some(1))).is_err());
+    assert!(SharedMemory::new(&engine, MemoryType::new64(1, None)).is_err());
+    assert!(SharedMemory::new(&engine, MemoryType::new64(1, Some(1))).is_err());
+    assert!(SharedMemory::new(&engine, MemoryType::shared(1, 0)).is_err());
+
+    let memory = SharedMemory::new(&engine, MemoryType::shared(1, 1))?;
+    assert!(memory.ty().is_shared());
+    assert_eq!(memory.ty().minimum(), 1);
+    assert_eq!(memory.ty().maximum(), Some(1));
+    assert_eq!(memory.size(), 1);
+    assert_eq!(memory.data_size(), 65536);
+    assert_eq!(memory.data().len(), 65536);
+    assert!(memory.grow(1).is_err());
+
+    // misaligned
+    assert_eq!(memory.atomic_notify(1, 100), Err(Trap::HeapMisaligned));
+    assert_eq!(
+        memory.atomic_wait32(1, 100, None),
+        Err(Trap::HeapMisaligned)
+    );
+    assert_eq!(
+        memory.atomic_wait64(1, 100, None),
+        Err(Trap::HeapMisaligned)
+    );
+
+    // oob
+    assert_eq!(
+        memory.atomic_notify(1 << 20, 100),
+        Err(Trap::MemoryOutOfBounds)
+    );
+    assert_eq!(
+        memory.atomic_wait32(1 << 20, 100, None),
+        Err(Trap::MemoryOutOfBounds)
+    );
+    assert_eq!(
+        memory.atomic_wait64(1 << 20, 100, None),
+        Err(Trap::MemoryOutOfBounds)
+    );
+
+    // ok
+    assert_eq!(memory.atomic_notify(8, 100), Ok(0));
+    assert_eq!(memory.atomic_wait32(8, 1, None), Ok(WaitResult::Mismatch));
+    assert_eq!(memory.atomic_wait64(8, 1, None), Ok(WaitResult::Mismatch));
+
+    // timeout
+    let near_future = Instant::now() + Duration::new(0, 100);
+    assert_eq!(
+        memory.atomic_wait32(8, 0, Some(near_future)),
+        Ok(WaitResult::TimedOut)
+    );
+    assert_eq!(
+        memory.atomic_wait64(8, 0, Some(near_future)),
+        Ok(WaitResult::TimedOut)
+    );
+
+    Ok(())
+}
+
+#[test]
+fn shared_memory_wait_notify() -> Result<()> {
+    const THREADS: usize = 8;
+    const COUNT: usize = 100_000;
+
+    let engine = Engine::default();
+    let memory = SharedMemory::new(&engine, MemoryType::shared(1, 1))?;
+    let data = unsafe { &*(memory.data().as_ptr() as *const AtomicU32) };
+    let locked = unsafe { &*(memory.data().as_ptr().add(4) as *const AtomicU32) };
+
+    // Note that `SeqCst` is used here to not think much about the orderings
+    // here, and it also somewhat more closely mirrors what's happening in wasm.
+    let lock = || {
+        while locked.swap(1, SeqCst) == 1 {
+            memory.atomic_wait32(0, 1, None).unwrap();
+        }
+    };
+    let unlock = || {
+        locked.store(0, SeqCst);
+        memory.atomic_notify(0, 1).unwrap();
+    };
+
+    std::thread::scope(|s| {
+        for _ in 0..THREADS {
+            s.spawn(|| {
+                for _ in 0..COUNT {
+                    lock();
+                    let next = data.load(SeqCst) + 1;
+                    data.store(next, SeqCst);
+                    unlock();
+                }
+            });
+        }
+    });
+
+    assert_eq!(data.load(SeqCst), (THREADS * COUNT) as u32);
+
+    Ok(())
+}
diff --git a/tests/all/threads.rs b/tests/all/threads.rs
index fe497beec..30ca86bdd 100644
--- a/tests/all/threads.rs
+++ b/tests/all/threads.rs
@@ -107,7 +107,7 @@ fn test_probe_shared_memory_size() -> Result<()> {
     let mut store = Store::new(&engine, ());
     let instance = Instance::new(&mut store, &module, &[])?;
     let size_fn = instance.get_typed_func::<(), i32>(&mut store, "size")?;
-    let mut shared_memory = instance.get_shared_memory(&mut store, "memory").unwrap();
+    let shared_memory = instance.get_shared_memory(&mut store, "memory").unwrap();
 
     assert_eq!(size_fn.call(&mut store, ())?, 1);
     assert_eq!(shared_memory.size(), 1);
@@ -244,7 +244,7 @@ fn test_memory_size_accessibility() -> Result<()> {
     let shared_memory = SharedMemory::new(&engine, MemoryType::shared(1, NUM_GROW_OPS as u32))?;
     let done = Arc::new(AtomicBool::new(false));
 
-    let mut grow_memory = shared_memory.clone();
+    let grow_memory = shared_memory.clone();
     let grow_thread = std::thread::spawn(move || {
         for i in 0..NUM_GROW_OPS {
             if grow_memory.grow(1).is_err() {