Split GC type layout computation into a shared trait in `wasmtime_env…

…iron`

This stuff was previously living in the `GcRuntime` trait, but it turns out that the
compiler will also need to know the layout of GC types, who'd have thunk it.

crates/cranelift/src/gc.rs

@@ -19,6 +19,7 @@ use enabled as imp;
 mod disabled;
 #[cfg(not(feature = "gc"))]
 use disabled as imp;
+use wasmtime_environ::GcTypeLayouts;
 
 /// Get the GC compiler configured for the given function environment.
 pub fn gc_compiler(func_env: &FuncEnvironment<'_>) -> Box<dyn GcCompiler> {
@@ -82,6 +83,10 @@ pub fn gc_ref_table_fill_builtin(
 
 /// A trait for different collectors to emit any GC barriers they might require.
 pub trait GcCompiler {
+    /// Get the GC type layouts for this GC compiler.
+    #[allow(dead_code)] // Used in future PRs.
+    fn layouts(&self) -> &dyn GcTypeLayouts;
+
     /// Emit a read barrier for when we are cloning a GC reference onto the Wasm
     /// stack.
     ///

crates/cranelift/src/gc/enabled.rs

@@ -3,11 +3,13 @@ use crate::func_environ::FuncEnvironment;
 use cranelift_codegen::ir::{self, condcodes::IntCC, InstBuilder};
 use cranelift_frontend::FunctionBuilder;
 use cranelift_wasm::{TargetEnvironment, WasmHeapTopType, WasmHeapType, WasmRefType, WasmResult};
-use wasmtime_environ::{PtrSize, I31_DISCRIMINANT, NON_NULL_NON_I31_MASK};
+use wasmtime_environ::{
+    drc::DrcTypeLayouts, GcTypeLayouts, PtrSize, I31_DISCRIMINANT, NON_NULL_NON_I31_MASK,
+};
 
 /// Get the default GC compiler.
 pub fn gc_compiler(_func_env: &FuncEnvironment<'_>) -> Box<dyn GcCompiler> {
-    Box::new(DrcCompiler)
+    Box::new(DrcCompiler::default())
 }
 
 pub fn unbarriered_load_gc_ref(
@@ -184,7 +186,10 @@ impl FuncEnvironment<'_> {
     }
 }
 
-struct DrcCompiler;
+#[derive(Default)]
+struct DrcCompiler {
+    layouts: DrcTypeLayouts,
+}
 
 impl DrcCompiler {
     /// Generate code to load the given GC reference's ref count.
@@ -276,6 +281,10 @@ impl DrcCompiler {
 }
 
 impl GcCompiler for DrcCompiler {
+    fn layouts(&self) -> &dyn GcTypeLayouts {
+        &self.layouts
+    }
+
     fn translate_read_gc_reference(
         &mut self,
         func_env: &mut FuncEnvironment<'_>,

crates/environ/src/gc.rs

@@ -9,13 +9,138 @@
 //! on our various `gc` cargo features is the actual garbage collection
 //! functions and their associated impact on binary size anyways.
 
+#[cfg(feature = "gc")]
+pub mod drc;
+
+use crate::prelude::*;
+use wasmtime_types::{WasmArrayType, WasmStructType};
+
 /// Discriminant to check whether GC reference is an `i31ref` or not.
 pub const I31_DISCRIMINANT: u64 = 1;
 
 /// A mask that can be used to check for non-null and non-i31ref GC references
 /// with a single bitwise-and operation.
 pub const NON_NULL_NON_I31_MASK: u64 = !I31_DISCRIMINANT;
 
+/// TODO FITZGEN
+pub trait GcTypeLayouts {
+    /// Get this collector's layout for the given array type.
+    fn array_layout(&self, ty: &WasmArrayType) -> GcArrayLayout;
+
+    /// Get this collector's layout for the given struct type.
+    fn struct_layout(&self, ty: &WasmStructType) -> GcStructLayout;
+}
+
+/// The layout of a GC-managed object.
+#[derive(Clone, Debug)]
+pub enum GcLayout {
+    /// The layout of a GC-managed array object.
+    Array(GcArrayLayout),
+
+    /// The layout of a GC-managed struct object.
+    Struct(GcStructLayout),
+}
+
+impl From<GcArrayLayout> for GcLayout {
+    fn from(layout: GcArrayLayout) -> Self {
+        Self::Array(layout)
+    }
+}
+
+impl From<GcStructLayout> for GcLayout {
+    fn from(layout: GcStructLayout) -> Self {
+        Self::Struct(layout)
+    }
+}
+
+impl GcLayout {
+    /// Get the underlying `GcStructLayout`, or panic.
+    #[track_caller]
+    pub fn unwrap_struct(&self) -> &GcStructLayout {
+        match self {
+            Self::Struct(s) => s,
+            _ => panic!("GcLayout::unwrap_struct on non-struct GC layout"),
+        }
+    }
+
+    /// Get the underlying `GcArrayLayout`, or panic.
+    #[track_caller]
+    pub fn unwrap_array(&self) -> &GcArrayLayout {
+        match self {
+            Self::Array(a) => a,
+            _ => panic!("GcLayout::unwrap_array on non-array GC layout"),
+        }
+    }
+}
+
+/// The layout of a GC-managed array.
+///
+/// This layout is only valid for use with the GC runtime that created it. It is
+/// not valid to use one GC runtime's layout with another GC runtime, doing so
+/// is memory safe but will lead to general incorrectness like panics and wrong
+/// results.
+///
+/// All offsets are from the start of the object; that is, the size of the GC
+/// header (for example) is included in the offset.
+///
+/// All arrays are composed of the generic `VMGcHeader`, followed by
+/// collector-specific fields, followed by the contiguous array elements
+/// themselves. The array elements must be aligned to the element type's natural
+/// alignment.
+#[derive(Clone, Debug)]
+#[allow(dead_code)] // Not used yet, but added for completeness.
+pub struct GcArrayLayout {
+    /// The size of this array object, ignoring its elements.
+    pub size: u32,
+
+    /// The alignment of this array.
+    pub align: u32,
+
+    /// The offset of the array's length.
+    pub length_field_offset: u32,
+
+    /// The offset from where this array's contiguous elements begin.
+    pub elems_offset: u32,
+
+    /// The size and natural alignment of each element in this array.
+    pub elem_size: u32,
+}
+
+impl GcArrayLayout {
+    /// Get the total size of this array for a given length of elements.
+    pub fn size_for_len(&self, len: u32) -> u32 {
+        self.size + len * self.elem_size
+    }
+
+    /// Get the offset of the `i`th element in an array with this layout.
+    #[inline]
+    pub fn elem_offset(&self, i: u32, elem_size: u32) -> u32 {
+        self.elems_offset + i * elem_size
+    }
+}
+
+/// The layout for a GC-managed struct type.
+///
+/// This layout is only valid for use with the GC runtime that created it. It is
+/// not valid to use one GC runtime's layout with another GC runtime, doing so
+/// is memory safe but will lead to general incorrectness like panics and wrong
+/// results.
+///
+/// All offsets are from the start of the object; that is, the size of the GC
+/// header (for example) is included in the offset.
+#[derive(Clone, Debug)]
+pub struct GcStructLayout {
+    /// The size of this struct.
+    pub size: u32,
+
+    /// The alignment of this struct.
+    pub align: u32,
+
+    /// The fields of this struct. The `i`th entry is the `i`th struct field's
+    /// offset in the struct.
+    pub fields: Vec<u32>,
+}
+
 /// 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

@@ -0,0 +1,100 @@
+//! Layout of Wasm GC objects in the deferred reference-counting collector.
+
+use super::*;
+use wasmtime_types::{WasmStorageType, WasmValType};
+
+/// The size of the `VMDrcHeader` header for GC objects.
+pub const HEADER_SIZE: u32 = 16;
+
+/// The align of the `VMDrcHeader` header for GC objects.
+pub const HEADER_ALIGN: u32 = 8;
+
+/// The offset of the length field in a `VMDrcArrayHeader`.
+pub const ARRAY_LENGTH_OFFSET: u32 = HEADER_SIZE;
+
+/// Align `offset` up to `bytes`, updating `max_align` if `align` is the
+/// new maximum alignment, and returning the aligned offset.
+fn align_up(offset: &mut u32, max_align: &mut u32, align: u32) -> u32 {
+    debug_assert!(max_align.is_power_of_two());
+    debug_assert!(align.is_power_of_two());
+    *offset = offset.checked_add(align - 1).unwrap() & !(align - 1);
+    *max_align = core::cmp::max(*max_align, align);
+    *offset
+}
+
+/// Define a new field of size and alignment `bytes`, updating the object's
+/// total `size` and `align` as necessary. The offset of the new field is
+/// returned.
+fn field(size: &mut u32, align: &mut u32, bytes: u32) -> u32 {
+    let offset = align_up(size, align, bytes);
+    *size += bytes;
+    offset
+}
+
+fn size_of_wasm_ty(ty: &WasmStorageType) -> u32 {
+    match ty {
+        WasmStorageType::I8 => 1,
+        WasmStorageType::I16 => 2,
+        WasmStorageType::Val(ty) => match ty {
+            WasmValType::I32 | WasmValType::F32 | WasmValType::Ref(_) => 4,
+            WasmValType::I64 | WasmValType::F64 => 8,
+            WasmValType::V128 => 16,
+        },
+    }
+}
+
+/// The layout of Wasm GC objects in the deferred reference-counting collector.
+#[derive(Default)]
+pub struct DrcTypeLayouts;
+
+impl GcTypeLayouts for DrcTypeLayouts {
+    fn array_layout(&self, ty: &WasmArrayType) -> GcArrayLayout {
+        let mut size = HEADER_SIZE;
+        let mut align = HEADER_ALIGN;
+
+        let length_field_offset = field(&mut size, &mut align, 4);
+        debug_assert_eq!(length_field_offset, ARRAY_LENGTH_OFFSET);
+
+        let elem_size = size_of_wasm_ty(&ty.0.element_type);
+        let elems_offset = align_up(&mut size, &mut align, elem_size);
+
+        GcArrayLayout {
+            size,
+            align,
+            length_field_offset,
+            elems_offset,
+            elem_size,
+        }
+    }
+
+    fn struct_layout(&self, ty: &WasmStructType) -> GcStructLayout {
+        // Process each field, aligning it to its natural alignment.
+        //
+        // We don't try and do any fancy field reordering to minimize padding
+        // (yet?) because (a) the toolchain probably already did that and (b)
+        // we're just doing the simple thing first. We can come back and improve
+        // things here if we find that (a) isn't actually holding true in
+        // practice.
+        let mut size = HEADER_SIZE;
+        let mut align = HEADER_ALIGN;
+
+        let fields = ty
+            .fields
+            .iter()
+            .map(|f| {
+                let field_size = size_of_wasm_ty(&f.element_type);
+                field(&mut size, &mut align, field_size)
+            })
+            .collect();
+
+        // Ensure that the final size is a multiple of the alignment, for
+        // simplicity.
+        align_up(&mut size, &mut 16, align);
+
+        GcStructLayout {
+            size,
+            align,
+            fields,
+        }
+    }
+}

crates/wasmtime/src/runtime/gc/enabled/arrayref.rs

@@ -2,7 +2,7 @@
 
 use crate::runtime::vm::VMGcRef;
 use crate::store::StoreId;
-use crate::vm::{GcArrayLayout, GcLayout, VMArrayRef, VMGcHeader};
+use crate::vm::{VMArrayRef, VMGcHeader};
 use crate::{
     prelude::*,
     store::{AutoAssertNoGc, StoreContextMut, StoreOpaque},
@@ -11,7 +11,7 @@ use crate::{
 };
 use crate::{AnyRef, FieldType};
 use core::mem::{self, MaybeUninit};
-use wasmtime_environ::{VMGcKind, VMSharedTypeIndex};
+use wasmtime_environ::{GcArrayLayout, GcLayout, VMGcKind, VMSharedTypeIndex};
 
 /// An allocator for a particular Wasm GC array type.
 ///

crates/wasmtime/src/runtime/gc/enabled/structref.rs

@@ -2,7 +2,7 @@
 
 use crate::runtime::vm::VMGcRef;
 use crate::store::StoreId;
-use crate::vm::{GcLayout, GcStructLayout, VMGcHeader, VMStructRef};
+use crate::vm::{VMGcHeader, VMStructRef};
 use crate::{
     prelude::*,
     store::{AutoAssertNoGc, StoreContextMut, StoreOpaque},
@@ -11,7 +11,7 @@ use crate::{
 };
 use crate::{AnyRef, FieldType};
 use core::mem::{self, MaybeUninit};
-use wasmtime_environ::{VMGcKind, VMSharedTypeIndex};
+use wasmtime_environ::{GcLayout, GcStructLayout, VMGcKind, VMSharedTypeIndex};
 
 /// An allocator for a particular Wasm GC struct type.
 ///

crates/wasmtime/src/runtime/type_registry.rs

@@ -5,7 +5,7 @@
 
 use crate::prelude::*;
 use crate::sync::RwLock;
-use crate::vm::{GcLayout, GcRuntime};
+use crate::vm::GcRuntime;
 use crate::Engine;
 use alloc::borrow::Cow;
 use alloc::sync::Arc;
@@ -22,7 +22,7 @@ use core::{
 };
 use hashbrown::HashSet;
 use wasmtime_environ::{
-    iter_entity_range, packed_option::PackedOption, EngineOrModuleTypeIndex,
+    iter_entity_range, packed_option::PackedOption, EngineOrModuleTypeIndex, GcLayout,
     ModuleInternedTypeIndex, ModuleTypes, PrimaryMap, SecondaryMap, TypeTrace, VMSharedTypeIndex,
     WasmRecGroup, WasmSubType,
 };
@@ -749,10 +749,10 @@ impl TypeRegistryInner {
         let gc_layout = match &ty.composite_type {
             wasmtime_environ::WasmCompositeType::Func(_) => None,
             wasmtime_environ::WasmCompositeType::Array(a) => {
-                Some(gc_runtime.array_layout(a).into())
+                Some(gc_runtime.layouts().array_layout(a).into())
             }
             wasmtime_environ::WasmCompositeType::Struct(s) => {
-                Some(gc_runtime.struct_layout(s).into())
+                Some(gc_runtime.layouts().struct_layout(s).into())
             }
         };
 

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

@@ -22,7 +22,7 @@ use crate::prelude::*;
 use crate::runtime::vm::GcHeapAllocationIndex;
 use core::ptr;
 use core::{any::Any, num::NonZeroUsize};
-use wasmtime_environ::{StackMap, VMGcKind, VMSharedTypeIndex};
+use wasmtime_environ::{GcArrayLayout, GcStructLayout, StackMap, VMGcKind, VMSharedTypeIndex};
 
 /// Used by the runtime to lookup information about a module given a
 /// program counter value.

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

@@ -17,7 +17,7 @@ use crate::runtime::vm::GcRuntime;
 
 /// Get the default GC runtime.
 pub fn default_gc_runtime() -> impl GcRuntime {
-    DrcCollector
+    DrcCollector::default()
 }
 
 /// The default GC heap capacity: 512KiB.

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

@@ -3,11 +3,10 @@ use crate::{
     prelude::*,
     runtime::vm::{GcHeap, GcStore, VMGcRef},
     store::{AutoAssertNoGc, StoreOpaque},
-    vm::GcArrayLayout,
     AnyRef, ExternRef, HeapType, RootedGcRefImpl, StorageType, Val, ValType,
 };
 use core::fmt;
-use wasmtime_environ::VMGcKind;
+use wasmtime_environ::{GcArrayLayout, VMGcKind};
 
 /// A `VMGcRef` that we know points to a `array`.
 ///