Refactor exception objects to share more implementation with structs. (#11467)

* Refactor exception objects to share more implementation with structs.

This PR updates the implementation of Wasm exception objects to share
more layout-computation code with structs, except for a single fixed
tag slot. The former is simply a nice refactor, and the latter is in
preparation for full exception support, which will require accessing
an exception's tag without knowing its layout dynamically at runtime.

* Review feedback.

Author: cfallin

crates/environ/src/compile/module_environ.rs

@@ -354,7 +354,13 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                         TypeRef::Tag(ty) => {
                             let index = TypeIndex::from_u32(ty.func_type_idx);
                             let signature = self.result.module.types[index];
-                            let tag = Tag { signature };
+                            let exception = self.types.define_exception_type_for_tag(
+                                signature.unwrap_module_type_index(),
+                            );
+                            let tag = Tag {
+                                signature,
+                                exception: EngineOrModuleTypeIndex::Module(exception),
+                            };
                             self.result.module.num_imported_tags += 1;
                             EntityType::Tag(tag)
                         }
@@ -426,7 +432,10 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                     let sigindex = entry?.func_type_idx;
                     let ty = TypeIndex::from_u32(sigindex);
                     let interned_index = self.result.module.types[ty];
-                    self.result.module.push_tag(interned_index);
+                    let exception = self
+                        .types
+                        .define_exception_type_for_tag(interned_index.unwrap_module_type_index());
+                    self.result.module.push_tag(interned_index, exception);
                 }
             }
 

crates/environ/src/compile/module_types.rs

@@ -1,9 +1,14 @@
 use crate::{
     EngineOrModuleTypeIndex, EntityRef, ModuleInternedRecGroupIndex, ModuleInternedTypeIndex,
     ModuleTypes, TypeConvert, TypeIndex, WasmArrayType, WasmCompositeInnerType, WasmCompositeType,
-    WasmFuncType, WasmHeapType, WasmResult, WasmStructType, WasmSubType, wasm_unsupported,
+    WasmExnType, WasmFieldType, WasmFuncType, WasmHeapType, WasmResult, WasmStorageType,
+    WasmStructType, WasmSubType, wasm_unsupported,
+};
+use std::{
+    borrow::Cow,
+    collections::{HashMap, hash_map::Entry},
+    ops::Index,
 };
-use std::{borrow::Cow, collections::HashMap, ops::Index};
 use wasmparser::{UnpackedIndex, Validator, ValidatorId};
 
 /// A type marking the start of a recursion group's definition.
