Refactor internals of PassiveElementSegment (#13187)

In looking over #13179 I found it a bit brittle to access `anyref`
fields unconditionally where the determination that the type of the
segment was a GC reference was made much further away. I also found it
somewhat confusing to have methods like `clone_gc_ref` and
`drop_gc_ref` applied to all values which happened to be noops for
non-gc-ref types. I've refactored things a bit internally here to plumb
a `WasmRefType` around and have additionally added comments to why
`anyref` accessors are used despite the value possibly having an
`exnref` or `externref` type.

Author: alexcrichton

crates/environ/src/compile/module_environ.rs

@@ -3,6 +3,7 @@ use crate::module::{
     FuncRefIndex, Initializer, MemoryInitialization, MemoryInitializer, Module, TableSegment,
     TableSegmentElements,
 };
+use crate::prelude::*;
 use crate::{
     ConstExpr, ConstOp, DataIndex, DefinedFuncIndex, ElemIndex, EngineOrModuleTypeIndex,
     EntityIndex, EntityType, FuncIndex, FuncKey, GlobalIndex, IndexType, InitMemory, MemoryIndex,
@@ -11,7 +12,6 @@ use crate::{
     Tunables, TypeConvert, TypeIndex, WasmError, WasmHeapTopType, WasmHeapType, WasmResult,
     WasmValType, WasmparserTypeConverter,
 };
-use crate::{NeedsGcRooting, prelude::*};
 use cranelift_entity::SecondaryMap;
 use cranelift_entity::packed_option::ReservedValue;
 use std::borrow::Cow;
@@ -543,11 +543,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                                 }
                             }
                             TableSegmentElements::Expressions {
-                                needs_gc_rooting: if ty.is_vmgcref_type_and_not_i31() {
-                                    NeedsGcRooting::Yes
-                                } else {
-                                    NeedsGcRooting::No
-                                },
+                                ty,
                                 exprs: exprs.into(),
                             }
                         }

crates/environ/src/module.rs

@@ -266,15 +266,6 @@ pub struct TableSegment {
     pub elements: TableSegmentElements,
 }
 
