Fix select missing stack map declarations for GC refs

crates/cranelift/src/func_environ.rs

@@ -1371,6 +1371,17 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
             builder.ins().stack_store(vmctx, slot, 0);
         }
     }
+
+    pub(crate) fn val_ty_needs_stack_map(&self, ty: WasmValType) -> bool {
+        match ty {
+            WasmValType::Ref(r) => self.heap_ty_needs_stack_map(r.heap_type),
+            _ => false,
+        }
+    }
+
+    pub(crate) fn heap_ty_needs_stack_map(&self, ty: WasmHeapType) -> bool {
+        ty.is_vmgcref_type_and_not_i31() && !ty.is_bottom()
+    }
 }
 
 impl TranslateTrap for FuncEnvironment<'_> {
@@ -2552,13 +2563,7 @@ impl<'module_environment> TargetEnvironment for FuncEnvironment<'module_environm
 
     fn reference_type(&self, wasm_ty: WasmHeapType) -> (ir::Type, bool) {
         let ty = crate::reference_type(wasm_ty, self.pointer_type());
-        let needs_stack_map = match wasm_ty.top() {
-            WasmHeapTopType::Extern | WasmHeapTopType::Any | WasmHeapTopType::Exn => true,
-            WasmHeapTopType::Func => false,
-            // TODO(#10248) Once continuations can be stored on the GC heap, we
-            // will need stack maps for continuation objects.
-            WasmHeapTopType::Cont => false,
-        };
+        let needs_stack_map = self.heap_ty_needs_stack_map(wasm_ty);
         (ty, needs_stack_map)
     }
 

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