@@ -37,6 +42,12 @@ pub struct ModuleTypesBuilder {
     /// trampoline type for a given function type if we've already interned one.
     trampoline_types: HashMap<WasmFuncType, ModuleInternedTypeIndex>,
 
+    /// An interning map for exception types corresponding to function
+    /// types used by tags. Tags are nominal, but the underlying
+    /// Wasmtime types describe only the object layout and so are
+    /// structural.
+    exception_types: HashMap<ModuleInternedTypeIndex, ModuleInternedTypeIndex>,
+
     /// A map from already-interned `wasmparser` types to their corresponding
     /// Wasmtime type.
     wasmparser_to_wasmtime: HashMap<wasmparser::types::CoreTypeId, ModuleInternedTypeIndex>,
@@ -56,6 +67,7 @@ impl ModuleTypesBuilder {
             validator_id: validator.id(),
             types: ModuleTypes::default(),
             trampoline_types: HashMap::default(),
+            exception_types: HashMap::default(),
             wasmparser_to_wasmtime: HashMap::default(),
             already_seen: HashMap::default(),
             defining_rec_group: None,
@@ -301,6 +313,50 @@ impl ModuleTypesBuilder {
         module_interned_index
     }
 
+    /// Define a new exception type when we see a function type used
+    /// in a tag.
+    ///
+    /// The returned `ModuleInternedTypeIndex` gives us a Wasmtime
+    /// type which corresponds to the exception object layout, but
+    /// note that these types do not exist in the Wasm spec: at the
+    /// Wasm level, only function types exist (and tags and exception
+    /// instructions reference them). For implementation reasons, we
+    /// need a separate type to describe the exception object layout,
+    /// and this registers and provides that type.
+    pub fn define_exception_type_for_tag(
+        &mut self,
+        for_func_ty: ModuleInternedTypeIndex,
+    ) -> ModuleInternedTypeIndex {
+        match self.exception_types.entry(for_func_ty) {
+            Entry::Occupied(o) => *o.get(),
+            Entry::Vacant(v) => {
+                let fields = self.types[for_func_ty]
+                    .unwrap_func()
+                    .params()
+                    .iter()
+                    .map(|valtype| WasmFieldType {
+                        element_type: WasmStorageType::Val(*valtype),
+                        mutable: false,
+                    })
+                    .collect();
+                let idx = self.types.push(WasmSubType {
+                    is_final: true,
+                    supertype: None,
+                    composite_type: WasmCompositeType {
+                        inner: WasmCompositeInnerType::Exn(WasmExnType {
+                            func_ty: EngineOrModuleTypeIndex::Module(for_func_ty),
+                            fields,
+                        }),
+                        shared: false,
+                    },
+                });
+                let next = self.types.next_ty();
+                self.types.push_rec_group(idx..next);
+                *v.insert(idx)
+            }
+        }
+    }
+
     /// Returns the result [`ModuleTypes`] of this builder.
     pub fn finish(self) -> ModuleTypes {
         self.types
@@ -347,10 +403,7 @@ impl ModuleTypesBuilder {
         if composite_type.shared {
             return Err(wasm_unsupported!("shared structs are not yet implemented"));
         }
-        match &composite_type.inner {
-            WasmCompositeInnerType::Struct(s) => Ok(s),
-            _ => unreachable!(),
-        }
+        Ok(composite_type.inner.unwrap_struct())
     }
 
     /// Get and unwrap a [`WasmArrayType`] for the given array index.
@@ -367,10 +420,45 @@ impl ModuleTypesBuilder {
         if composite_type.shared {
             return Err(wasm_unsupported!("shared arrays are not yet implemented"));
         }
-        match &composite_type.inner {
-            WasmCompositeInnerType::Array(a) => Ok(a),
-            _ => unreachable!(),
+        Ok(composite_type.inner.unwrap_array())
+    }
+
+    /// Get and unwrap a [`WasmExnType`] for the given exception-type index.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the unwrapped type is not an exception type.
+    ///
+    /// # Errors
+    ///
+    /// For now, fails with an unsupported error if the type is shared.
+    pub fn unwrap_exn(&self, interned_ty: ModuleInternedTypeIndex) -> WasmResult<&WasmExnType> {
+        let composite_type = &self.types[interned_ty].composite_type;
+        if composite_type.shared {
+            return Err(wasm_unsupported!(
+                "shared exceptions are not yet implemented"
+            ));
+        }
+        Ok(composite_type.inner.unwrap_exn())
+    }
+
+    /// Get and unwrap a [`WasmFuncType`] for the given function-type index.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the unwrapped type is not a function type.
+    ///
+    /// # Errors
+    ///
+    /// For now, fails with an unsupported error if the type is shared.
+    pub fn unwrap_func(&self, interned_ty: ModuleInternedTypeIndex) -> WasmResult<&WasmFuncType> {
+        let composite_type = &self.types[interned_ty].composite_type;
+        if composite_type.shared {
+            return Err(wasm_unsupported!(
+                "shared functions are not yet implemented"
+            ));
         }
+        Ok(composite_type.inner.unwrap_func())
     }
 }
 

crates/environ/src/gc.rs

@@ -175,30 +175,30 @@ fn common_struct_layout(
         size,
         align,
         fields,
+        is_exception: false,
     }
 }
 
 /// Common code to define a GC exception object's layout, given the
 /// size and alignment of the collector's GC header and its expected
 /// offset of the array length field.
 #[cfg(any(feature = "gc-null", feature = "gc-drc"))]
-fn common_exn_layout(ty: &WasmExnType, header_size: u32, header_align: u32) -> GcExceptionLayout {
+fn common_exn_layout(ty: &WasmExnType, header_size: u32, header_align: u32) -> GcStructLayout {
     assert!(header_size >= crate::VM_GC_HEADER_SIZE);
     assert!(header_align >= crate::VM_GC_HEADER_ALIGN);
 
     // Compute a struct layout, with extra header size for the
     // `(instance_idx, tag_idx)` fields.
-    let tag_offset = header_size;
     assert!(header_align >= 8);
     let header_size = header_size + 2 * u32::try_from(core::mem::size_of::<u32>()).unwrap();
 
     let (size, align, fields) = common_struct_or_exn_layout(&ty.fields, header_size, header_align);
 
-    GcExceptionLayout {
+    GcStructLayout {
         size,
         align,
-        tag_offset,
         fields,
+        is_exception: true,
     }
 }
 
@@ -211,6 +211,12 @@ pub trait GcTypeLayouts {
     /// element type.
     fn array_length_field_offset(&self) -> u32;
 
+    /// The offset of an exception object's tag reference.
+    ///
+    /// This msut be the same for all exception objects in the heap,
+    /// regardless of their specific signature.
+    fn exception_tag_offset(&self) -> u32;
+
     /// Get this collector's layout for the given composite type.
     ///
     /// Returns `None` if the type is a function type, as functions are not
@@ -235,7 +241,7 @@ pub trait GcTypeLayouts {
     fn struct_layout(&self, ty: &WasmStructType) -> GcStructLayout;
 
     /// Get this collector's layout for the given exception type.
-    fn exn_layout(&self, ty: &WasmExnType) -> GcExceptionLayout;
+    fn exn_layout(&self, ty: &WasmExnType) -> GcStructLayout;
 }
 
 /// The layout of a GC-managed object.
@@ -244,11 +250,8 @@ pub enum GcLayout {
     /// The layout of a GC-managed array object.
     Array(GcArrayLayout),
 
-    /// The layout of a GC-managed struct object.
+    /// The layout of a GC-managed struct or exception object.
     Struct(GcStructLayout),
-
-    /// The layout of a GC-managed exception object.
-    Exception(GcExceptionLayout),
 }
 
 impl From<GcArrayLayout> for GcLayout {
@@ -263,12 +266,6 @@ impl From<GcStructLayout> for GcLayout {
     }
 }
 
-impl From<GcExceptionLayout> for GcLayout {
-    fn from(layout: GcExceptionLayout) -> Self {
-        Self::Exception(layout)
-    }
-}
-
 impl GcLayout {
     /// Get the underlying `GcStructLayout`, or panic.
     #[track_caller]
@@ -287,15 +284,6 @@ impl GcLayout {
             _ => panic!("GcLayout::unwrap_array on non-array GC layout"),
         }
     }
-
-    /// Get the underlying `GcExceptionLayout`, or panic.
-    #[track_caller]
-    pub fn unwrap_exception(&self) -> &GcExceptionLayout {
-        match self {
-            Self::Exception(e) => e,
-            _ => panic!("GcLayout::unwrap_exception on a non-exception GC layout"),
-        }
-    }
 }
 
 /// The layout of a GC-managed array.
@@ -352,7 +340,7 @@ impl GcArrayLayout {
     }
 }
 
-/// The layout for a GC-managed struct type.
+/// The layout for a GC-managed struct type or exception 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
@@ -361,6 +349,12 @@ impl GcArrayLayout {
 ///
 /// All offsets are from the start of the object; that is, the size of the GC
 /// header (for example) is included in the offset.
+///
+/// Note that these are reused between structs and exceptions to avoid
+/// unnecessary code duplication. In both cases, the objects are
+/// tuples of typed fields with a certain size. The only difference in
+/// practice is that an exception object also carries a tag reference
+/// (at a fixed offset as per `GcTypeLayouts::exception_tag_offset`).
 #[derive(Clone, Debug)]
 pub struct GcStructLayout {
     /// The size (in bytes) of this struct.
@@ -372,6 +366,9 @@ pub struct GcStructLayout {
     /// The fields of this struct. The `i`th entry contains information about
     /// the `i`th struct field's layout.
     pub fields: Vec<GcStructLayoutField>,
+
+    /// Whether this is an exception object layout.
+    pub is_exception: bool,
 }
 
 impl GcStructLayout {
@@ -397,41 +394,6 @@ pub struct GcStructLayoutField {
     pub is_gc_ref: bool,
 }
 
-/// The layout for a GC-managed exception object.
-///
-/// 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 GcExceptionLayout {
-    /// The size (in bytes) of this struct.
-    pub size: u32,
-
-    /// The alignment (in bytes) of this struct.
-    pub align: u32,
-
-    /// The offset of the VMTagImport pointer.
-    pub tag_offset: u32,
-
-    /// The fields of this exception object. The `i`th entry contains
-    /// information about the `i`th parameter in the associated tag
-    /// type.
-    pub fields: Vec<GcStructLayoutField>,
-}
-
-impl GcExceptionLayout {
-    /// Get a `core::alloc::Layout` for an exception object of this type.
-    pub fn layout(&self) -> Layout {
-        let size = usize::try_from(self.size).unwrap();
-        let align = usize::try_from(self.align).unwrap();
-        Layout::from_size_align(size, align).unwrap()
-    }
-}
-
 /// 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

@@ -11,6 +11,9 @@ pub const HEADER_ALIGN: u32 = 8;
 /// The offset of the length field in a `VMDrcArrayHeader`.
 pub const ARRAY_LENGTH_OFFSET: u32 = HEADER_SIZE;
 
+/// The offset of the tag fields in an exception header.
+pub const EXCEPTION_TAG_OFFSET: u32 = HEADER_SIZE;
+
 /// The bit within a `VMDrcHeader`'s reserved bits that is the mark
 /// bit. Collectively, this bit in all the heap's objects' headers implements
 /// the precise-stack-roots set.
@@ -29,6 +32,10 @@ impl GcTypeLayouts for DrcTypeLayouts {
         ARRAY_LENGTH_OFFSET
     }
 
+    fn exception_tag_offset(&self) -> u32 {
+        EXCEPTION_TAG_OFFSET
+    }
+
     fn array_layout(&self, ty: &WasmArrayType) -> GcArrayLayout {
         common_array_layout(ty, HEADER_SIZE, HEADER_ALIGN, ARRAY_LENGTH_OFFSET)
     }
@@ -37,7 +44,7 @@ impl GcTypeLayouts for DrcTypeLayouts {
         common_struct_layout(ty, HEADER_SIZE, HEADER_ALIGN)
     }
 
-    fn exn_layout(&self, ty: &WasmExnType) -> GcExceptionLayout {
+    fn exn_layout(&self, ty: &WasmExnType) -> GcStructLayout {
         common_exn_layout(ty, HEADER_SIZE, HEADER_ALIGN)
     }
 }