WIP drc transitive dec ref

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

@@ -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 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 {