@@ -261,8 +261,7 @@ impl DrcCompiler {
         new_val: ir::Value,
         flags: ir::MemFlags,
     ) -> WasmResult<()> {
-        let (ref_ty, needs_stack_map) = func_env.reference_type(ty.heap_type);
-        debug_assert!(needs_stack_map);
+        let (ref_ty, _) = func_env.reference_type(ty.heap_type);
 
         // Special case for references to uninhabited bottom types: see
         // `translate_write_gc_reference` for details.
@@ -585,8 +584,7 @@ impl GcCompiler for DrcCompiler {
 
         assert!(ty.is_vmgcref_type());
 
-        let (reference_type, needs_stack_map) = func_env.reference_type(ty.heap_type);
-        debug_assert!(needs_stack_map);
+        let (reference_type, _) = func_env.reference_type(ty.heap_type);
 
         // Special case for references to uninhabited bottom types: the
         // reference must either be nullable and we can just eagerly return
@@ -737,8 +735,7 @@ impl GcCompiler for DrcCompiler {
     ) -> WasmResult<()> {
         assert!(ty.is_vmgcref_type());
 
-        let (ref_ty, needs_stack_map) = func_env.reference_type(ty.heap_type);
-        debug_assert!(needs_stack_map);
+        let (ref_ty, _) = func_env.reference_type(ty.heap_type);
 
         // Special case for references to uninhabited bottom types: either the
         // reference is nullable and we can just eagerly store null into `dst`

crates/cranelift/src/translate/code_translator.rs

@@ -203,31 +203,47 @@ pub fn translate_operator(
         Operator::Drop => {
             environ.stacks.pop1();
         }
-        Operator::Select => {
-            let (mut arg1, mut arg2, cond) = environ.stacks.pop3();
-            if builder.func.dfg.value_type(arg1).is_vector() {
-                arg1 = optionally_bitcast_vector(arg1, I8X16, builder);
-            }
-            if builder.func.dfg.value_type(arg2).is_vector() {
-                arg2 = optionally_bitcast_vector(arg2, I8X16, builder);
-            }
-            environ.stacks.push1(builder.ins().select(cond, arg1, arg2));
+        Operator::Nop => {
+            // We do nothing
         }
-        Operator::TypedSelect { ty: _ } => {
+        Operator::Select
+        | Operator::TypedSelect {
             // We ignore the explicit type parameter as it is only needed for
             // validation, which we require to have been performed before
             // translation.
+            ty: _,
+        } => {
             let (mut arg1, mut arg2, cond) = environ.stacks.pop3();
+
             if builder.func.dfg.value_type(arg1).is_vector() {
                 arg1 = optionally_bitcast_vector(arg1, I8X16, builder);
             }
             if builder.func.dfg.value_type(arg2).is_vector() {
                 arg2 = optionally_bitcast_vector(arg2, I8X16, builder);
             }
-            environ.stacks.push1(builder.ins().select(cond, arg1, arg2));
-        }
-        Operator::Nop => {
-            // We do nothing
+
+            let val = builder.ins().select(cond, arg1, arg2);
+
+            // If either of the input types need inclusion in stack maps, then
+            // the result will as well.
+            //
+            // Note that we don't need to check whether the result's type needs
+            // inclusion in stack maps (that would be a conservative over
+            // approximation) because the input types give us more-precise
+            // information than the result type does. For example, the result
+            // does not need inclusion in stack maps in the scenario where both
+            // inputs are `i31ref`s and the result is an `anyref`. Even though
+            // `anyref`s normally do require inclusion in stack maps, in this
+            // particular case, we know that we are dealing with an `anyref`
+            // that doesn't actually require inclusion.
+            if operand_types
+                .iter()
+                .any(|ty| environ.val_ty_needs_stack_map(*ty))
+            {
+                builder.declare_value_needs_stack_map(val);
+            }
+
+            environ.stacks.push1(val);
         }
         Operator::Unreachable => {
             environ.trap(builder, crate::TRAP_UNREACHABLE);

tests/all/gc.rs

@@ -1800,3 +1800,80 @@ fn array_init_elem_oom() -> Result<()> {
 
     Ok(())
 }
+
+// The result of a `select` instruction with a GC reference type should be
+// declared as needing a stack map.
+#[test]
+#[cfg_attr(miri, ignore)]
+fn select_gc_ref_stack_map() -> Result<()> {
+    let _ = env_logger::try_init();
+
+    let mut config = Config::new();
+    config.wasm_function_references(true);
+    config.wasm_gc(true);
+    // Use pooling allocator with a tiny GC heap so that GC is triggered
+    // frequently and freed memory is reused quickly.
+    let mut pool = crate::small_pool_config();
+    pool.max_memory_size(1 << 16); // 64 KiB
+    config.allocation_strategy(pool);
+
+    let engine = Engine::new(&config)?;
+    let module = Module::new(
+        &engine,
+        r#"
+            (module
+                (type $pair (struct (field (mut i32))))
+                (type $arr (array (mut i8)))
+
+                ;; Allocate many objects to fill the GC heap and trigger
+                ;; collection. After GC frees everything, subsequent
+                ;; allocations reuse the freed memory.
+                (func $force_gc
+                    (local i32)
+                    (local.set 0 (i32.const 64))
+                    (loop $l
+                        (drop (array.new $arr (i32.const 0) (i32.const 1024)))
+                        (local.set 0 (i32.sub (local.get 0) (i32.const 1)))
+                        (br_if $l (local.get 0))
+                    )
+                )
+
+                (func (export "test") (param $cond i32) (result i32)
+                    ;; The select result stays on the Wasm operand stack (never
+                    ;; stored in a local variable). If the new SSA value created
+                    ;; by select is not declared for stack maps, the GC will
+                    ;; free it.
+                    (select (result (ref null $pair))
+                        (struct.new $pair (i32.const 111))
+                        (ref.null $pair)
+                        (local.get $cond)
+                    )
+
+                    ;; This call is a safepoint. The select result is live
+                    ;; across it. The called function triggers GC which will
+                    ;; free the struct if it is incorrectly omitted from the
+                    ;; safepoint's stack map, and the new allocations would
+                    ;; overwrite the freed memory.
+                    (call $force_gc)
+
+                    ;; Use the select result. If it was incorrectly freed, then
+                    ;; this will have the wrong value.
+                    (struct.get $pair 0)
+                )
+            )
+        "#,
+    )?;
+
+    let mut store = Store::new(&engine, ());
+    let instance = Instance::new(&mut store, &module, &[])?;
+    let test = instance.get_typed_func::<(i32,), i32>(&mut store, "test")?;
+
+    // Run multiple times to increase chance of triggering GC at the right
+    // moment.
+    for _ in 0..30 {
+        let result = test.call(&mut store, (1,))?;
+        assert_eq!(result, 111);
+    }
+
+    Ok(())
+}

tests/disas/gc/typed-select-and-stack-maps.wat

@@ -0,0 +1,90 @@
+;;! target = "x86_64"
+;;! flags = "-W function-references,gc"
+;;! test = "optimize"
+
+(module
+  (import "" "observe" (func $observe (param anyref)))
+  (import "" "safepoint" (func $safepoint))
+
+  (func (param structref i31ref i32)
+    ;; Select from two types, one of which requires inclusion in stack maps,
+    ;; resulting in a value that also requires inclusion in stack maps.
+    (select (result anyref)
+            (local.get 0)
+            (local.get 1)
+            (local.get 2))
+
+    ;; Make a call, which is a safepoint and has stack maps.
+    (call $safepoint)
+
+    ;; Observe the result of the select to keep it alive across the call.
+    call $observe
+  )
+
+  (func (param i31ref i31ref i32)
+        ;; Select from two types that do not require inclusion in stack maps,
+        ;; resulting in one that (normally) does. In this case, however, we
+        ;; shouldn't include the value in a stack map, because we know that the
+        ;; anyref cannot be an instance of a subtype that actually does require
+        ;; inclusion in stack maps.
+        (select (result anyref)
+                (local.get 0)
+                (local.get 1)
+                (local.get 2))
+
+        ;; Make a call, which is a safepoint and has stack maps.
+        (call $safepoint)
+
+        ;; Observe the result of the select to keep it alive across the call.
+        call $observe
+        )
+)
+;; function u0:0(i64 vmctx, i64, i32, i32, i32) tail {
+;;     ss0 = explicit_slot 4, align = 4
+;;     gv0 = vmctx
+;;     gv1 = load.i64 notrap aligned readonly gv0+8
+;;     gv2 = load.i64 notrap aligned gv1+24
+;;     gv3 = vmctx
+;;     sig0 = (i64 vmctx, i64) tail
+;;     sig1 = (i64 vmctx, i64, i32) tail
+;;     stack_limit = gv2
+;;
+;;                                 block0(v0: i64, v1: i64, v2: i32, v3: i32, v4: i32):
+;; @0049                               v5 = select v4, v2, v3
+;;                                     v14 = stack_addr.i64 ss0
+;;                                     store notrap v5, v14
+;; @004c                               v8 = load.i64 notrap aligned readonly can_move v0+72
+;; @004c                               v7 = load.i64 notrap aligned readonly can_move v0+88
+;; @004c                               call_indirect sig0, v8(v7, v0), stack_map=[i32 @ ss0+0]
+;;                                     v12 = load.i32 notrap v14
+;; @004e                               v11 = load.i64 notrap aligned readonly can_move v0+48
+;; @004e                               v10 = load.i64 notrap aligned readonly can_move v0+64
+;; @004e                               call_indirect sig1, v11(v10, v0, v12)
+;; @0050                               jump block1
+;;
+;;                                 block1:
+;; @0050                               return
+;; }
+;;
+;; function u0:1(i64 vmctx, i64, i32, i32, i32) tail {
+;;     gv0 = vmctx
+;;     gv1 = load.i64 notrap aligned readonly gv0+8
+;;     gv2 = load.i64 notrap aligned gv1+24
+;;     gv3 = vmctx
+;;     sig0 = (i64 vmctx, i64) tail
+;;     sig1 = (i64 vmctx, i64, i32) tail
+;;     stack_limit = gv2
+;;
+;;                                 block0(v0: i64, v1: i64, v2: i32, v3: i32, v4: i32):
+;; @005c                               v8 = load.i64 notrap aligned readonly can_move v0+72
+;; @005c                               v7 = load.i64 notrap aligned readonly can_move v0+88
+;; @005c                               call_indirect sig0, v8(v7, v0)
+;; @005e                               v11 = load.i64 notrap aligned readonly can_move v0+48
+;; @005e                               v10 = load.i64 notrap aligned readonly can_move v0+64
+;; @0059                               v5 = select v4, v2, v3
+;; @005e                               call_indirect sig1, v11(v10, v0, v5)
+;; @0060                               jump block1
+;;
+;;                                 block1:
+;; @0060                               return
+;; }