Change the representation of `VMGcKind` to allow for fast subtyping checks (#8926)

crates/environ/src/gc.rs

@@ -18,43 +18,108 @@ pub const NON_NULL_NON_I31_MASK: u64 = !I31_DISCRIMINANT;
 
 /// The kind of an object in a GC heap.
 ///
-/// This is accessed from Wasm JIT code.
+/// Note that this type is accessed from Wasm JIT code.
+///
+/// `VMGcKind` is a bitset where to test if `a` is a subtype of an
+/// "abstract-ish" type `b`, we can simply use a single bitwise-and operation:
+///
+/// ```ignore
+/// a <: b   iff   a & b == b
+/// ```
+///
+/// For example, because `VMGcKind::AnyRef` has the high bit set, every kind
+/// representing some subtype of `anyref` also has its high bit set.
+///
+/// We say "abstract-ish" type because in addition to the abstract heap types
+/// (other than `i31`) we also have variants for `externref`s that have been
+/// converted into an `anyref` via `extern.convert_any` and `externref`s that
+/// have been convered into an `anyref` via `any.convert_extern`. Note that in
+/// the latter case, because `any.convert_extern $foo` produces a value that is
+/// not an instance of `eqref`, `VMGcKind::AnyOfExternRef & VMGcKind::EqRef !=
+/// VMGcKind::EqRef`.
+///
+/// Furthermore, this type only uses the highest 6 bits of its `u32`
+/// representation, allowing the lower 26 bytes to be bitpacked with other stuff
+/// as users see fit.
 #[repr(u32)]
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[rustfmt::skip]
+#[allow(missing_docs)]
 pub enum VMGcKind {
-    /// An `externref` holding some kind of host data.
-    ExternRef = 0b00 << 30,
-    //
-    // When we support more GC types, we will complete this type with:
-    //
-    // /// An `anyref` (or one of its subtypes, other than `i31ref`).
-    // AnyRef = 0b01 << 30,
-    //
-    // /// An `anyref` that was wrapped into an `externref` via
-    // /// `extern.convert_any`.
-    // ExternOfAnyRef = 0b10 << 30,
-    //
-    // /// An `externref` that was wrapped into an `anyref` via
-    // /// `any.convert_extern`.
-    // AnyOfExternRef = 0b11 << 30,
+    ExternRef      = 0b010000 << 26,
+    ExternOfAnyRef = 0b011000 << 26,
+    AnyRef         = 0b100000 << 26,
+    AnyOfExternRef = 0b100100 << 26,
+    EqRef          = 0b101000 << 26,
+    ArrayRef       = 0b101001 << 26,
+    StructRef      = 0b101010 << 26,
 }
 
 impl VMGcKind {
-    /// Mask this value with a `u32` to turn it into a valid `VMGcKind`.
-    pub const MASK: u32 = 0b11 << 30;
+    /// Mask this value with a `u32` to get just the bits that `VMGcKind` uses.
+    pub const MASK: u32 = 0b111111 << 26;
 
     /// Mask this value with a `u32` that potentially contains a `VMGcKind` to
     /// get the bits that `VMGcKind` doesn't use.
     pub const UNUSED_MASK: u32 = !Self::MASK;
 
     /// Convert the given value into a `VMGcKind` by masking off the unused
-    /// bits.
-    pub const fn from_u32(val: u32) -> VMGcKind {
+    /// bottom bits.
+    pub fn from_high_bits_of_u32(val: u32) -> VMGcKind {
         let masked = val & Self::MASK;
-        assert!(
-            masked == Self::ExternRef as u32,
-            "not all masked bit patterns are valid `VMGcKind`s yet"
-        );
-        Self::ExternRef
+        match masked {
+            x if x == Self::ExternRef as u32 => Self::ExternRef,
+            x if x == Self::ExternOfAnyRef as u32 => Self::ExternOfAnyRef,
+            x if x == Self::AnyRef as u32 => Self::AnyRef,
+            x if x == Self::AnyOfExternRef as u32 => Self::AnyOfExternRef,
+            x if x == Self::EqRef as u32 => Self::EqRef,
+            x if x == Self::ArrayRef as u32 => Self::ArrayRef,
+            x if x == Self::StructRef as u32 => Self::StructRef,
+            _ => panic!("invalid `VMGcKind`: {masked:#032b}"),
+        }
+    }
+
+    /// Does this kind match the other kind?
+    ///
+    /// That is, is this kind a subtype of the other kind?
+    pub fn matches(self, other: Self) -> bool {
+        (self as u32) & (other as u32) == (other as u32)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::VMGcKind::*;
+    use crate::prelude::*;
+
+    #[test]
+    fn kind_matches() {
+        let all = [
+            ExternRef,
+            ExternOfAnyRef,
+            AnyRef,
+            AnyOfExternRef,
+            EqRef,
+            ArrayRef,
+            StructRef,
+        ];
+
+        for (sup, subs) in [
+            (ExternRef, vec![ExternOfAnyRef]),
+            (ExternOfAnyRef, vec![]),
+            (AnyRef, vec![AnyOfExternRef, EqRef, ArrayRef, StructRef]),
+            (AnyOfExternRef, vec![]),
+            (EqRef, vec![ArrayRef, StructRef]),
+            (ArrayRef, vec![]),
+            (StructRef, vec![]),
+        ] {
+            assert!(sup.matches(sup));
+            for sub in &subs {
+                assert!(sub.matches(sup));
+            }
+            for kind in all.iter().filter(|k| **k != sup && !subs.contains(k)) {
+                assert!(!kind.matches(sup));
+            }
+        }
     }
 }

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

@@ -49,39 +49,39 @@ impl VMGcHeader {
     /// Get the kind of GC object that this is.
     pub fn kind(&self) -> VMGcKind {
         let upper = u32::try_from(self.0 >> 32).unwrap();
-        VMGcKind::from_u32(upper)
+        VMGcKind::from_high_bits_of_u32(upper)
     }
 
-    /// Get the reserved 30 bits in this header.
+    /// Get the reserved 26 bits in this header.
     ///
     /// These are bits are reserved for `GcRuntime` implementations to make use
     /// of however they see fit.
-    pub fn reserved_u30(&self) -> u32 {
+    pub fn reserved_u26(&self) -> u32 {
         let upper = u32::try_from(self.0 >> 32).unwrap();
         upper & VMGcKind::UNUSED_MASK
     }
 
-    /// Set the 30-bit reserved value.
+    /// Set the 26-bit reserved value.
     ///
     /// # Panics
     ///
     /// Panics if the given `value` has any of the upper 2 bits set.
-    pub fn set_reserved_u30(&mut self, value: u32) {
+    pub fn set_reserved_u26(&mut self, value: u32) {
         assert_eq!(
             value & VMGcKind::MASK,
             0,
-            "VMGcHeader::set_reserved_u30 with value using more than 30 bits"
+            "VMGcHeader::set_reserved_u26 with value using more than 26 bits"
         );
         self.0 |= u64::from(value) << 32;
     }
 
-    /// Set the 30-bit reserved value.
+    /// Set the 26-bit reserved value.
     ///
     /// # Safety
     ///
-    /// The given `value` must only use the lower 30 bits; its upper 2 bits must
+    /// The given `value` must only use the lower 26 bits; its upper 2 bits must
     /// be unset.
-    pub unsafe fn unchecked_set_reserved_u30(&mut self, value: u32) {
+    pub unsafe fn unchecked_set_reserved_u26(&mut self, value: u32) {
         self.0 |= u64::from(value) << 32;
     }