Copying collector: Store most tracing info inline in the object header

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

@@ -124,13 +124,14 @@ fn emit_gc_raw_alloc(
     ty: ModuleInternedTypeIndex,
     size: ir::Value,
     align: u32,
+    reserved_bits: u32,
 ) -> ir::Value {
     let gc_alloc_raw_builtin = func_env.builtin_functions.gc_alloc_raw(builder.func);
     let vmctx = func_env.vmctx_val(&mut builder.cursor());
 
     let kind = builder
         .ins()
-        .iconst(ir::types::I32, i64::from(kind.as_u32()));
+        .iconst(ir::types::I32, i64::from(kind.as_u32() | reserved_bits));
 
     let ty = func_env.module_interned_to_shared_ty(&mut builder.cursor(), ty);
 

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

@@ -13,7 +13,7 @@ use crate::translate::TargetEnvironment;
 use cranelift_codegen::ir::{self, InstBuilder};
 use cranelift_frontend::FunctionBuilder;
 use wasmtime_environ::copying::{
-    ALIGN, EXCEPTION_TAG_DEFINED_OFFSET, EXCEPTION_TAG_INSTANCE_OFFSET,
+    ALIGN, EXCEPTION_TAG_DEFINED_OFFSET, EXCEPTION_TAG_INSTANCE_OFFSET, InlineTraceInfo,
 };
 use wasmtime_environ::{
     GcTypeLayouts, ModuleInternedTypeIndex, PtrSize, TypeIndex, VMGcKind, WasmHeapTopType,
@@ -92,6 +92,7 @@ impl CopyingCompiler {
         kind: VMGcKind,
         ty: ModuleInternedTypeIndex,
         size: ir::Value,
+        reserved_bits: u32,
     ) -> WasmResult<(ir::Value, ir::Value)> {
         debug_assert_ne!(kind, VMGcKind::ExternRef);
         debug_assert!(!ty.is_reserved_value());
@@ -132,7 +133,7 @@ impl CopyingCompiler {
             builder.switch_to_block(slow_block);
             builder.seal_block(slow_block);
             builder.set_cold_block(slow_block);
-            let gc_ref = emit_gc_raw_alloc(func_env, builder, kind, ty, size, ALIGN);
+            let gc_ref = emit_gc_raw_alloc(func_env, builder, kind, ty, size, ALIGN, reserved_bits);
             let base = func_env.get_gc_heap_base(builder)?;
             let heap_offset = uextend_i32_to_pointer_type(builder, pointer_type, gc_ref);
             let obj_ptr = builder.ins().iadd(base, heap_offset);
@@ -166,10 +167,10 @@ impl CopyingCompiler {
             let heap_offset = uextend_i32_to_pointer_type(builder, pointer_type, gc_ref);
             let obj_ptr = builder.ins().iadd(base, heap_offset);
 
-            // Write `VMGcHeader::kind`.
+            // Write `VMGcHeader::kind` with inline trace info bits included.
             let kind_val = builder
                 .ins()
-                .iconst(ir::types::I32, i64::from(kind.as_u32()));
+                .iconst(ir::types::I32, i64::from(kind.as_u32() | reserved_bits));
             builder.ins().store(
                 ir::MemFlagsData::trusted(),
                 kind_val,
@@ -232,6 +233,7 @@ impl GcCompiler for CopyingCompiler {
 
         let len_offset = gc_compiler(func_env)?.layouts().array_length_field_offset();
         let array_layout = func_env.array_layout(interned_type_index)?.clone();
+        let reserved_bits = InlineTraceInfo::array(&array_layout).encode();
 
         // First, compute the array's total size.
         let size = emit_array_size(func_env, builder, &array_layout, len);
@@ -243,6 +245,7 @@ impl GcCompiler for CopyingCompiler {
             VMGcKind::ArrayRef,
             interned_type_index,
             size,
+            reserved_bits,
         )?;
         let len_addr = builder.ins().iadd_imm(object_addr, i64::from(len_offset));
         builder.ins().store(GC_MEMFLAGS, len, len_addr, 0);
@@ -262,6 +265,7 @@ impl GcCompiler for CopyingCompiler {
         let struct_layout = func_env.struct_or_exn_layout(interned_type_index);
 
         let struct_size = struct_layout.size;
+        let reserved_bits = InlineTraceInfo::r#struct(&struct_layout).encode();
 
         let struct_size_val = builder.ins().iconst(ir::types::I32, i64::from(struct_size));
 
@@ -271,6 +275,7 @@ impl GcCompiler for CopyingCompiler {
             VMGcKind::StructRef,
             interned_type_index,
             struct_size_val,
+            reserved_bits,
         )?;
 
         // Initialize fields.
@@ -300,6 +305,7 @@ impl GcCompiler for CopyingCompiler {
         let exn_layout = func_env.struct_or_exn_layout(interned_type_index);
 
         let exn_size = exn_layout.size;
+        let reserved_bits = InlineTraceInfo::r#struct(&exn_layout).encode();
 
         let exn_size_val = builder.ins().iconst(ir::types::I32, i64::from(exn_size));
 
@@ -309,6 +315,7 @@ impl GcCompiler for CopyingCompiler {
             VMGcKind::ExnRef,
             interned_type_index,
             exn_size_val,
+            reserved_bits,
         )?;
 
         // Initialize fields.

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

@@ -471,6 +471,7 @@ impl GcCompiler for DrcCompiler {
             interned_type_index,
             size,
             align,
+            0,
         );
 
         // Write the array's length into the appropriate slot.
@@ -512,6 +513,7 @@ impl GcCompiler for DrcCompiler {
             interned_type_index,
             struct_size_val,
             struct_align,
+            0,
         );
 
         // Second, initialize each of the newly-allocated struct's fields.
@@ -562,6 +564,7 @@ impl GcCompiler for DrcCompiler {
             interned_type_index,
             exn_size_val,
             exn_align,
+            0,
         );
 
         // Second, initialize each of the newly-allocated exception

crates/environ/src/gc/copying.rs

@@ -42,6 +42,122 @@ pub const MIN_OBJECT_SIZE: u32 = FORWARDING_REF_OFFSET + mem::size_of::<u32>() a
 #[derive(Default)]
 pub struct CopyingTypeLayouts;
 
+/// The inline trace info encoded in a `VMCopyingHeader`'s reserved bits.
+#[derive(Clone, Copy, Debug)]
+pub enum InlineTraceInfo {
+    /// The trace info is not encoded inline. Look it up from `TraceInfos`.
+    OutOfLine,
+
+    /// Inline trace info for an array type.
+    Array {
+        /// Whether the array's elements are GC references that need tracing.
+        elems_are_gc_refs: bool,
+    },
+
+    /// Inline trace info for a struct, exception, or externref type.
+    Struct {
+        /// A bitmap where the `i`th bit is set iff the `i`th `u32` in the
+        /// object's data (after the header) is a `VMGcRef` that needs
+        /// tracing. Only the lower 23 bits are meaningful.
+        gc_ref_bitmap: u32,
+    },
+}
+
+impl InlineTraceInfo {
+    const IS_INLINE_BIT: u32 = 1 << 1;
+    const IS_ARRAY_BIT: u32 = 1 << 2;
+    const ELEMS_ARE_GC_REFS_BIT: u32 = 1 << 3;
+    const STRUCT_BITMAP_SHIFT: u32 = 3;
+    const STRUCT_BITMAP_BITS: u32 = 23;
+
+    /// Inline trace info for an object with no GC-reference edges (e.g.
+    /// externrefs). This is a `Struct` with an empty bitmap.
+    pub const NO_GC_EDGES: InlineTraceInfo = InlineTraceInfo::Struct { gc_ref_bitmap: 0 };
+
+    /// Create inline trace info for the given GC layout.
+    pub fn new(layout: &GcLayout) -> Self {
+        match layout {
+            GcLayout::Array(a) => Self::array(a),
+            GcLayout::Struct(s) => Self::r#struct(s),
+        }
+    }
+
+    /// Create inline trace info for an array type.
+    ///
+    /// Arrays can always be represented inline (only one bit of payload is
+    /// needed), so this never returns `OutOfLine`.
+    pub fn array(layout: &GcArrayLayout) -> Self {
+        InlineTraceInfo::Array {
+            elems_are_gc_refs: layout.elems_are_gc_refs,
+        }
+    }
+
+    /// Create inline trace info for a struct or exception type.
+    ///
+    /// Returns `OutOfLine` only when there is a GC-reference field whose
+    /// offset cannot be represented in the 23-bit bitmap. Non-GC-reference
+    /// fields are ignored regardless of offset.
+    pub fn r#struct(layout: &GcStructLayout) -> Self {
+        let mut bitmap: u32 = 0;
+        for f in layout.fields.iter() {
+            if !f.is_gc_ref {
+                continue;
+            }
+            let Some(data_offset) = f.offset.checked_sub(HEADER_SIZE) else {
+                return Self::OutOfLine;
+            };
+            if data_offset % 4 != 0 {
+                return Self::OutOfLine;
+            }
+            let slot = data_offset / 4;
+            if slot >= Self::STRUCT_BITMAP_BITS {
+                return Self::OutOfLine;
+            }
+            bitmap |= 1u32 << slot;
+        }
+        InlineTraceInfo::Struct {
+            gc_ref_bitmap: bitmap,
+        }
+    }
+
+    /// Encode this inline trace info as its bit-packed representation for
+    /// storage in a `VMGcHeader`'s reserved bits.
+    pub fn encode(&self) -> u32 {
+        match self {
+            Self::OutOfLine => 0,
+            Self::Array { elems_are_gc_refs } => {
+                Self::IS_INLINE_BIT
+                    | Self::IS_ARRAY_BIT
+                    | if *elems_are_gc_refs {
+                        Self::ELEMS_ARE_GC_REFS_BIT
+                    } else {
+                        0
+                    }
+            }
+            Self::Struct { gc_ref_bitmap } => {
+                Self::IS_INLINE_BIT | (*gc_ref_bitmap << Self::STRUCT_BITMAP_SHIFT)
+            }
+        }
+    }
+
+    /// Decode inline trace info from the reserved bits of a `VMGcHeader`.
+    pub fn decode(reserved: u32) -> Self {
+        if reserved & Self::IS_INLINE_BIT == 0 {
+            return Self::OutOfLine;
+        }
+        if reserved & Self::IS_ARRAY_BIT != 0 {
+            Self::Array {
+                elems_are_gc_refs: reserved & Self::ELEMS_ARE_GC_REFS_BIT != 0,
+            }
+        } else {
+            Self::Struct {
+                gc_ref_bitmap: (reserved >> Self::STRUCT_BITMAP_SHIFT)
+                    & ((1u32 << Self::STRUCT_BITMAP_BITS) - 1),
+            }
+        }
+    }
+}
+
 impl GcTypeLayouts for CopyingTypeLayouts {
     fn array_length_field_offset(&self) -> u32 {
         ARRAY_LENGTH_OFFSET

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

@@ -22,7 +22,7 @@ use core::{
     alloc::Layout, any::Any, mem, num::NonZeroU32, ptr::NonNull, sync::atomic::AtomicUsize,
 };
 use wasmtime_environ::copying::{
-    ALIGN, ARRAY_LENGTH_OFFSET, CopyingTypeLayouts, HEADER_COPIED_BIT,
+    ALIGN, ARRAY_LENGTH_OFFSET, CopyingTypeLayouts, HEADER_COPIED_BIT, HEADER_SIZE, InlineTraceInfo,
 };
 use wasmtime_environ::{
     GcArrayLayout, GcStructLayout, GcTypeLayouts, POISON, VMGcKind, VMSharedTypeIndex, gc_assert,
@@ -83,6 +83,12 @@ impl VMCopyingHeader {
         let reserved = self.header.reserved_u26();
         self.header.set_reserved_u26(reserved | HEADER_COPIED_BIT);
     }
+
+    /// Decode the inline trace info from this header's reserved bits.
+    #[inline]
+    fn inline_trace_info(&self) -> InlineTraceInfo {
+        InlineTraceInfo::decode(self.header.reserved_u26())
+    }
 }
 
 /// A copying collector header together with a forwarding reference.
@@ -586,22 +592,20 @@ survived collection, since the active space is the same size as the idle space",
         let index = gc_ref.heap_index()?.get();
         debug_assert!(self.is_in_active_space(index));
 
-        let ty = self.index(copying_ref(gc_ref))?.header.ty();
-
-        // `externref`s have no GC edges.
-        let Some(ty) = ty else {
-            return Ok(());
-        };
-
+        let inline_info = self.index(copying_ref(gc_ref))?.inline_trace_info();
         let object_start = usize::try_from(index)?;
-        match trace_infos.trace_info(&ty) {
-            TraceInfo::Struct { gc_ref_offsets } => {
-                for &offset in gc_ref_offsets {
-                    self.scan_field(object_start, offset)?;
+
+        match inline_info {
+            InlineTraceInfo::Struct { gc_ref_bitmap } => {
+                let mut remaining = gc_ref_bitmap;
+                while remaining != 0 {
+                    let slot = remaining.trailing_zeros();
+                    self.scan_field(object_start, HEADER_SIZE + slot * 4)?;
+                    remaining &= remaining - 1;
                 }
             }
-            TraceInfo::Array { gc_ref_elems } => {
-                if *gc_ref_elems {
+            InlineTraceInfo::Array { elems_are_gc_refs } => {
+                if elems_are_gc_refs {
                     let array_ref = gc_ref.as_arrayref_unchecked();
                     let len = self.array_len(array_ref)?;
 
@@ -617,6 +621,22 @@ survived collection, since the active space is the same size as the idle space",
                     }
                 }
             }
+            InlineTraceInfo::OutOfLine => {
+                let ty = self.index(copying_ref(gc_ref))?.header.ty();
+                let Some(ty) = ty else {
+                    bail_bug!("out-of-line trace info but no type index");
+                };
+                match trace_infos.trace_info(&ty) {
+                    TraceInfo::Struct { gc_ref_offsets } => {
+                        for &offset in gc_ref_offsets {
+                            self.scan_field(object_start, offset)?;
+                        }
+                    }
+                    TraceInfo::Array { .. } => {
+                        bail_bug!("arrays should always have inline trace info");
+                    }
+                }
+            }
         }
         Ok(())
     }
@@ -756,10 +776,9 @@ unsafe impl GcHeap for CopyingHeap {
         let align = usize::try_from(ALIGN).unwrap();
         let size = core::mem::size_of::<VMCopyingExternRef>();
         let size = (size + align - 1) & !(align - 1);
-        let gc_ref = match self.alloc_raw(
-            VMGcHeader::externref(),
-            Layout::from_size_align(size, align)?,
-        )? {
+        let mut header = VMGcHeader::externref();
+        header.set_reserved_u26(InlineTraceInfo::NO_GC_EDGES.encode());
+        let gc_ref = match self.alloc_raw(header, Layout::from_size_align(size, align)?)? {
             Err(n) => return Ok(Err(n)),
             Ok(gc_ref) => gc_ref,
         };
@@ -822,7 +841,7 @@ unsafe impl GcHeap for CopyingHeap {
             0,
             "bump_ptr is not aligned to ALIGN"
         );
-        debug_assert_eq!(header.reserved_u26(), 0);
+        debug_assert_eq!(header.reserved_u26() & HEADER_COPIED_BIT, 0);
 
         // We must have trace info for every GC type that we allocate in this
         // heap. Trace info is eagerly registered during module instantiation
@@ -885,11 +904,12 @@ unsafe impl GcHeap for CopyingHeap {
         } else {
             VMGcKind::StructRef
         };
-        let gc_ref =
-            match self.alloc_raw(VMGcHeader::from_kind_and_index(kind, ty), layout.layout())? {
-                Err(n) => return Ok(Err(n)),
-                Ok(gc_ref) => gc_ref,
-            };
+        let mut header = VMGcHeader::from_kind_and_index(kind, ty);
+        header.set_reserved_u26(InlineTraceInfo::r#struct(layout).encode());
+        let gc_ref = match self.alloc_raw(header, layout.layout())? {
+            Err(n) => return Ok(Err(n)),
+            Ok(gc_ref) => gc_ref,
+        };
 
         Ok(Ok(gc_ref))
     }
@@ -906,14 +926,13 @@ unsafe impl GcHeap for CopyingHeap {
         length: u32,
         layout: &GcArrayLayout,
     ) -> Result<Result<VMArrayRef, u64>> {
+        let mut header = VMGcHeader::from_kind_and_index(VMGcKind::ArrayRef, ty);
+        header.set_reserved_u26(InlineTraceInfo::array(layout).encode());
         let layout = layout
             .layout(length)
             .ok_or_else(|| crate::Trap::AllocationTooLarge)?;
 
-        let gc_ref = match self.alloc_raw(
-            VMGcHeader::from_kind_and_index(VMGcKind::ArrayRef, ty),
-            layout,
-        )? {
+        let gc_ref = match self.alloc_raw(header, layout)? {
             Err(n) => return Ok(Err(n)),
             Ok(gc_ref) => gc_ref,
         };

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

@@ -16,6 +16,10 @@ pub(super) enum TraceInfo {
     Array {
         /// Whether this array type's elements are GC references, and need
         /// tracing.
+        #[cfg_attr(
+            not(feature = "gc-drc"),
+            allow(dead_code, reason = "easier not to cfg on/off")
+        )]
         gc_ref_elems: bool,
     },