review feedback and fix tests

fitzgen · fitzgen · commit e5234e6bda8c · 2025-04-09T16:04:00.000-07:00
diff --git a/crates/wasmtime/src/runtime/gc/enabled/arrayref.rs b/crates/wasmtime/src/runtime/gc/enabled/arrayref.rs
@@ -306,11 +306,6 @@ impl ArrayRef {
         elem: &Val,
         len: u32,
     ) -> Result<Rooted<ArrayRef>> {
-        assert_eq!(
-            store.id(),
-            allocator.store_id,
-            "attempted to use a `ArrayRefPre` with the wrong store"
-        );
         assert!(
             !store.async_support(),
             " use `ArrayRef::new_async` with asynchronous stores"
@@ -374,11 +369,6 @@ impl ArrayRef {
         elem: &Val,
         len: u32,
     ) -> Result<Rooted<ArrayRef>> {
-        assert_eq!(
-            store.id(),
-            allocator.store_id,
-            "attempted to use a `ArrayRefPre` with the wrong store"
-        );
         assert!(
             store.async_support(),
             "use `ArrayRef::new` with synchronous stores"
@@ -395,13 +385,38 @@ impl ArrayRef {
             .await
     }
 
+    /// Like `ArrayRef::new` but when async is configured must only ever be
+    /// called from on a fiber stack.
+    pub(crate) unsafe fn new_maybe_async(
+        store: &mut StoreOpaque,
+        allocator: &ArrayRefPre,
+        elem: &Val,
+        len: u32,
+    ) -> Result<Rooted<ArrayRef>> {
+        // Type check the initial element value against the element type.
+        elem.ensure_matches_ty(store, allocator.ty.element_type().unpack())
+            .context("element type mismatch")?;
+
+        unsafe {
+            store.retry_after_gc_maybe_async((), |store, ()| {
+                Self::_new_unchecked(store, allocator, RepeatN(elem, len))
+            })
+        }
+    }
+
     /// Allocate a new array of the given elements, without checking that the
     /// elements' types match the array's element type.
     fn _new_unchecked<'a>(
         store: &mut StoreOpaque,
         allocator: &ArrayRefPre,
         elems: impl ExactSizeIterator<Item = &'a Val>,
     ) -> Result<Rooted<ArrayRef>> {
+        assert_eq!(
+            store.id(),
+            allocator.store_id,
+            "attempted to use a `ArrayRefPre` with the wrong store"
+        );
+
         let len = u32::try_from(elems.len()).unwrap();
 
         // Allocate the array and write each field value into the appropriate
@@ -572,6 +587,34 @@ impl ArrayRef {
             .await
     }
 
+    /// Like `ArrayRef::new_fixed[_async]` but it is the caller's responsibility
+    /// to ensure that when async is enabled, this is only called from on a
+    /// fiber stack.
+    #[cfg(feature = "async")]
+    pub(crate) unsafe fn _new_fixed_maybe_async(
+        store: &mut StoreOpaque,
+        allocator: &ArrayRefPre,
+        elems: &[Val],
+    ) -> Result<Rooted<ArrayRef>> {
+        assert_eq!(
+            store.id(),
+            allocator.store_id,
+            "attempted to use a `ArrayRefPre` with the wrong store"
+        );
+
+        // Type check the elements against the element type.
+        for elem in elems {
+            elem.ensure_matches_ty(store, allocator.ty.element_type().unpack())
+                .context("element type mismatch")?;
+        }
+
+        unsafe {
+            store.retry_after_gc_maybe_async((), |store, ()| {
+                Self::_new_unchecked(store, allocator, elems.iter())
+            })
+        }
+    }
+
     #[inline]
     pub(crate) fn comes_from_same_store(&self, store: &StoreOpaque) -> bool {
         self.inner.comes_from_same_store(store)
diff --git a/crates/wasmtime/src/runtime/gc/enabled/structref.rs b/crates/wasmtime/src/runtime/gc/enabled/structref.rs
@@ -240,11 +240,6 @@ impl StructRef {
         allocator: &StructRefPre,
         fields: &[Val],
     ) -> Result<Rooted<StructRef>> {
-        assert_eq!(
-            store.id(),
-            allocator.store_id,
-            "attempted to use a `StructRefPre` with the wrong store"
-        );
         assert!(
             !store.async_support(),
             "use `StructRef::new_async` with asynchronous stores"
@@ -296,11 +291,6 @@ impl StructRef {
         allocator: &StructRefPre,
         fields: &[Val],
     ) -> Result<Rooted<StructRef>> {
-        assert_eq!(
-            store.id(),
-            allocator.store_id,
-            "attempted to use a `StructRefPre` with the wrong store"
-        );
         assert!(
             store.async_support(),
             "use `StructRef::new` with synchronous stores"
@@ -313,6 +303,25 @@ impl StructRef {
             .await
     }
 
+    /// Like `Self::new` but caller's must ensure that if the store is
+    /// configured for async, this is only ever called from on a fiber stack.
+    pub(crate) unsafe fn new_maybe_async(
+        store: &mut StoreOpaque,
+        allocator: &StructRefPre,
+        fields: &[Val],
+    ) -> Result<Rooted<StructRef>> {
+        assert!(
+            store.async_support(),
+            "use `StructRef::new` with synchronous stores"
+        );
+        Self::type_check_fields(store, allocator, fields)?;
+        unsafe {
+            store.retry_after_gc_maybe_async((), |store, ()| {
+                Self::new_unchecked(store, allocator, fields)
+            })
+        }
+    }
+
     /// Type check the field values before allocating a new struct.
     fn type_check_fields(
         store: &mut StoreOpaque,
@@ -346,6 +355,12 @@ impl StructRef {
         allocator: &StructRefPre,
         fields: &[Val],
     ) -> Result<Rooted<StructRef>> {
+        assert_eq!(
+            store.id(),
+            allocator.store_id,
+            "attempted to use a `StructRefPre` with the wrong store"
+        );
+
         // Allocate the struct and write each field value into the appropriate
         // offset.
         let structref = store
diff --git a/crates/wasmtime/src/runtime/store.rs b/crates/wasmtime/src/runtime/store.rs
@@ -1543,6 +1543,64 @@ impl StoreOpaque {
         }
     }
 
+    /// Like `retry_after_gc` but must be called from a fiber stack if this is
+    /// an async store.
+    #[cfg(feature = "gc")]
+    pub(crate) unsafe fn retry_after_gc_maybe_async<T, U>(
+        &mut self,
+        value: T,
+        alloc_func: impl Fn(&mut Self, T) -> Result<U>,
+    ) -> Result<U>
+    where
+        T: Send + Sync + 'static,
+    {
+        match alloc_func(self, value) {
+            Ok(x) => Ok(x),
+            Err(e) => match e.downcast::<crate::GcHeapOutOfMemory<T>>() {
+                Ok(oom) => {
+                    let value = oom.into_inner();
+                    self.maybe_async_gc(None)?;
+                    alloc_func(self, value)
+                }
+                Err(e) => Err(e),
+            },
+        }
+    }
+
+    /// Like `gc` but must be called from a fiber stack if this is an async
+    /// store.
+    unsafe fn maybe_async_gc(&mut self, root: Option<VMGcRef>) -> Result<Option<VMGcRef>> {
+        let mut scope = crate::OpaqueRootScope::new(self);
+        let store_id = scope.id();
+        let root = root.map(|r| scope.gc_roots_mut().push_lifo_root(store_id, r));
+
+        if scope.async_support() {
+            #[cfg(feature = "async")]
+            unsafe {
+                let async_cx = scope.async_cx();
+                let future = scope.gc_async();
+                async_cx
+                    .expect("attempted to pull async context during shutdown")
+                    .block_on(future)?;
+            }
+        } else {
+            scope.gc();
+        }
+
+        let root = match root {
+            None => None,
+            Some(r) => {
+                let r = r
+                    .get_gc_ref(&scope)
+                    .expect("still in scope")
+                    .unchecked_copy();
+                Some(scope.gc_store_mut()?.clone_gc_ref(&r))
+            }
+        };
+
+        Ok(root)
+    }
+
     #[cfg(feature = "gc")]
     pub fn gc(&mut self) {
         // If the GC heap hasn't been initialized, there is nothing to collect.
@@ -2096,41 +2154,12 @@ unsafe impl<T> crate::runtime::vm::VMStore for StoreInner<T> {
     }
 
     #[cfg(feature = "gc")]
-    fn maybe_async_gc(&mut self, root: Option<VMGcRef>) -> Result<Option<VMGcRef>> {
-        let mut scope = crate::RootScope::new(self);
-        let store = scope.as_context_mut().0;
-        let store_id = store.id();
-        let root = root.map(|r| store.gc_roots_mut().push_lifo_root(store_id, r));
-
-        if store.async_support() {
-            #[cfg(feature = "async")]
-            unsafe {
-                let async_cx = store.async_cx();
-                let future = store.gc_async();
-                async_cx
-                    .expect("attempted to pull async context during shutdown")
-                    .block_on(future)?;
-            }
-        } else {
-            (**store).gc();
-        }
-
-        let root = match root {
-            None => None,
-            Some(r) => {
-                let r = r
-                    .get_gc_ref(store)
-                    .expect("still in scope")
-                    .unchecked_copy();
-                Some(store.gc_store_mut()?.clone_gc_ref(&r))
-            }
-        };
-
-        Ok(root)
+    unsafe fn maybe_async_gc(&mut self, root: Option<VMGcRef>) -> Result<Option<VMGcRef>> {
+        (**self).maybe_async_gc(root)
     }
 
     #[cfg(not(feature = "gc"))]
-    fn maybe_async_gc(&mut self, root: Option<VMGcRef>) -> Result<Option<VMGcRef>> {
+    unsafe fn maybe_async_gc(&mut self, root: Option<VMGcRef>) -> Result<Option<VMGcRef>> {
         Ok(root)
     }
 
diff --git a/crates/wasmtime/src/runtime/vm.rs b/crates/wasmtime/src/runtime/vm.rs
@@ -200,7 +200,7 @@ pub unsafe trait VMStore {
     ///
     /// If the async GC was cancelled, returns an error. This should be raised
     /// as a trap to clean up Wasm execution.
-    fn maybe_async_gc(&mut self, root: Option<VMGcRef>) -> Result<Option<VMGcRef>>;
+    unsafe fn maybe_async_gc(&mut self, root: Option<VMGcRef>) -> Result<Option<VMGcRef>>;
 
     /// Metadata required for resources for the component model.
     #[cfg(feature = "component-model")]
diff --git a/crates/wasmtime/src/runtime/vm/const_expr.rs b/crates/wasmtime/src/runtime/vm/const_expr.rs
@@ -73,7 +73,7 @@ impl<'a> ConstEvalContext<'a> {
             .collect::<Vec<_>>();
 
         let allocator = StructRefPre::_new(store, struct_ty);
-        let struct_ref = StructRef::_new(store, &allocator, &fields)?;
+        let struct_ref = unsafe { StructRef::new_maybe_async(store, &allocator, &fields)? };
         let raw = struct_ref.to_anyref()._to_raw(store)?;
         Ok(ValRaw::anyref(raw))
     }
@@ -141,6 +141,9 @@ impl ConstExprEvaluator {
     /// global values of the expected types. This evaluator operates directly on
     /// untyped `ValRaw`s and does not and cannot check that its operands are of
     /// the correct type.
+    ///
+    /// If given async store, then this must be called from on an async fiber
+    /// stack.
     pub unsafe fn eval(
         &mut self,
         store: &mut StoreOpaque,
@@ -269,7 +272,7 @@ impl ConstExprEvaluator {
                     let elem = Val::_from_raw(&mut store, self.pop()?, ty.element_type().unpack());
 
                     let pre = ArrayRefPre::_new(&mut store, ty);
-                    let array = ArrayRef::_new(&mut store, &pre, &elem, len)?;
+                    let array = unsafe { ArrayRef::new_maybe_async(&mut store, &pre, &elem, len)? };
 
                     self.stack
                         .push(ValRaw::anyref(array.to_anyref()._to_raw(&mut store)?));
@@ -288,7 +291,7 @@ impl ConstExprEvaluator {
                         .expect("type should have a default value");
 
                     let pre = ArrayRefPre::_new(&mut store, ty);
-                    let array = ArrayRef::_new(&mut store, &pre, &elem, len)?;
+                    let array = unsafe { ArrayRef::new_maybe_async(&mut store, &pre, &elem, len)? };
 
                     self.stack
                         .push(ValRaw::anyref(array.to_anyref()._to_raw(&mut store)?));
@@ -323,7 +326,8 @@ impl ConstExprEvaluator {
                         .collect::<SmallVec<[_; 8]>>();
 
                     let pre = ArrayRefPre::_new(&mut store, ty);
-                    let array = ArrayRef::_new_fixed(&mut store, &pre, &elems)?;
+                    let array =
+                        unsafe { ArrayRef::_new_fixed_maybe_async(&mut store, &pre, &elems)? };
 
                     self.stack
                         .push(ValRaw::anyref(array.to_anyref()._to_raw(&mut store)?));
diff --git a/crates/wasmtime/src/runtime/vm/libcalls.rs b/crates/wasmtime/src/runtime/vm/libcalls.rs
@@ -772,7 +772,7 @@ unsafe fn array_new_elem(
     use crate::{
         store::AutoAssertNoGc,
         vm::const_expr::{ConstEvalContext, ConstExprEvaluator},
-        ArrayRef, ArrayRefPre, ArrayType, Func, GcHeapOutOfMemory, RootSet, RootedGcRefImpl, Val,
+        ArrayRef, ArrayRefPre, ArrayType, Func, RootSet, RootedGcRefImpl, Val,
     };
     use wasmtime_environ::{ModuleInternedTypeIndex, TableSegmentElements};
 
@@ -827,16 +827,7 @@ unsafe fn array_new_elem(
             }
         }
 
-        let array = match ArrayRef::_new_fixed(store, &pre, &vals) {
-            Ok(a) => a,
-            Err(e) if e.is::<GcHeapOutOfMemory<()>>() => {
-                // Collect garbage to hopefully free up space, then try the
-                // allocation again.
-                store.maybe_async_gc(None)?;
-                ArrayRef::_new_fixed(store, &pre, &vals)?
-            }
-            Err(e) => return Err(e),
-        };
+        let array = unsafe { ArrayRef::_new_fixed_maybe_async(store, &pre, &vals)? };
 
         let mut store = AutoAssertNoGc::new(store);
         let gc_ref = array.try_clone_gc_ref(&mut store)?;
diff --git a/crates/wast/src/core.rs b/crates/wast/src/core.rs
@@ -1,11 +1,12 @@
+use crate::WastContext;
 use anyhow::{anyhow, bail, Context, Result};
 use std::fmt::{Display, LowerHex};
 use wasmtime::{AnyRef, ExternRef, Store, Val};
 use wast::core::{AbstractHeapType, HeapType, NanPattern, V128Pattern, WastArgCore, WastRetCore};
 use wast::token::{F32, F64};
 
 /// Translate from a `script::Value` to a `RuntimeValue`.
-pub fn val<T>(store: &mut Store<T>, v: &WastArgCore<'_>) -> Result<Val> {
+pub fn val<T>(ctx: &mut WastContext<T>, v: &WastArgCore<'_>) -> Result<Val> {
     use wast::core::WastArgCore::*;
 
     Ok(match v {
@@ -30,10 +31,18 @@ pub fn val<T>(store: &mut Store<T>, v: &WastArgCore<'_>) -> Result<Val> {
             shared: false,
             ty: AbstractHeapType::None,
         }) => Val::AnyRef(None),
-        RefExtern(x) => Val::ExternRef(Some(ExternRef::new(store, *x)?)),
+        RefExtern(x) => Val::ExternRef(if let Some(rt) = ctx.async_runtime.as_ref() {
+            Some(rt.block_on(ExternRef::new_async(&mut ctx.store, *x))?)
+        } else {
+            Some(ExternRef::new(&mut ctx.store, *x)?)
+        }),
         RefHost(x) => {
-            let x = ExternRef::new(&mut *store, *x)?;
-            let x = AnyRef::convert_extern(&mut *store, x)?;
+            let x = if let Some(rt) = ctx.async_runtime.as_ref() {
+                rt.block_on(ExternRef::new_async(&mut ctx.store, *x))?
+            } else {
+                ExternRef::new(&mut ctx.store, *x)?
+            };
+            let x = AnyRef::convert_extern(&mut ctx.store, x)?;
             Val::AnyRef(Some(x))
         }
         other => bail!("couldn't convert {:?} to a runtime value", other),
diff --git a/crates/wast/src/wast.rs b/crates/wast/src/wast.rs