-/// Does something need GC rooting?
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
-pub enum NeedsGcRooting {
-    /// GC rooting is needed.
-    Yes,
-    /// GC rooting is not needed.
-    No,
-}
-
 /// Elements of a table segment, either a list of functions or list of arbitrary
 /// expressions.
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -284,9 +275,8 @@ pub enum TableSegmentElements {
     Functions(Box<[FuncIndex]>),
     /// Arbitrary expressions, aka either functions, null or a load of a global.
     Expressions {
-        /// Is this expression's result of a type that needs GC rooting and
-        /// tracing?
-        needs_gc_rooting: NeedsGcRooting,
+        /// The type of each element in `exprs`.
+        ty: WasmRefType,
         /// The const expressions for this segment's elements.
         exprs: Box<[ConstExpr]>,
     },

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

@@ -36,9 +36,9 @@ use wasmtime_environ::ModuleInternedTypeIndex;
 use wasmtime_environ::error::OutOfMemory;
 use wasmtime_environ::{
     DataIndex, DefinedGlobalIndex, DefinedMemoryIndex, DefinedTableIndex, DefinedTagIndex,
-    ElemIndex, EntityIndex, EntityRef, FuncIndex, GlobalIndex, HostPtr, MemoryIndex,
-    NeedsGcRooting, PtrSize, TableIndex, TableInitialValue, TagIndex, Trap, VMCONTEXT_MAGIC,
-    VMOffsets, VMSharedTypeIndex, packed_option::ReservedValue,
+    ElemIndex, EntityIndex, EntityRef, FuncIndex, GlobalIndex, HostPtr, MemoryIndex, PtrSize,
+    TableIndex, TableInitialValue, TagIndex, Trap, VMCONTEXT_MAGIC, VMOffsets, VMSharedTypeIndex,
+    WasmRefType, packed_option::ReservedValue,
 };
 #[cfg(feature = "wmemcheck")]
 use wasmtime_wmemcheck::Wmemcheck;
@@ -225,7 +225,7 @@ impl Instance {
         gc_roots: &mut crate::vm::GcRootsList,
     ) {
         for segment in self.passive_elements_mut().iter_mut() {
-            if segment.needs_gc_rooting() == NeedsGcRooting::Yes {
+            if segment.needs_gc_rooting {
                 for e in segment.elements() {
                     let Some(root) = e.as_vmgc_ref_ptr() else {
                         continue;
@@ -1922,58 +1922,59 @@ impl<T: InstanceLayout> Drop for OwnedInstance<T> {
 
 #[derive(Debug)]
 pub(crate) struct PassiveElementSegment {
-    needs_gc_rooting: NeedsGcRooting,
+    needs_gc_rooting: bool,
     elements: TryVec<ValRaw>,
 }
 
 impl PassiveElementSegment {
     /// Create a new passive element segment with the given capacity.
-    pub(crate) fn new(
-        needs_gc_rooting: NeedsGcRooting,
-        capacity: usize,
-    ) -> Result<Self, OutOfMemory> {
+    pub(crate) fn new(ty: WasmRefType, capacity: usize) -> Result<Self, OutOfMemory> {
         Ok(Self {
-            needs_gc_rooting,
+            needs_gc_rooting: ty.is_vmgcref_type_and_not_i31(),
             elements: TryVec::with_capacity(capacity)?,
         })
     }
 
     /// Push a value onto this passive element segment.
     ///
-    /// NB: Does not type check the value, relies on callers to ensure the value
-    /// is of the correct type (generally, due to validation).
+    /// NB: Does not type check the value, relies on callers to ensure the
+    /// value is of the correct type (generally, due to validation).
     pub(crate) fn push(&mut self, store: &mut StoreOpaque, val: Val) -> Result<()> {
-        let val = {
+        let mut val = {
             let mut store = AutoAssertNoGc::new(store);
             val.to_raw_(&mut store)?
         };
-        let val = self.clone_gc_ref(store, val);
-        self.elements.push(val)?;
-        Ok(())
-    }
-
-    fn clone_gc_ref(&mut self, store: &mut StoreOpaque, val: ValRaw) -> ValRaw {
-        if let NeedsGcRooting::Yes = self.needs_gc_rooting {
+        if self.needs_gc_rooting {
+            // Note that `anyref` accessors and constructors are used here
+            // without actually checking the type of this segment or value. The
+            // representation and handling of all three is the same which means
+            // that this should work out.
             let gc_ref = val.get_anyref();
+            debug_assert_eq!(gc_ref, val.get_exnref());
+            debug_assert_eq!(gc_ref, val.get_externref());
             if let Some(gc_ref) = VMGcRef::from_raw_u32(gc_ref) {
                 if let Some(gc_store) = store.optional_gc_store_mut() {
-                    return ValRaw::anyref(gc_store.clone_gc_ref(&gc_ref).as_raw_u32());
+                    val = ValRaw::anyref(gc_store.clone_gc_ref(&gc_ref).as_raw_u32());
                 }
             }
         }
-        val
+        self.elements.push(val)?;
+        Ok(())
     }
 
     /// Clear this segment's elements.
     pub(crate) fn clear(&mut self, mut gc_store: Option<&mut GcStore>) {
-        for val in mem::take(&mut self.elements) {
-            self.drop_gc_ref(&mut gc_store, val);
+        let elements = mem::take(&mut self.elements);
+        if !self.needs_gc_rooting {
+            return;
         }
-    }
-
-    fn drop_gc_ref(&mut self, gc_store: &mut Option<&mut GcStore>, val: ValRaw) {
-        if let NeedsGcRooting::Yes = self.needs_gc_rooting {
+        for val in elements {
+            // Like above, `anyref` accessors are used here even if this
+            // element segment has a different type because all of the vmgcref
+            // types are treated the same way.
             let gc_ref = val.get_anyref();
+            debug_assert_eq!(gc_ref, val.get_exnref());
+            debug_assert_eq!(gc_ref, val.get_externref());
             if let Some(gc_ref) = VMGcRef::from_raw_u32(gc_ref) {
                 if let Some(gc_store) = gc_store.as_deref_mut() {
                     let _ = gc_store.drop_gc_ref(gc_ref);
@@ -1982,12 +1983,6 @@ impl PassiveElementSegment {
         }
     }
 
-    /// Whether this segment needs GC rooting and tracing.
-    #[cfg(feature = "gc")]
-    pub(crate) fn needs_gc_rooting(&self) -> NeedsGcRooting {
-        self.needs_gc_rooting
-    }
-
     /// The elements of this segment.
     pub(crate) fn elements(&self) -> &[ValRaw] {
         &self.elements

crates/wasmtime/src/runtime/vm/instance/allocator.rs

@@ -14,8 +14,8 @@ use core::pin::Pin;
 use core::{mem, ptr};
 use wasmtime_environ::{
     DefinedMemoryIndex, DefinedTableIndex, EntityRef, HostPtr, InitMemory, MemoryInitialization,
-    MemoryInitializer, Module, NeedsGcRooting, SizeOverflow, TableInitialValue,
-    TableSegmentElements, Trap, VMOffsets,
+    MemoryInitializer, Module, SizeOverflow, TableInitialValue, TableSegmentElements, Trap,
+    VMOffsets, WasmRefType,
 };
 
 #[cfg(feature = "gc")]
@@ -597,10 +597,7 @@ async fn initialize_tables(
                     table.set_(&mut store, i, func.into())?;
                 }
             }
-            TableSegmentElements::Expressions {
-                exprs,
-                needs_gc_rooting: _,
-            } => {
+            TableSegmentElements::Expressions { exprs, ty: _ } => {
                 for (i, expr) in positions.zip(exprs) {
                     let val = const_evaluator
                         .eval(&mut store, limiter.as_deref_mut(), context, expr)
@@ -861,7 +858,7 @@ async fn initialize_passive_elements(
         match segment {
             TableSegmentElements::Functions(func_indices) => {
                 let mut segment =
-                    PassiveElementSegment::new(NeedsGcRooting::No, func_indices.len())?;
+                    PassiveElementSegment::new(WasmRefType::FUNCREF, func_indices.len())?;
                 for func_idx in func_indices {
                     let (instance, registry) =
                         store.instance_and_module_registry_mut(context.instance);
@@ -873,11 +870,8 @@ async fn initialize_passive_elements(
                 debug_assert_eq!(instance.passive_elements.len(), idx.index());
                 instance.passive_elements_mut().push(segment)?;
             }
-            TableSegmentElements::Expressions {
-                needs_gc_rooting,
-                exprs,
-            } => {
-                let mut segment = PassiveElementSegment::new(*needs_gc_rooting, exprs.len())?;
+            TableSegmentElements::Expressions { ty, exprs } => {
+                let mut segment = PassiveElementSegment::new(*ty, exprs.len())?;
                 for expr in exprs {
                     let mut store = OpaqueRootScope::new(&mut *store);
                     let val = const_evaluator