Add GC zeal assertions (#12933)

* Add GC zeal assertions

- Poison freed GC objects and new heap memory
- Assert newly-allocated objects are filled with poison pattern
- Add `gc_assert!` checks for valid `VMGcKind` on `GcHeap::index[_mut]`
- Add `gc_assert!` checks for valid `VMGcKind` during tracing
- Add `VMGcKind::try_from_u32()` for fallible kind validation
- Add over-approximated stack roots list integrity checks, called before and
  after trace and sweep. Validates kind, in-list bit, ref count, and that the
  list is not cyclic.
- Add assertion that all free blocks of memory contain the poison pattern,
  before and after trace and sweep

* review feedback

Author: fitzgen

crates/cranelift/src/func_environ/gc/enabled.rs

@@ -94,8 +94,55 @@ fn unbarriered_store_gc_ref(
     Ok(())
 }
 
-/// Emit code to read a struct field or array element from its raw address in
-/// the GC heap.
+/// Emit inline CLIF code that asserts an object's `VMGcKind` matches the
+/// expected kind. Only emits code when `cfg(gc_zeal)` is enabled.
+///
+/// `gc_ref` must be a non-null, non-i31 GC reference (i32 heap index).
+fn emit_gc_kind_assert(
+    func_env: &mut FuncEnvironment<'_>,
+    builder: &mut FunctionBuilder<'_>,
+    gc_ref: ir::Value,
+    expected_kind: VMGcKind,
+) {
+    if !cfg!(gc_zeal) {
+        return;
+    }
+
+    func_env.trapz(builder, gc_ref, crate::TRAP_NULL_REFERENCE);
+
+    let kind_addr = func_env.prepare_gc_ref_access(
+        builder,
+        gc_ref,
+        BoundsCheck::StaticObjectField {
+            offset: wasmtime_environ::VM_GC_HEADER_KIND_OFFSET,
+            access_size: wasmtime_environ::VM_GC_KIND_SIZE,
+            object_size: wasmtime_environ::VM_GC_HEADER_SIZE,
+        },
+    );
+    let kind_and_reserved_bits = builder.ins().load(
+        ir::types::I32,
+        ir::MemFlags::trusted().with_readonly(),
+        kind_addr,
+        0,
+    );
+    let kind_mask = builder
+        .ins()
+        .iconst(ir::types::I32, i64::from(VMGcKind::MASK));
+    let actual_kind = builder.ins().band(kind_and_reserved_bits, kind_mask);
+
+    let expected_kind = builder
+        .ins()
+        .iconst(ir::types::I32, i64::from(expected_kind.as_u32()));
+
+    // NB: Do a subtype check rather than a strict equality check. See
+    // `VMGcKind::matches` for details.
+    let and = builder.ins().band(actual_kind, expected_kind);
+    let matches = builder.ins().icmp(IntCC::Equal, and, expected_kind);
+
+    builder.ins().trapz(matches, TRAP_INTERNAL_ASSERT);
+}
+
+/// Read a struct field or array element from its raw address in the GC heap.
 ///
 /// The given address MUST have already been bounds-checked via
 /// `prepare_gc_ref_access`.
@@ -331,6 +378,8 @@ pub fn translate_struct_get(
     // type info from `wasmparser` and through to here is a bit funky.
     func_env.trapz(builder, struct_ref, crate::TRAP_NULL_REFERENCE);
 
+    emit_gc_kind_assert(func_env, builder, struct_ref, VMGcKind::StructRef);
+
     let field_index = usize::try_from(field_index).unwrap();
     let interned_type_index = func_env.module.types[struct_type_index].unwrap_module_type_index();
 
@@ -378,6 +427,8 @@ pub fn translate_struct_set(
     // TODO: See comment in `translate_struct_get` about the `trapz`.
     func_env.trapz(builder, struct_ref, crate::TRAP_NULL_REFERENCE);
 
+    emit_gc_kind_assert(func_env, builder, struct_ref, VMGcKind::StructRef);
+
     let field_index = usize::try_from(field_index).unwrap();
     let interned_type_index = func_env.module.types[struct_type_index].unwrap_module_type_index();
 
@@ -979,6 +1030,8 @@ pub fn translate_array_get(
 ) -> WasmResult<ir::Value> {
     log::trace!("translate_array_get({array_type_index:?}, {array_ref:?}, {index:?})");
 
+    emit_gc_kind_assert(func_env, builder, array_ref, VMGcKind::ArrayRef);
+
     let array_type_index = func_env.module.types[array_type_index].unwrap_module_type_index();
     let elem_addr = array_elem_addr(func_env, builder, array_type_index, array_ref, index);
 
@@ -1000,6 +1053,8 @@ pub fn translate_array_set(
 ) -> WasmResult<()> {
     log::trace!("translate_array_set({array_type_index:?}, {array_ref:?}, {index:?}, {value:?})");
 
+    emit_gc_kind_assert(func_env, builder, array_ref, VMGcKind::ArrayRef);
+
     let array_type_index = func_env.module.types[array_type_index].unwrap_module_type_index();
     let elem_addr = array_elem_addr(func_env, builder, array_type_index, array_ref, index);
 

crates/environ/src/gc.rs

@@ -499,15 +499,8 @@ impl VMGcKind {
     #[inline]
     pub fn from_high_bits_of_u32(val: u32) -> VMGcKind {
         let masked = val & Self::MASK;
-        let result = match masked {
-            x if x == Self::ExternRef.as_u32() => Self::ExternRef,
-            x if x == Self::AnyRef.as_u32() => Self::AnyRef,
-            x if x == Self::EqRef.as_u32() => Self::EqRef,
-            x if x == Self::ArrayRef.as_u32() => Self::ArrayRef,
-            x if x == Self::StructRef.as_u32() => Self::StructRef,
-            x if x == Self::ExnRef.as_u32() => Self::ExnRef,
-            _ => panic!("invalid `VMGcKind`: {masked:#032b}"),
-        };
+        let result = Self::try_from_u32(masked)
+            .unwrap_or_else(|| panic!("invalid `VMGcKind`: {masked:#032b}"));
 
         let poison_kind = u32::from_le_bytes([POISON, POISON, POISON, POISON]) & VMGcKind::MASK;
         debug_assert_ne!(
@@ -531,6 +524,22 @@ impl VMGcKind {
     pub fn as_u32(self) -> u32 {
         self as u32
     }
+
+    /// Try to convert a `u32` into a `VMGcKind`.
+    ///
+    /// Returns `None` if the value doesn't match any known kind.
+    #[inline]
+    pub fn try_from_u32(x: u32) -> Option<VMGcKind> {
+        match x {
+            _ if x == Self::ExternRef.as_u32() => Some(Self::ExternRef),
+            _ if x == Self::AnyRef.as_u32() => Some(Self::AnyRef),
+            _ if x == Self::EqRef.as_u32() => Some(Self::EqRef),
+            _ if x == Self::ArrayRef.as_u32() => Some(Self::ArrayRef),
+            _ if x == Self::StructRef.as_u32() => Some(Self::StructRef),
+            _ if x == Self::ExnRef.as_u32() => Some(Self::ExnRef),
+            _ => None,
+        }
+    }
 }
 
 #[cfg(test)]

crates/wasmtime/src/runtime/store/gc.rs

@@ -56,11 +56,16 @@ impl StoreOpaque {
         bytes_needed: Option<u64>,
         asyncness: Asyncness,
     ) {
-        if let Some(n) = bytes_needed {
+        if let Some(n) = bytes_needed
+            // The gc_zeal's allocation counter will pass `bytes_needed == 0` to
+            // signify that we shouldn't grow the GC heap, just do a collection.
+            && n > 0
+        {
             if self.grow_gc_heap(limiter, n).await.is_ok() {
                 return;
             }
         }
+
         self.do_gc(asyncness).await;
     }
 

crates/wasmtime/src/runtime/vm/gc.rs

@@ -48,18 +48,48 @@ pub struct GcStore {
 
     /// The function-references table for this GC heap.
     pub func_ref_table: FuncRefTable,
+
+    /// An allocation counter that triggers GC when it reaches zero.
+    ///
+    /// Initialized from the `WASMTIME_GC_ZEAL_ALLOC_COUNTER` environment
+    /// variable. Decremented on every allocation and when it hits zero, a GC is
+    /// forced and the counter is reset.
+    #[cfg(all(gc_zeal, feature = "std"))]
+    gc_zeal_alloc_counter: Option<NonZeroU32>,
+
+    /// The initial value to reset the counter to after it triggers.
+    #[cfg(all(gc_zeal, feature = "std"))]
+    gc_zeal_alloc_counter_init: Option<NonZeroU32>,
 }
 
 impl GcStore {
     /// Create a new `GcStore`.
     pub fn new(allocation_index: GcHeapAllocationIndex, gc_heap: Box<dyn GcHeap>) -> Self {
         let host_data_table = ExternRefHostDataTable::default();
         let func_ref_table = FuncRefTable::default();
+
+        #[cfg(all(gc_zeal, feature = "std"))]
+        let gc_zeal_alloc_counter_init =
+            std::env::var("WASMTIME_GC_ZEAL_ALLOC_COUNTER")
+                .ok()
+                .map(|v| {
+                    v.parse::<NonZeroU32>().unwrap_or_else(|_| {
+                        panic!(
+                            "`WASMTIME_GC_ZEAL_ALLOC_COUNTER` must be a non-zero \
+                             `u32` value, got: {v}"
+                        )
+                    })
+                });
+
         Self {
             allocation_index,
             gc_heap,
             host_data_table,
             func_ref_table,
+            #[cfg(all(gc_zeal, feature = "std"))]
+            gc_zeal_alloc_counter: gc_zeal_alloc_counter_init,
+            #[cfg(all(gc_zeal, feature = "std"))]
+            gc_zeal_alloc_counter_init,
         }
     }
 
@@ -231,6 +261,20 @@ impl GcStore {
         header: VMGcHeader,
         layout: Layout,
     ) -> Result<Result<VMGcRef, u64>> {
+        // When gc_zeal is enabled with an allocation counter, decrement it and
+        // force a GC cycle when it reaches zero by returning a fake OOM.
+        #[cfg(all(gc_zeal, feature = "std"))]
+        if let Some(counter) = self.gc_zeal_alloc_counter.take() {
+            match NonZeroU32::new(counter.get() - 1) {
+                Some(c) => self.gc_zeal_alloc_counter = Some(c),
+                None => {
+                    log::trace!("gc_zeal: allocation counter reached zero, forcing GC");
+                    self.gc_zeal_alloc_counter = self.gc_zeal_alloc_counter_init;
+                    return Ok(Err(0));
+                }
+            }
+        }
+
         self.gc_heap.alloc_raw(header, layout)
     }