Transitively dec-ref objects in the DRC collector

Previously, the DRC collector only performed shallow reference counting where if
an object's ref count was decremented and reached zero, then we would deallocate
that object, but we would not decrement the reference counts of any other object
that was referenced by it. This commit changes that, so that now we will
decrement the reference counts of other objects referenced by an object that is
being deallocated.

This requires some method of tracing an object's outgoing edges. To avoid
needing to keep around and dispatch upon type-specific information about which
struct fields are GC references, for example, we instead pack all GC references
at the start of the object, just after its header. We additionally keep a count
of how many GC references an object has in its header. This allows us to create
a slice of any object's GC references, and we can uniformly trace any GC
object's outgoing edges. We don't need to reflect on the GC object's actual
type.

A final detail: the DRC collector was previously storing the object size in the
common GC header's reserved bits. It now stores the number of GC refs there
instead, and the object size is stored in its own field in the DRC-specific
header, next to the ref count. This is nice because, since we are making the DRC
header larger anyways, it raises the DRC collector's implementation limit on
object size from `u27::MAX` to `u32::MAX`.

Fixes bytecodealliance#9701

Author: fitzgen

crates/cranelift/src/gc/enabled.rs

@@ -316,7 +316,7 @@ pub fn translate_struct_get(
     let struct_size = struct_layout.size;
     let struct_size_val = builder.ins().iconst(ir::types::I32, i64::from(struct_size));
 
-    let field_offset = struct_layout.fields[field_index];
+    let field_offset = struct_layout.fields[field_index].offset;
     let field_ty = &func_env.types.unwrap_struct(interned_type_index)?.fields[field_index];
     let field_size = wasmtime_environ::byte_size_of_wasm_ty_in_gc_heap(&field_ty.element_type);
     assert!(field_offset + field_size <= struct_size);
@@ -361,7 +361,7 @@ pub fn translate_struct_set(
     let struct_size = struct_layout.size;
     let struct_size_val = builder.ins().iconst(ir::types::I32, i64::from(struct_size));
 
-    let field_offset = struct_layout.fields[field_index];
+    let field_offset = struct_layout.fields[field_index].offset;
     let field_ty = &func_env.types.unwrap_struct(interned_type_index)?.fields[field_index];
     let field_size = wasmtime_environ::byte_size_of_wasm_ty_in_gc_heap(&field_ty.element_type);
     assert!(field_offset + field_size <= struct_size);
@@ -1194,7 +1194,7 @@ fn initialize_struct_fields(
 ) -> WasmResult<()> {
     let struct_layout = func_env.struct_layout(struct_ty);
     let struct_size = struct_layout.size;
-    let field_offsets: SmallVec<[_; 8]> = struct_layout.fields.iter().copied().collect();
+    let field_offsets: SmallVec<[_; 8]> = struct_layout.fields.iter().map(|f| f.offset).collect();
     assert_eq!(field_offsets.len(), field_values.len());
 
     assert!(!func_env.types[struct_ty].composite_type.shared);

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

@@ -261,12 +261,13 @@ impl DrcCompiler {
 
 /// Emit CLIF to call the `gc_raw_alloc` libcall.
 ///
-/// It is the caller's responsibility to ensure that `size` fits within the
-/// `VMGcKind`'s unused bits.
+/// It is the caller's responsibility to ensure that `num_gc_refs` fits within
+/// the `VMGcKind`'s unused bits.
 fn emit_gc_raw_alloc(
     func_env: &mut FuncEnvironment<'_>,
     builder: &mut FunctionBuilder<'_>,
     kind: VMGcKind,
+    num_gc_refs: ir::Value,
     ty: ModuleInternedTypeIndex,
     size: ir::Value,
     align: u32,
@@ -277,15 +278,17 @@ fn emit_gc_raw_alloc(
     let kind = builder
         .ins()
         .iconst(ir::types::I32, i64::from(kind.as_u32()));
+    let kind_and_reserved = builder.ins().bor(kind, num_gc_refs);
 
     let ty = builder.ins().iconst(ir::types::I32, i64::from(ty.as_u32()));
 
     assert!(align.is_power_of_two());
     let align = builder.ins().iconst(ir::types::I32, i64::from(align));
 
-    let call_inst = builder
-        .ins()
-        .call(gc_alloc_raw_builtin, &[vmctx, kind, ty, size, align]);
+    let call_inst = builder.ins().call(
+        gc_alloc_raw_builtin,
+        &[vmctx, kind_and_reserved, ty, size, align],
+    );
 
     let gc_ref = builder.func.dfg.first_result(call_inst);
     let gc_ref = builder.ins().ireduce(ir::types::I32, gc_ref);
@@ -318,13 +321,19 @@ impl GcCompiler for DrcCompiler {
         // First, compute the array's total size from its base size, element
         // size, and length.
         let size = emit_array_size(func_env, builder, &array_layout, init);
+        let num_gc_refs = if array_layout.elems_are_gc_refs {
+            size
+        } else {
+            builder.ins().iconst(ir::types::I32, 0)
+        };
 
         // Second, now that we have the array object's total size, call the
         // `gc_alloc_raw` builtin libcall to allocate the array.
         let array_ref = emit_gc_raw_alloc(
             func_env,
             builder,
             VMGcKind::ArrayRef,
+            num_gc_refs,
             interned_type_index,
             size,
             align,
@@ -385,10 +394,15 @@ impl GcCompiler for DrcCompiler {
         assert_eq!(VMGcKind::UNUSED_MASK & struct_size, struct_size);
         let struct_size_val = builder.ins().iconst(ir::types::I32, i64::from(struct_size));
 
+        let num_gc_refs = struct_layout.fields.iter().filter(|f| f.is_gc_ref).count();
+        let num_gc_refs = u32::try_from(num_gc_refs).unwrap();
+        let num_gc_refs = builder.ins().iconst(ir::types::I32, i64::from(num_gc_refs));
+
         let struct_ref = emit_gc_raw_alloc(
             func_env,
             builder,
             VMGcKind::StructRef,
+            num_gc_refs,
             interned_type_index,
             struct_size_val,
             struct_align,

crates/environ/src/gc.rs

@@ -81,23 +81,36 @@ fn common_array_layout(
     header_align: u32,
     expected_array_length_offset: u32,
 ) -> GcArrayLayout {
+    use core::mem;
+
     assert!(header_size >= crate::VM_GC_HEADER_SIZE);
     assert!(header_align >= crate::VM_GC_HEADER_ALIGN);
 
     let mut size = header_size;
     let mut align = header_align;
 
-    let length_field_offset = field(&mut size, &mut align, 4);
+    let length_field_size = u32::try_from(mem::size_of::<u32>()).unwrap();
+    let length_field_offset = field(&mut size, &mut align, length_field_size);
     assert_eq!(length_field_offset, expected_array_length_offset);
 
     let elem_size = byte_size_of_wasm_ty_in_gc_heap(&ty.0.element_type);
     let elems_offset = align_up(&mut size, &mut align, elem_size);
     assert_eq!(elems_offset, size);
 
+    let elems_are_gc_refs = ty.0.element_type.is_vmgcref_type_and_not_i31();
+    if elems_are_gc_refs {
+        debug_assert_eq!(
+            length_field_offset + length_field_size,
+            elems_offset,
+            "DRC collector relies on GC ref elements appearing directly after the length field, without any padding",
+        );
+    }
+
     GcArrayLayout {
         base_size: size,
         align,
         elem_size,
+        elems_are_gc_refs,
     }
 }
 
@@ -127,7 +140,9 @@ fn common_struct_layout(
         .iter()
         .map(|f| {
             let field_size = byte_size_of_wasm_ty_in_gc_heap(&f.element_type);
-            field(&mut size, &mut align, field_size)
+            let offset = field(&mut size, &mut align, field_size);
+            let is_gc_ref = f.element_type.is_vmgcref_type_and_not_i31();
+            GcStructLayoutField { offset, is_gc_ref }
         })
         .collect();
 
@@ -241,6 +256,9 @@ pub struct GcArrayLayout {
 
     /// The size and natural alignment of each element in this array.
     pub elem_size: u32,
+
+    /// Whether or not the elements of this array are GC references or not.
+    pub elems_are_gc_refs: bool,
 }
 
 impl GcArrayLayout {
@@ -283,9 +301,9 @@ pub struct GcStructLayout {
     /// The alignment (in bytes) of this struct.
     pub align: u32,
 
-    /// The fields of this struct. The `i`th entry is the `i`th struct field's
-    /// offset (in bytes) in the struct.
-    pub fields: Vec<u32>,
+    /// The fields of this struct. The `i`th entry contains information about
+    /// the `i`th struct field's layout.
+    pub fields: Vec<GcStructLayoutField>,
 }
 
 impl GcStructLayout {
@@ -297,6 +315,20 @@ impl GcStructLayout {
     }
 }
 
+/// A field in a `GcStructLayout`.
+#[derive(Clone, Copy, Debug)]
+pub struct GcStructLayoutField {
+    /// The offset (in bytes) of this field inside instances of this type.
+    pub offset: u32,
+
+    /// Whether or not this field might contain a reference to another GC
+    /// object.
+    ///
+    /// Note: it is okay for this to be `false` for `i31ref`s, since they never
+    /// actually reference another GC object.
+    pub is_gc_ref: bool,
+}
+
 /// The kind of an object in a GC heap.
 ///
 /// Note that this type is accessed from Wasm JIT code.

crates/environ/src/gc/drc.rs

@@ -1,9 +1,10 @@
 //! Layout of Wasm GC objects in the deferred reference-counting collector.
 
 use super::*;
+use core::cmp;
 
 /// The size of the `VMDrcHeader` header for GC objects.
-pub const HEADER_SIZE: u32 = 16;
+pub const HEADER_SIZE: u32 = 24;
 
 /// The align of the `VMDrcHeader` header for GC objects.
 pub const HEADER_ALIGN: u32 = 8;
@@ -25,6 +26,54 @@ impl GcTypeLayouts for DrcTypeLayouts {
     }
 
     fn struct_layout(&self, ty: &WasmStructType) -> GcStructLayout {
-        common_struct_layout(ty, HEADER_SIZE, HEADER_ALIGN)
+        // Sort the struct fields into the order in which we will lay them out.
+        //
+        // NB: we put all GC refs first so that we can simply store the number
+        // of GC refs in any object in the `VMDrcHeader` and then uniformly
+        // trace all structs types.
+        let mut fields: Vec<_> = ty.fields.iter().enumerate().collect();
+        fields.sort_by_key(|(i, f)| {
+            let is_gc_ref = f.element_type.is_vmgcref_type_and_not_i31();
+            let size = byte_size_of_wasm_ty_in_gc_heap(&f.element_type);
+            (cmp::Reverse(is_gc_ref), cmp::Reverse(size), *i)
+        });
+
+        // Compute the offset of each field as well as the size and alignment of
+        // the whole struct.
+        let mut size = HEADER_SIZE;
+        let mut align = HEADER_ALIGN;
+        let mut fields: Vec<_> = fields
+            .into_iter()
+            .map(|(i, f)| {
+                let field_size = byte_size_of_wasm_ty_in_gc_heap(&f.element_type);
+                let offset = field(&mut size, &mut align, field_size);
+                let is_gc_ref = f.element_type.is_vmgcref_type_and_not_i31();
+                (i, GcStructLayoutField { offset, is_gc_ref })
+            })
+            .collect();
+        if let Some((_i, f)) = fields.get(0) {
+            if f.is_gc_ref {
+                debug_assert_eq!(
+                    f.offset, HEADER_SIZE,
+                    "GC refs should come directly after the header, without any padding",
+                );
+            }
+        }
+
+        // Re-sort the fields into their definition (rather than layout) order
+        // and throw away the definition index.
+        fields.sort_by_key(|(i, _f)| *i);
+        let fields: Vec<_> = fields.into_iter().map(|(_i, f)| f).collect();
+
+        // Ensure that the final size is a multiple of the alignment, for
+        // simplicity.
+        let align_size_to = align;
+        align_up(&mut size, &mut align, align_size_to);
+
+        GcStructLayout {
+            size,
+            align,
+            fields,
+        }
     }
 }

crates/environ/src/types.rs

@@ -198,8 +198,8 @@ impl WasmValType {
     /// Is this a type that is represented as a `VMGcRef` and is additionally
     /// not an `i31`?
     ///
-    /// That is, is this a a type that actually refers to an object allocated in
-    /// a GC heap?
+    /// That is, is this a type that actually refers to an object allocated in a
+    /// GC heap?
     #[inline]
     pub fn is_vmgcref_type_and_not_i31(&self) -> bool {
         match self {
@@ -280,8 +280,8 @@ impl WasmRefType {
     /// Is this a type that is represented as a `VMGcRef` and is additionally
     /// not an `i31`?
     ///
-    /// That is, is this a a type that actually refers to an object allocated in
-    /// a GC heap?
+    /// That is, is this a type that actually refers to an object allocated in a
+    /// GC heap?
     #[inline]
     pub fn is_vmgcref_type_and_not_i31(&self) -> bool {
         self.heap_type.is_vmgcref_type_and_not_i31()
@@ -545,8 +545,8 @@ impl WasmHeapType {
     /// Is this a type that is represented as a `VMGcRef` and is additionally
     /// not an `i31`?
     ///
-    /// That is, is this a a type that actually refers to an object allocated in
-    /// a GC heap?
+    /// That is, is this a type that actually refers to an object allocated in a
+    /// GC heap?
     #[inline]
     pub fn is_vmgcref_type_and_not_i31(&self) -> bool {
         self.is_vmgcref_type() && *self != Self::I31
@@ -862,6 +862,20 @@ impl TypeTrace for WasmStorageType {
     }
 }
 
+impl WasmStorageType {
+    /// Is this a type that is represented as a `VMGcRef` and is additionally
+    /// not an `i31`?
+    ///
+    /// That is, is this a type that actually refers to an object allocated in a
+    /// GC heap?
+    pub fn is_vmgcref_type_and_not_i31(&self) -> bool {
+        match self {
+            WasmStorageType::I8 | WasmStorageType::I16 => false,
+            WasmStorageType::Val(v) => v.is_vmgcref_type_and_not_i31(),
+        }
+    }
+}
+
 /// The type of a struct field or array element.
 #[derive(Debug, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)]
 pub struct WasmFieldType {