Fix sret for AArch64 (#4634)

* Fix sret for AArch64

AArch64 requires the struct return address argument to be stored in the x8
register. This register is never used for regular arguments.

* Add extra sret tests for x86_64

cranelift/codegen/src/isa/aarch64/abi.rs

@@ -143,6 +143,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
             let (rcs, reg_types) = Inst::rc_for_type(param.value_type)?;
 
             if let ir::ArgumentPurpose::StructArgument(size) = param.purpose {
+                assert_eq!(args_or_rets, ArgsOrRets::Args);
                 let offset = next_stack as i64;
                 let size = size as u64;
                 assert!(size % 8 == 0, "StructArgument size is not properly aligned");
@@ -156,6 +157,24 @@ impl ABIMachineSpec for AArch64MachineDeps {
                 continue;
             }
 
+            if let ir::ArgumentPurpose::StructReturn = param.purpose {
+                // FIXME add assert_eq!(args_or_rets, ArgsOrRets::Args); once
+                // ensure_struct_return_ptr_is_returned is gone.
+                assert!(
+                    param.value_type == types::I64,
+                    "StructReturn must be a pointer sized integer"
+                );
+                ret.push(ABIArg::Slots {
+                    slots: smallvec![ABIArgSlot::Reg {
+                        reg: xreg(8).to_real_reg().unwrap(),
+                        ty: types::I64,
+                        extension: param.extension,
+                    },],
+                    purpose: ir::ArgumentPurpose::StructReturn,
+                });
+                continue;
+            }
+
             // Handle multi register params
             //
             // See AArch64 ABI (https://github.com/ARM-software/abi-aa/blob/2021Q1/aapcs64/aapcs64.rst#642parameter-passing-rules), (Section 6.4.2 Stage C).

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

@@ -452,3 +452,54 @@ block0:
 ;   str w7, [x11]
 ;   ret
 
+function %f17(i64 sret) {
+block0(v0: i64):
+    v1 = iconst.i64 42
+    store v1, v0
+    return
+}
+
+; block0:
+;   mov x5, x8
+;   movz x4, #42
+;   str x4, [x8]
+;   ret
+
+function %f18(i64) -> i64 {
+    fn0 = %g(i64 sret) -> i64
+
+block0(v0: i64):
+    v1 = call fn0(v0)
+    return v1
+}
+
+;   stp fp, lr, [sp, #-16]!
+;   mov fp, sp
+; block0:
+;   mov x8, x0
+;   ldr x5, 8 ; b 12 ; data TestCase { length: 1, ascii: [103, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] } + 0
+;   blr x5
+;   mov x0, x8
+;   ldp fp, lr, [sp], #16
+;   ret
+
+function %f18(i64 sret) {
+    fn0 = %g(i64 sret)
+
+block0(v0: i64):
+    call fn0(v0)
+    return
+}
+
+;   stp fp, lr, [sp, #-16]!
+;   mov fp, sp
+;   str x24, [sp, #-16]!
+; block0:
+;   mov x24, x8
+;   ldr x5, 8 ; b 12 ; data TestCase { length: 1, ascii: [103, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] } + 0
+;   blr x5
+;   mov x8, x24
+;   ldr x24, [sp], #16
+;   ldp fp, lr, [sp], #16
+;   ret
+

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

@@ -18,3 +18,45 @@ block0(v0: i64):
 ;   popq    %rbp
 ;   ret
 
+
+function %f1(i64, i64) -> i64 {
+    fn0 = %f2(i64 sret) -> i64
+
+block0(v0: i64, v1: i64):
+    v2 = call fn0(v1)
+    return v2
+}
+
+;   pushq   %rbp
+;   movq    %rsp, %rbp
+; block0:
+;   movq    %rsi, %rdi
+;   load_ext_name %f2+0, %r9
+;   call    *%r9
+;   movq    %rbp, %rsp
+;   popq    %rbp
+;   ret
+
+function %f3(i64 sret) {
+    fn0 = %f4(i64 sret)
+
+block0(v0: i64):
+    call fn0(v0)
+    return
+}
+
+;   pushq   %rbp
+;   movq    %rsp, %rbp
+;   subq    %rsp, $16, %rsp
+;   movq    %r15, 0(%rsp)
+; block0:
+;   movq    %rdi, %r15
+;   load_ext_name %f4+0, %r8
+;   call    *%r8
+;   movq    %r15, %rax
+;   movq    0(%rsp), %r15
+;   addq    %rsp, $16, %rsp
+;   movq    %rbp, %rsp
+;   popq    %rbp
+;   ret
+