Repository: bytecodealliance/wasmtime
Pull request #9260 (2024-09-16 19:17:01)

Stop special-casing `i31ref` accesses in many places
Consolidate all `i31ref` special casing into the `translate_{read,write}_gc_reference` methods, which already emit more or less inline GC barriers depending on the reference's heap type, so it makes sense to also handle `i31ref` there.

<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->


fitzgen: Thought I landed this on Friday, but I guess not.

Patch:
diff --git a/cranelift/codegen/src/isa/aarch64/abi.rs b/cranelift/codegen/src/isa/aarch64/abi.rs
index 0fd19f973..adb66e62d 100644
--- a/cranelift/codegen/src/isa/aarch64/abi.rs
+++ b/cranelift/codegen/src/isa/aarch64/abi.rs
@@ -161,31 +161,19 @@ impl ABIMachineSpec for AArch64MachineDeps {
 
             let (rcs, reg_types) = Inst::rc_for_type(param.value_type)?;
 
-            if matches!(
-                param.purpose,
-                ir::ArgumentPurpose::StructArgument(_) | ir::ArgumentPurpose::StructReturn
-            ) {
+            if let ir::ArgumentPurpose::StructReturn = param.purpose {
                 assert!(
                     call_conv != isa::CallConv::Tail,
-                    "support for {:?} parameters is not implemented for the `tail` calling \
-                    convention yet",
-                    param.purpose,
+                    "support for StructReturn parameters is not implemented for the `tail` \
+                    calling convention yet",
                 );
             }
 
-            if let ir::ArgumentPurpose::StructArgument(size) = param.purpose {
-                assert_eq!(args_or_rets, ArgsOrRets::Args);
-                let offset = next_stack as i64;
-                let size = size;
-                assert!(size % 8 == 0, "StructArgument size is not properly aligned");
-                next_stack += size;
-                args.push(ABIArg::StructArg {
-                    pointer: None,
-                    offset,
-                    size: size as u64,
-                    purpose: param.purpose,
-                });
-                continue;
+            if let ir::ArgumentPurpose::StructArgument(_) = param.purpose {
+                panic!(
+                    "StructArgument parameters are not supported on arm64. \
+                    Use regular pointer arguments instead."
+                );
             }
 
             if let ir::ArgumentPurpose::StructReturn = param.purpose {
diff --git a/cranelift/codegen/src/isa/riscv64/abi.rs b/cranelift/codegen/src/isa/riscv64/abi.rs
index 258bac24d..7c0d944ee 100644
--- a/cranelift/codegen/src/isa/riscv64/abi.rs
+++ b/cranelift/codegen/src/isa/riscv64/abi.rs
@@ -113,17 +113,11 @@ impl ABIMachineSpec for Riscv64MachineDeps {
         let mut next_stack: u32 = 0;
 
         for param in params {
-            if let ir::ArgumentPurpose::StructArgument(size) = param.purpose {
-                let offset = next_stack;
-                assert!(size % 8 == 0, "StructArgument size is not properly aligned");
-                next_stack += size;
-                args.push(ABIArg::StructArg {
-                    pointer: None,
-                    offset: offset as i64,
-                    size: size as u64,
-                    purpose: param.purpose,
-                });
-                continue;
+            if let ir::ArgumentPurpose::StructArgument(_) = param.purpose {
+                panic!(
+                    "StructArgument parameters are not supported on riscv64. \
+                    Use regular pointer arguments instead."
+                );
             }
 
             // Find regclass(es) of the register(s) used to store a value of this type.
diff --git a/cranelift/codegen/src/isa/s390x/abi.rs b/cranelift/codegen/src/isa/s390x/abi.rs
index f6de879b1..89d0db256 100644
--- a/cranelift/codegen/src/isa/s390x/abi.rs
+++ b/cranelift/codegen/src/isa/s390x/abi.rs
@@ -322,6 +322,13 @@ impl ABIMachineSpec for S390xMachineDeps {
         }
 
         for mut param in params.into_iter().copied() {
+            if let ir::ArgumentPurpose::StructArgument(_) = param.purpose {
+                panic!(
+                    "StructArgument parameters are not supported on s390x. \
+                    Use regular pointer arguments instead."
+                );
+            }
+
             let intreg = in_int_reg(param.value_type);
             let fltreg = in_flt_reg(param.value_type);
             let vecreg = in_vec_reg(param.value_type);
@@ -393,22 +400,14 @@ impl ABIMachineSpec for S390xMachineDeps {
                 }
             };
 
-            if let ir::ArgumentPurpose::StructArgument(size) = param.purpose {
-                assert!(size % 8 == 0, "StructArgument size is not properly aligned");
-                args.push(ABIArg::StructArg {
-                    pointer: Some(slot),
-                    offset: 0,
-                    size: size as u64,
-                    purpose: param.purpose,
-                });
-            } else if let Some(ty) = implicit_ref {
+            if let Some(ty) = implicit_ref {
                 assert!(
                     (ty_bits(ty) / 8) % 8 == 0,
                     "implicit argument size is not properly aligned"
                 );
                 args.push(ABIArg::ImplicitPtrArg {
                     pointer: slot,
-                    offset: 0,
+                    offset: 0, // Will be filled in later
                     ty,
                     purpose: param.purpose,
                 });
@@ -447,13 +446,10 @@ impl ABIMachineSpec for S390xMachineDeps {
         };
 
         // After all arguments are in their well-defined location,
-        // allocate buffers for all StructArg or ImplicitPtrArg arguments.
+        // allocate buffers for all ImplicitPtrArg arguments.
         for arg in args.args_mut() {
             match arg {
-                ABIArg::StructArg { offset, size, .. } => {
-                    *offset = next_stack as i64;
-                    next_stack += *size as u32;
-                }
+                ABIArg::StructArg { .. } => unreachable!(),
                 ABIArg::ImplicitPtrArg { offset, ty, .. } => {
                     *offset = next_stack as i64;
                     next_stack += (ty_bits(*ty) / 8) as u32;
@@ -484,9 +480,7 @@ impl ABIMachineSpec for S390xMachineDeps {
                             }
                         }
                     }
-                    ABIArg::StructArg { offset, .. } => {
-                        *offset -= next_stack as i64;
-                    }
+                    ABIArg::StructArg { .. } => unreachable!(),
                     ABIArg::ImplicitPtrArg { offset, .. } => {
                         *offset -= next_stack as i64;
                     }
diff --git a/cranelift/codegen/src/isa/s390x/inst.isle b/cranelift/codegen/src/isa/s390x/inst.isle
index 36af4ba75..74022b0e5 100644
--- a/cranelift/codegen/src/isa/s390x/inst.isle
+++ b/cranelift/codegen/src/isa/s390x/inst.isle
@@ -2927,10 +2927,6 @@
 ;; Prepare a stack copy of a single (oversized) argument.
 (decl copy_to_buffer (MemArg ABIArg Value) InstOutput)
 (rule 2 (copy_to_buffer base (abi_arg_only_slot slot) _) (output_none))
-(rule 1 (copy_to_buffer base (abi_arg_struct_pointer _ offset size) val)
-      (let ((dst MemArg (memarg_offset base offset))
-            (src MemArg (memarg_reg_plus_off val 0 0 (memflags_trusted))))
-        (side_effect (memcpy dst src size))))
 (rule 0 (copy_to_buffer base (abi_arg_implicit_pointer _ offset ty)
                       val @ (value_type ty))
       (side_effect (arg_store ty val (memarg_offset base offset))))
@@ -2940,9 +2936,6 @@
 (decl copy_to_arg (CallArgListBuilder LaneOrder MemArg ABIArg Value) InstOutput)
 (rule 2 (copy_to_arg uses lo base (abi_arg_only_slot slot) val)
       (copy_reg_to_arg_slot uses lo base slot (prepare_arg_val slot val)))
-(rule 1 (copy_to_arg uses lo base (abi_arg_struct_pointer slot offset _) _)
-      (let ((ptr Reg (load_addr (memarg_offset base offset))))
-        (copy_reg_to_arg_slot uses lo base slot ptr)))
 (rule 0 (copy_to_arg uses lo base (abi_arg_implicit_pointer slot offset _) _)
       (let ((ptr Reg (load_addr (memarg_offset base offset))))
         (copy_reg_to_arg_slot uses lo base slot ptr)))
diff --git a/cranelift/codegen/src/isa/x64/abi.rs b/cranelift/codegen/src/isa/x64/abi.rs
index 3568a5b88..781cf1099 100644
--- a/cranelift/codegen/src/isa/x64/abi.rs
+++ b/cranelift/codegen/src/isa/x64/abi.rs
@@ -140,7 +140,6 @@ impl ABIMachineSpec for X64ABIMachineSpec {
                 assert!(size % 8 == 0, "StructArgument size is not properly aligned");
                 next_stack += size;
                 args.push(ABIArg::StructArg {
-                    pointer: None,
                     offset,
                     size: size as u64,
                     purpose: param.purpose,
diff --git a/cranelift/codegen/src/machinst/abi.rs b/cranelift/codegen/src/machinst/abi.rs
index e965d167b..b9752a947 100644
--- a/cranelift/codegen/src/machinst/abi.rs
+++ b/cranelift/codegen/src/machinst/abi.rs
@@ -205,10 +205,6 @@ pub enum ABIArg {
     /// area; on the callee side, we compute a pointer to this stack area and
     /// provide that as the argument's value.
     StructArg {
-        /// Register or stack slot holding a pointer to the buffer as passed
-        /// by the caller to the callee.  If None, the ABI defines the buffer
-        /// to reside at a well-known location (i.e. at `offset` below).
-        pointer: Option<ABIArgSlot>,
         /// Offset of this arg relative to base of stack args.
         offset: i64,
         /// Size of this arg on the stack.
@@ -1483,20 +1479,13 @@ impl<M: ABIMachineSpec> Callee<M> {
                     copy_arg_slot_to_reg(&slot, &into_reg);
                 }
             }
-            &ABIArg::StructArg {
-                pointer, offset, ..
-            } => {
+            &ABIArg::StructArg { offset, .. } => {
                 let into_reg = into_regs.only_reg().unwrap();
-                if let Some(slot) = pointer {
-                    // Buffer address is passed in a register or stack slot.
-                    copy_arg_slot_to_reg(&slot, &into_reg);
-                } else {
-                    // Buffer address is implicitly defined by the ABI.
-                    insts.push(M::gen_get_stack_addr(
-                        StackAMode::IncomingArg(offset, sigs[self.sig].sized_stack_arg_space),
-                        into_reg,
-                    ));
-                }
+                // Buffer address is implicitly defined by the ABI.
+                insts.push(M::gen_get_stack_addr(
+                    StackAMode::IncomingArg(offset, sigs[self.sig].sized_stack_arg_space),
+                    into_reg,
+                ));
             }
             &ABIArg::ImplicitPtrArg { pointer, ty, .. } => {
                 let into_reg = into_regs.only_reg().unwrap();
@@ -2199,8 +2188,8 @@ impl<M: ABIMachineSpec> CallSite<M> {
                     }
                 }
             }
-            ABIArg::StructArg { pointer, .. } => {
-                assert!(pointer.is_none()); // Only supported via ISLE.
+            ABIArg::StructArg { .. } => {
+                // Only supported via ISLE.
             }
             ABIArg::ImplicitPtrArg {
                 offset,
diff --git a/cranelift/codegen/src/machinst/isle.rs b/cranelift/codegen/src/machinst/isle.rs
index 7f12892a1..6954dab2a 100644
--- a/cranelift/codegen/src/machinst/isle.rs
+++ b/cranelift/codegen/src/machinst/isle.rs
@@ -491,24 +491,6 @@ macro_rules! isle_lower_prelude_methods {
             }
         }
 
-        fn abi_arg_struct_pointer(&mut self, arg: &ABIArg) -> Option<(ABIArgSlot, i64, u64)> {
-            match arg {
-                &ABIArg::StructArg {
-                    pointer,
-                    offset,
-                    size,
-                    ..
-                } => {
-                    if let Some(pointer) = pointer {
-                        Some((pointer, offset, size))
-                    } else {
-                        None
-                    }
-                }
-                _ => None,
-            }
-        }
-
         fn abi_arg_implicit_pointer(&mut self, arg: &ABIArg) -> Option<(ABIArgSlot, i64, Type)> {
             match arg {
                 &ABIArg::ImplicitPtrArg {
diff --git a/cranelift/codegen/src/prelude_lower.isle b/cranelift/codegen/src/prelude_lower.isle
index b9c16d68a..5d5a9a964 100644
--- a/cranelift/codegen/src/prelude_lower.isle
+++ b/cranelift/codegen/src/prelude_lower.isle
@@ -1030,11 +1030,6 @@
 (decl abi_arg_only_slot (ABIArgSlot) ABIArg)
 (extern extractor abi_arg_only_slot abi_arg_only_slot)
 
-;; Extractor to detect the special case where a struct argument
-;; is explicitly passed by reference using a hidden pointer.
-(decl abi_arg_struct_pointer (ABIArgSlot i64 u64) ABIArg)
-(extern extractor abi_arg_struct_pointer abi_arg_struct_pointer)
-
 ;; Extractor to detect the special case where a non-struct argument
 ;; is implicitly passed by reference using a hidden pointer.
 (decl abi_arg_implicit_pointer (ABIArgSlot i64 Type) ABIArg)
diff --git a/crates/cranelift/src/func_environ.rs b/crates/cranelift/src/func_environ.rs
index 180852eaf..4cc8449b8 100644
--- a/crates/cranelift/src/func_environ.rs
+++ b/crates/cranelift/src/func_environ.rs
@@ -1716,13 +1716,6 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         let table_data = self.tables[table_index].clone().unwrap();
         let heap_ty = table.ref_type.heap_type;
         match heap_ty.top() {
-            // `i31ref`s never need barriers, and therefore don't need to go
-            // through the GC compiler.
-            WasmHeapTopType::Any if heap_ty == WasmHeapType::I31 => {
-                let (src, flags) = table_data.prepare_table_addr(self, builder, index);
-                gc::unbarriered_load_gc_ref(self, builder, WasmHeapType::I31, src, flags)
-            }
-
             // GC-managed types.
             WasmHeapTopType::Any | WasmHeapTopType::Extern => {
                 let (src, flags) = table_data.prepare_table_addr(self, builder, index);
@@ -1762,13 +1755,6 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         let table_data = self.tables[table_index].clone().unwrap();
         let heap_ty = table.ref_type.heap_type;
         match heap_ty.top() {
-            // `i31ref`s never need GC barriers, and therefore don't need to go
-            // through the GC compiler.
-            WasmHeapTopType::Any if heap_ty == WasmHeapType::I31 => {
-                let (addr, flags) = table_data.prepare_table_addr(self, builder, index);
-                gc::unbarriered_store_gc_ref(self, builder, WasmHeapType::I31, addr, value, flags)
-            }
-
             // GC-managed types.
             WasmHeapTopType::Any | WasmHeapTopType::Extern => {
                 let (dst, flags) = table_data.prepare_table_addr(self, builder, index);
@@ -1920,17 +1906,13 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         let gv = builder.ins().global_value(self.pointer_type(), gv);
         let src = builder.ins().iadd_imm(gv, i64::from(offset));
 
-        if let WasmHeapType::I31 = ty.heap_type {
-            gc::unbarriered_load_gc_ref(self, builder, ty.heap_type, src, ir::MemFlags::trusted())
-        } else {
-            gc::gc_compiler(self).translate_read_gc_reference(
-                self,
-                builder,
-                ty,
-                src,
-                ir::MemFlags::trusted(),
-            )
-        }
+        gc::gc_compiler(self).translate_read_gc_reference(
+            self,
+            builder,
+            ty,
+            src,
+            ir::MemFlags::trusted(),
+        )
     }
 
     fn translate_custom_global_set(
@@ -1952,25 +1934,14 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
         let gv = builder.ins().global_value(self.pointer_type(), gv);
         let src = builder.ins().iadd_imm(gv, i64::from(offset));
 
-        if let WasmHeapType::I31 = ty.heap_type {
-            gc::unbarriered_store_gc_ref(
-                self,
-                builder,
-                ty.heap_type,
-                src,
-                value,
-                ir::MemFlags::trusted(),
-            )
-        } else {
-            gc::gc_compiler(self).translate_write_gc_reference(
-                self,
-                builder,
-                ty,
-                src,
-                value,
-                ir::MemFlags::trusted(),
-            )
-        }
+        gc::gc_compiler(self).translate_write_gc_reference(
+            self,
+            builder,
+            ty,
+            src,
+            value,
+            ir::MemFlags::trusted(),
+        )
     }
 
     fn make_heap(&mut self, func: &mut ir::Function, index: MemoryIndex) -> WasmResult<Heap> {
diff --git a/crates/cranelift/src/gc.rs b/crates/cranelift/src/gc.rs
index 8cc9f0d26..96aaedff7 100644
--- a/crates/cranelift/src/gc.rs
+++ b/crates/cranelift/src/gc.rs
@@ -26,39 +26,6 @@ pub fn gc_compiler(func_env: &FuncEnvironment<'_>) -> Box<dyn GcCompiler> {
     imp::gc_compiler(func_env)
 }
 
-/// Load a `*mut VMGcRef` into a virtual register, without any GC barriers.
-///
-/// The resulting value is an instance of the function environment's type for
-/// GC-managed references, aka `i32`. Note that a `VMGcRef` is always 4-bytes
-/// large, even when targeting 64-bit architectures.
-pub fn unbarriered_load_gc_ref(
-    func_env: &FuncEnvironment<'_>,
-    builder: &mut FunctionBuilder<'_>,
-    ty: WasmHeapType,
-    src: ir::Value,
-    flags: ir::MemFlags,
-) -> WasmResult<ir::Value> {
-    imp::unbarriered_load_gc_ref(func_env, builder, ty, src, flags)
-}
-
-/// Store `*dst = gc_ref`, without any GC barriers.
-///
-/// `dst` is a `*mut VMGcRef`.
-///
-/// `gc_ref` is an instance of the function environment's type for GC-managed
-/// references, aka `i32`. Note that a `VMGcRef` is always 4-bytes large, even
-/// when targeting 64-bit architectures.
-pub fn unbarriered_store_gc_ref(
-    func_env: &FuncEnvironment<'_>,
-    builder: &mut FunctionBuilder<'_>,
-    ty: WasmHeapType,
-    dst: ir::Value,
-    gc_ref: ir::Value,
-    flags: ir::MemFlags,
-) -> WasmResult<()> {
-    imp::unbarriered_store_gc_ref(func_env, builder, ty, dst, gc_ref, flags)
-}
-
 /// Get the index and signature of the built-in function for doing `table.grow`
 /// on GC reference tables.
 pub fn gc_ref_table_grow_builtin(
diff --git a/crates/cranelift/src/gc/disabled.rs b/crates/cranelift/src/gc/disabled.rs
index ef6847ad1..cc6250681 100644
--- a/crates/cranelift/src/gc/disabled.rs
+++ b/crates/cranelift/src/gc/disabled.rs
@@ -12,33 +12,6 @@ pub fn gc_compiler(_: &FuncEnvironment<'_>) -> Box<dyn GcCompiler> {
     Box::new(DisabledGcCompiler)
 }
 
-pub fn unbarriered_load_gc_ref(
-    _func_env: &FuncEnvironment<'_>,
-    _builder: &mut FunctionBuilder,
-    ty: WasmHeapType,
-    _ptr_to_gc_ref: ir::Value,
-    _flags: ir::MemFlags,
-) -> WasmResult<ir::Value> {
-    Err(wasm_unsupported!(
-        "support for `{ty}` references disabled at compile time because the `gc` cargo \
-         feature was not enabled"
-    ))
-}
-
-pub fn unbarriered_store_gc_ref(
-    _func_env: &FuncEnvironment<'_>,
-    _builder: &mut FunctionBuilder<'_>,
-    ty: WasmHeapType,
-    _dst: ir::Value,
-    _gc_ref: ir::Value,
-    _flags: ir::MemFlags,
-) -> WasmResult<()> {
-    Err(wasm_unsupported!(
-        "support for `{ty}` references disabled at compile time because the `gc` cargo \
-         feature was not enabled"
-    ))
-}
-
 pub fn gc_ref_table_grow_builtin(
     ty: WasmHeapType,
     _func_env: &mut FuncEnvironment<'_>,
diff --git a/crates/cranelift/src/gc/enabled.rs b/crates/cranelift/src/gc/enabled.rs
index 1d3fe0cf5..58de96489 100644
--- a/crates/cranelift/src/gc/enabled.rs
+++ b/crates/cranelift/src/gc/enabled.rs
@@ -12,34 +12,28 @@ pub fn gc_compiler(_func_env: &FuncEnvironment<'_>) -> Box<dyn GcCompiler> {
     Box::new(DrcCompiler::default())
 }
 
-pub fn unbarriered_load_gc_ref(
-    func_env: &FuncEnvironment<'_>,
+fn unbarriered_load_gc_ref(
     builder: &mut FunctionBuilder,
     ty: WasmHeapType,
     ptr_to_gc_ref: ir::Value,
     flags: ir::MemFlags,
 ) -> WasmResult<ir::Value> {
-    debug_assert!({
-        let (_, needs_stack_map) = func_env.reference_type(ty);
-        needs_stack_map
-    });
+    debug_assert!(ty.is_vmgcref_type());
     let gc_ref = builder.ins().load(ir::types::I32, flags, ptr_to_gc_ref, 0);
-    builder.declare_value_needs_stack_map(gc_ref);
+    if ty != WasmHeapType::I31 {
+        builder.declare_value_needs_stack_map(gc_ref);
+    }
     Ok(gc_ref)
 }
 
-pub fn unbarriered_store_gc_ref(
-    func_env: &FuncEnvironment<'_>,
+fn unbarriered_store_gc_ref(
     builder: &mut FunctionBuilder,
     ty: WasmHeapType,
     dst: ir::Value,
     gc_ref: ir::Value,
     flags: ir::MemFlags,
 ) -> WasmResult<()> {
-    debug_assert!({
-        let (_, needs_stack_map) = func_env.reference_type(ty);
-        needs_stack_map
-    });
+    debug_assert!(ty.is_vmgcref_type());
     builder.ins().store(flags, gc_ref, dst, 0);
     Ok(())
 }
@@ -293,7 +287,7 @@ impl GcCompiler for DrcCompiler {
         src: ir::Value,
         flags: ir::MemFlags,
     ) -> WasmResult<ir::Value> {
-        assert!(ty.is_vmgcref_type_and_not_i31());
+        assert!(ty.is_vmgcref_type());
 
         let (reference_type, needs_stack_map) = func_env.reference_type(ty.heap_type);
         debug_assert!(needs_stack_map);
@@ -315,6 +309,11 @@ impl GcCompiler for DrcCompiler {
             return Ok(null);
         };
 
+        // Special case for `i31` references: they don't need barriers.
+        if let WasmHeapType::I31 = ty.heap_type {
+            return unbarriered_load_gc_ref(builder, ty.heap_type, src, flags);
+        }
+
         // Our read barrier for GC references is roughly equivalent to the
         // following pseudo-CLIF:
         //
@@ -368,7 +367,7 @@ impl GcCompiler for DrcCompiler {
         builder.insert_block_after(continue_block, gc_block);
 
         // Load the GC reference and check for null/i31.
-        let gc_ref = unbarriered_load_gc_ref(func_env, builder, ty.heap_type, src, flags)?;
+        let gc_ref = unbarriered_load_gc_ref(builder, ty.heap_type, src, flags)?;
         let gc_ref_is_null_or_i31 = func_env.gc_ref_is_null_or_i31(builder, ty, gc_ref);
         builder.ins().brif(
             gc_ref_is_null_or_i31,
@@ -436,7 +435,7 @@ impl GcCompiler for DrcCompiler {
         new_val: ir::Value,
         flags: ir::MemFlags,
     ) -> WasmResult<()> {
-        assert!(ty.is_vmgcref_type_and_not_i31());
+        assert!(ty.is_vmgcref_type());
 
         let (ref_ty, needs_stack_map) = func_env.reference_type(ty.heap_type);
         debug_assert!(needs_stack_map);
@@ -460,6 +459,11 @@ impl GcCompiler for DrcCompiler {
             return Ok(());
         };
 
+        // Special case for `i31` references: they don't need barriers.
+        if let WasmHeapType::I31 = ty.heap_type {
+            return unbarriered_store_gc_ref(builder, ty.heap_type, dst, new_val, flags);
+        }
+
         // Our write barrier for GC references being copied out of the stack and
         // written into a table/global/etc... is roughly equivalent to the
         // following pseudo-CLIF:
@@ -542,7 +546,7 @@ impl GcCompiler for DrcCompiler {
 
         // Load the old value and then check whether the new value is non-null
         // and non-i31.
-        let old_val = unbarriered_load_gc_ref(func_env, builder, ty.heap_type, dst, flags)?;
+        let old_val = unbarriered_load_gc_ref(builder, ty.heap_type, dst, flags)?;
         let new_val_is_null_or_i31 = func_env.gc_ref_is_null_or_i31(builder, ty, new_val);
         builder.ins().brif(
             new_val_is_null_or_i31,
@@ -564,7 +568,7 @@ impl GcCompiler for DrcCompiler {
         // decremented.
         builder.switch_to_block(check_old_val_block);
         builder.seal_block(check_old_val_block);
-        unbarriered_store_gc_ref(func_env, builder, ty.heap_type, dst, new_val, flags)?;
+        unbarriered_store_gc_ref(builder, ty.heap_type, dst, new_val, flags)?;
         let old_val_is_null_or_i31 = func_env.gc_ref_is_null_or_i31(builder, ty, old_val);
         builder.ins().brif(
             old_val_is_null_or_i31,


Test patch:
diff --git a/cranelift/filetests/filetests/isa/s390x/struct-arg.clif b/cranelift/filetests/filetests/isa/s390x/struct-arg.clif
deleted file mode 100644
index b7b8f2097..000000000
--- a/cranelift/filetests/filetests/isa/s390x/struct-arg.clif
+++ /dev/null
@@ -1,203 +0,0 @@
-test compile precise-output
-target s390x
-
-function u0:0(i64 sarg(64)) -> i8 system_v {
-block0(v0: i64):
-    v1 = load.i8 v0
-    return v1
-}
-
-; VCode:
-; block0:
-;   llc %r2, 0(%r2)
-;   br %r14
-;
-; Disassembled:
-; block0: ; offset 0x0
-;   llc %r2, 0(%r2) ; trap: heap_oob
-;   br %r14
-
-function u0:1(i64 sarg(64), i64) -> i8 system_v {
-block0(v0: i64, v1: i64):
-    v2 = load.i8 v1
-    v3 = load.i8 v0
-    v4 = iadd.i8 v2, v3
-    return v4
-}
-
-; VCode:
-; block0:
-;   llc %r3, 0(%r3)
-;   llc %r4, 0(%r2)
-;   ark %r2, %r3, %r4
-;   br %r14
-;
-; Disassembled:
-; block0: ; offset 0x0
-;   llc %r3, 0(%r3) ; trap: heap_oob
-;   llc %r4, 0(%r2) ; trap: heap_oob
-;   ark %r2, %r3, %r4
-;   br %r14
-
-function u0:2(i64) -> i8 system_v {
-fn1 = colocated u0:0(i64 sarg(64)) -> i8 system_v
-
-block0(v0: i64):
-    v1 = call fn1(v0)
-    return v1
-}
-
-; VCode:
-;   stmg %r14, %r15, 112(%r15)
-;   aghi %r15, -224
-; block0:
-;   mvc 160(63,%r15), 0(%r2)
-;   la %r2, 160(%r15)
-;   brasl %r14, userextname0
-;   lmg %r14, %r15, 336(%r15)
-;   br %r14
-;
-; Disassembled:
-; block0: ; offset 0x0
-;   stmg %r14, %r15, 0x70(%r15)
-;   aghi %r15, -0xe0
-; block1: ; offset 0xa
-;   mvc 0xa0(0x40, %r15), 0(%r2)
-;   la %r2, 0xa0(%r15)
-;   brasl %r14, 0x14 ; reloc_external PLTRel32Dbl u0:0 2
-;   lmg %r14, %r15, 0x150(%r15)
-;   br %r14
-
-function u0:3(i64, i64) -> i8 system_v {
-fn1 = colocated u0:0(i64, i64 sarg(64)) -> i8 system_v
-
-block0(v0: i64, v1: i64):
-    v2 = call fn1(v0, v1)
-    return v2
-}
-
-; VCode:
-;   stmg %r14, %r15, 112(%r15)
-;   aghi %r15, -224
-; block0:
-;   mvc 160(63,%r15), 0(%r3)
-;   la %r3, 160(%r15)
-;   brasl %r14, userextname0
-;   lmg %r14, %r15, 336(%r15)
-;   br %r14
-;
-; Disassembled:
-; block0: ; offset 0x0
-;   stmg %r14, %r15, 0x70(%r15)
-;   aghi %r15, -0xe0
-; block1: ; offset 0xa
-;   mvc 0xa0(0x40, %r15), 0(%r3)
-;   la %r3, 0xa0(%r15)
-;   brasl %r14, 0x14 ; reloc_external PLTRel32Dbl u0:0 2
-;   lmg %r14, %r15, 0x150(%r15)
-;   br %r14
-
-function u0:4(i64 sarg(256), i64 sarg(64)) -> i8 system_v {
-block0(v0: i64, v1: i64):
-    v2 = load.i8 v0
-    v3 = load.i8 v1
-    v4 = iadd.i8 v2, v3
-    return v4
-}
-
-; VCode:
-; block0:
-;   lgr %r5, %r3
-;   llc %r3, 0(%r2)
-;   llc %r4, 0(%r5)
-;   ark %r2, %r3, %r4
-;   br %r14
-;
-; Disassembled:
-; block0: ; offset 0x0
-;   lgr %r5, %r3
-;   llc %r3, 0(%r2) ; trap: heap_oob
-;   llc %r4, 0(%r5) ; trap: heap_oob
-;   ark %r2, %r3, %r4
-;   br %r14
-
-function u0:5(i64, i64, i64) -> i8 system_v {
-fn1 = colocated u0:0(i64, i64 sarg(256), i64 sarg(64)) -> i8 system_v
-
-block0(v0: i64, v1: i64, v2: i64):
-    v3 = call fn1(v0, v1, v2)
-    return v3
-}
-
-; VCode:
-;   stmg %r14, %r15, 112(%r15)
-;   aghi %r15, -480
-; block0:
-;   mvc 160(255,%r15), 0(%r3)
-;   mvc 416(63,%r15), 0(%r4)
-;   la %r3, 160(%r15)
-;   la %r4, 416(%r15)
-;   brasl %r14, userextname0
-;   lmg %r14, %r15, 592(%r15)
-;   br %r14
-;
-; Disassembled:
-; block0: ; offset 0x0
-;   stmg %r14, %r15, 0x70(%r15)
-;   aghi %r15, -0x1e0
-; block1: ; offset 0xa
-;   mvc 0xa0(0x100, %r15), 0(%r3)
-;   mvc 0x1a0(0x40, %r15), 0(%r4)
-;   la %r3, 0xa0(%r15)
-;   la %r4, 0x1a0(%r15)
-;   brasl %r14, 0x1e ; reloc_external PLTRel32Dbl u0:0 2
-;   lmg %r14, %r15, 0x250(%r15)
-;   br %r14
-
-function u0:6(i64, i64, i64) -> i8 system_v {
-fn1 = colocated u0:0(i64, i64 sarg(1024), i64 sarg(64)) -> i8 system_v
-
-block0(v0: i64, v1: i64, v2: i64):
-    v3 = call fn1(v0, v1, v2)
-    return v3
-}
-
-; VCode:
-;   stmg %r8, %r15, 64(%r15)
-;   aghi %r15, -1248
-; block0:
-;   lgr %r8, %r2
-;   lgr %r10, %r4
-;   la %r2, 160(%r15)
-;   la %r3, 0(%r3)
-;   lghi %r4, 1024
-;   brasl %r14, %Memcpy
-;   lgr %r4, %r10
-;   mvc 1184(63,%r15), 0(%r4)
-;   la %r3, 160(%r15)
-;   la %r4, 1184(%r15)
-;   lgr %r2, %r8
-;   brasl %r14, userextname0
-;   lmg %r8, %r15, 1312(%r15)
-;   br %r14
-;
-; Disassembled:
-; block0: ; offset 0x0
-;   stmg %r8, %r15, 0x40(%r15)
-;   aghi %r15, -0x4e0
-; block1: ; offset 0xa
-;   lgr %r8, %r2
-;   lgr %r10, %r4
-;   la %r2, 0xa0(%r15)
-;   la %r3, 0(%r3)
-;   lghi %r4, 0x400
-;   brasl %r14, 0x1e ; reloc_external PLTRel32Dbl %Memcpy 2
-;   lgr %r4, %r10
-;   mvc 0x4a0(0x40, %r15), 0(%r4)
-;   la %r3, 0xa0(%r15)
-;   la %r4, 0x4a0(%r15)
-;   lgr %r2, %r8
-;   brasl %r14, 0x3a ; reloc_external PLTRel32Dbl u0:0 2
-;   lmg %r8, %r15, 0x520(%r15)
-;   br %r14
-