Fix StructReturn handling: properly mark the clobber, and offset actual rets. (#5023)

* Fix StructReturn handling: properly mark the clobber, and offset actual rets.

The legalization of `StructReturn` was causing issues in the new
call-handling code: the `StructReturn` ret was included in the `SigData` as
if it were an actual CLIF-level return value, but it is not.

Prior to using regalloc constraints for return values, we
unconditionally included rax (or the architecture's usual return
register) as a def, so it would be properly handled as "clobbered" by
the regalloc. With the new scheme, we include defs on the call only for
CLIF-level outputs. Callees with `StructReturn` args were thus not known
to clobber the return-value register, and values might be corrupted.

This PR updates the code to include a `StructReturn` ret as a clobber
rather than a returned value in the relevant spots. I observed it
causing saves/restores of rax in some CLIF that @bjorn3 provided me, but
I was having difficulty minimizing this into a test-case that I would be
comfortable including as a precise-output case (including the whole
thing verbatim would lock down a bunch of other irrelevant details and
cause test-update noise later). If we can find a more minimized example
I'm happy to include it as a filetest.

Fixes #5018.

cranelift/codegen/src/machinst/abi.rs

@@ -561,7 +561,7 @@ pub struct Sig(u32);
 cranelift_entity::entity_impl!(Sig);
 
 /// ABI information shared between body (callee) and caller.
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct SigData {
     /// Argument locations (regs or stack slots). Stack offsets are relative to
     /// SP on entry to function.
@@ -669,10 +669,17 @@ impl SigData {
         // regs, which we will remove from the clobber set below.
         let mut clobbers = M::get_regs_clobbered_by_call(self.call_conv);
 
-        // Compute defs: all retval regs, and all caller-save (clobbered) regs.
+        // Compute defs: all retval regs, and all caller-save
+        // (clobbered) regs, except for StructRet args.
         let mut defs = smallvec![];
         for ret in &self.rets {
-            if let &ABIArg::Slots { ref slots, .. } = ret {
+            if let &ABIArg::Slots {
+                ref slots, purpose, ..
+            } = ret
+            {
+                if purpose == ir::ArgumentPurpose::StructReturn {
+                    continue;
+                }
                 for slot in slots {
                     match slot {
                         &ABIArgSlot::Reg { reg, .. } => {
@@ -694,9 +701,17 @@ impl SigData {
         // regs, which we will remove from the clobber set below.
         let mut clobbers = M::get_regs_clobbered_by_call(self.call_conv);
 
-        // Remove retval regs from clobbers.
+        // Remove retval regs from clobbers. Skip StructRets: these
+        // are not, semantically, returns at the CLIF level, so we
+        // treat such a value as a clobber instead.
         for ret in &self.rets {
-            if let &ABIArg::Slots { ref slots, .. } = ret {
+            if let &ABIArg::Slots {
+                ref slots, purpose, ..
+            } = ret
+            {
+                if purpose == ir::ArgumentPurpose::StructReturn {
+                    continue;
+                }
                 for slot in slots {
                     match slot {
                         &ABIArgSlot::Reg { reg, .. } => {

cranelift/codegen/src/machinst/isle.rs

@@ -1206,13 +1206,22 @@ macro_rules! isle_prelude_method_helpers {
                     self.lower_ctx.emit(inst);
                 }
             }
+
             // Handle retvals prior to emitting call, so the
             // constraints are on the call instruction; but buffer the
             // instructions till after the call.
             let mut outputs = InstOutput::new();
             let mut retval_insts: crate::machinst::abi::SmallInstVec<_> = smallvec::smallvec![];
-            for i in 0..num_rets {
-                let ret = self.lower_ctx.sigs()[abi].get_ret(i);
+            // We take the *last* `num_rets` returns of the sig:
+            // this skips a StructReturn, if any, that is present.
+            let sigdata = &self.lower_ctx.sigs()[abi];
+            debug_assert!(num_rets <= sigdata.num_rets());
+            for i in (sigdata.num_rets() - num_rets)..sigdata.num_rets() {
+                // Borrow `sigdata` again so we don't hold a `self`
+                // borrow across the `&mut self` arg to
+                // `abi_arg_slot_regs()` below.
+                let sigdata = &self.lower_ctx.sigs()[abi];
+                let ret = sigdata.get_ret(i);
                 let retval_regs = self.abi_arg_slot_regs(&ret).unwrap();
                 retval_insts.extend(
                     caller

cranelift/filetests/filetests/isa/aarch64/call.clif

@@ -475,7 +475,6 @@ block0(v0: i64):
 ;   mov x8, x0
 ;   ldr x4, 8 ; b 12 ; data TestCase(%g) + 0
 ;   blr x4
-;   mov x0, x8
 ;   ldp fp, lr, [sp], #16
 ;   ret
 

cranelift/filetests/filetests/isa/x64/struct-ret.clif

@@ -33,6 +33,7 @@ block0(v0: i64, v1: i64):
 ;   movq    %rsi, %rdi
 ;   load_ext_name %f2+0, %r8
 ;   call    *%r8
+;   movq    %rdx, %rax
 ;   movq    %rbp, %rsp
 ;   popq    %rbp
 ;   ret
@@ -47,10 +48,15 @@ block0(v0: i64):
 
 ;   pushq   %rbp
 ;   movq    %rsp, %rbp
+;   subq    %rsp, $16, %rsp
+;   movq    %r15, 0(%rsp)
 ; block0:
-;   movq    %rdi, %rax
+;   movq    %rdi, %r15
 ;   load_ext_name %f4+0, %rdx
 ;   call    *%rdx
+;   movq    %r15, %rax
+;   movq    0(%rsp), %r15
+;   addq    %rsp, $16, %rsp
 ;   movq    %rbp, %rsp
 ;   popq    %rbp
 ;   ret