Fix varargs support on aarch64-apple-darwin

Author: beetrees

src/abi/mod.rs

@@ -5,8 +5,9 @@ mod pass_mode;
 mod returning;
 
 use std::borrow::Cow;
+use std::mem;
 
-use cranelift_codegen::ir::SigRef;
+use cranelift_codegen::ir::{ArgumentPurpose, SigRef};
 use cranelift_codegen::isa::CallConv;
 use cranelift_module::ModuleError;
 use rustc_codegen_ssa::errors::CompilerBuiltinsCannotCall;
@@ -17,6 +18,7 @@ use rustc_middle::ty::TypeVisitableExt;
 use rustc_monomorphize::is_call_from_compiler_builtins_to_upstream_monomorphization;
 use rustc_session::Session;
 use rustc_span::source_map::Spanned;
+use rustc_span::sym::assert;
 use rustc_target::abi::call::{Conv, FnAbi};
 use rustc_target::spec::abi::Abi;
 
@@ -425,7 +427,8 @@ pub(crate) fn codegen_terminator_call<'tcx>(
         None
     };
 
-    let extra_args = &args[fn_sig.inputs().skip_binder().len()..];
+    let fixed_arg_count = fn_sig.inputs().skip_binder().len();
+    let extra_args = &args[fixed_arg_count..];
     let extra_args = fx.tcx.mk_type_list_from_iter(
         extra_args.iter().map(|op_arg| fx.monomorphize(op_arg.node.ty(fx.mir, fx.tcx))),
     );
@@ -532,18 +535,59 @@ pub(crate) fn codegen_terminator_call<'tcx>(
     };
 
     self::returning::codegen_with_call_return_arg(fx, &fn_abi.ret, ret_place, |fx, return_ptr| {
-        let call_args = return_ptr
+        // The majority of this block is hacky workarounds for cranelift not supporting C variadic
+        // functions. The only bits that aren't are the `let mut call_args =` statement (excluding
+        // the `.take()`; `(&mut args_iter)` can be replace with `args.into_iter().enumerate()`)
+        // and the `let call_inst = ` statement.
+
+        let has_return_arg = return_ptr.is_some();
+        let mut args_iter = args.into_iter().enumerate();
+        let mut call_args = return_ptr
             .into_iter()
             .chain(first_arg_override.into_iter())
             .chain(
-                args.into_iter()
-                    .enumerate()
+                (&mut args_iter)
+                    // If the function is variadic, only add the fixed argument first so that how
+                    // many `Value`s the fixed arguments expand to can be counted.
+                    .take(if fn_sig.c_variadic() { fixed_arg_count } else { usize::MAX })
                     .skip(if first_arg_override.is_some() { 1 } else { 0 })
                     .flat_map(|(i, arg)| {
                         adjust_arg_for_abi(fx, arg.value, &fn_abi.args[i], arg.is_owned).into_iter()
                     }),
             )
             .collect::<Vec<Value>>();
+        // Count how many `Value`s represent the fixed arguments.
+        let fixed_arg_count = call_args.len();
+        // Add the rest of the arguments. This is a no-op unless the function is variadic.
+        call_args.extend(args_iter.flat_map(|(i, arg)| {
+            adjust_arg_for_abi(fx, arg.value, &fn_abi.args[i], arg.is_owned).into_iter()
+        }));
+
+        if fx.tcx.sess.target.is_like_osx
+            && fx.tcx.sess.target.arch == "aarch64"
+            && fn_sig.c_variadic()
+            // No need to pad the argument list unless variadic arguments are actually being passed.
+            && call_args.len() > fixed_arg_count
+        {
+            // 128-bit integers take 2 registers, and everything else takes 1.
+            // FIXME: Add support for non-integer types
+            // This relies on the checks below to ensure all arguments are integer types and that
+            // the ABI is "C".
+            // The return argument isn't counted as it goes in its own dedicated register.
+            let integer_registers_used: usize = call_args
+                [if has_return_arg { 1 } else { 0 }..fixed_arg_count]
+                .iter()
+                .map(|arg| if fx.bcx.func.dfg.value_type(*arg).bits() == 128 { 2 } else { 1 })
+                .sum();
+            // The ABI uses 8 registers before it starts pushing arguments to the stack. Pad out
+            // the registers if needed to ensure the variadic arguments are passed on the stack.
+            if integer_registers_used < 8 {
+                call_args.splice(
+                    fixed_arg_count..fixed_arg_count,
+                    (integer_registers_used..8).map(|_| fx.bcx.ins().iconst(types::I64, 0)),
+                );
+            }
+        }
 
         let call_inst = match func_ref {
             CallTarget::Direct(func_ref) => fx.bcx.ins().call(func_ref, &call_args),
@@ -561,20 +605,36 @@ pub(crate) fn codegen_terminator_call<'tcx>(
                 );
             }
             let sig_ref = fx.bcx.func.dfg.call_signature(call_inst).unwrap();
-            let abi_params = call_args
-                .into_iter()
-                .map(|arg| {
-                    let ty = fx.bcx.func.dfg.value_type(arg);
-                    if !ty.is_int() {
-                        // FIXME set %al to upperbound on float args once floats are supported
-                        fx.tcx.dcx().span_fatal(
-                            source_info.span,
-                            format!("Non int ty {:?} for variadic call", ty),
-                        );
-                    }
-                    AbiParam::new(ty)
-                })
-                .collect::<Vec<AbiParam>>();
+            let mut abi_params = mem::take(&mut fx.bcx.func.dfg.signatures[sig_ref].params);
+            assert_eq!(abi_params.len(), fixed_arg_count);
+            if fx.tcx.sess.target.is_like_osx && fx.tcx.sess.target.arch == "aarch64" {
+                // `StructArgument` is not currently used by the `aarch64`, and is therefore not
+                // handled when calculating how many padding arguments to use. Assert that this
+                // remains the case.
+                assert!(abi_params.iter().all(|param| matches!(
+                    param.purpose,
+                    // The only purposes used are `Normal` and `StructReturn`.
+                    ArgumentPurpose::Normal | ArgumentPurpose::StructReturn
+                )))
+            }
+            // Don't replace existing `AbiParam`s as they might have `ArgumentPurpose` or
+            // `ArgumentExtension`.
+            abi_params.extend(
+                call_args
+                    .into_iter()
+                    .map(|arg| {
+                        let ty = fx.bcx.func.dfg.value_type(arg);
+                        if !ty.is_int() {
+                            // FIXME set %al to upperbound on float args once floats are supported
+                            fx.tcx.dcx().span_fatal(
+                                source_info.span,
+                                format!("Non int ty {:?} for variadic call", ty),
+                            );
+                        }
+                        AbiParam::new(ty)
+                    })
+                    .skip(fixed_arg_count),
+            );
             fx.bcx.func.dfg.signatures[sig_ref].params = abi_params;
         }