From 097d1087e02e5e4b6a03289aae760948ae55de69 Mon Sep 17 00:00:00 2001
From: Nick Fitzgerald <fitzgen@gmail.com>
Date: Mon, 24 Oct 2022 17:21:34 -0700
Subject: [PATCH] Cranelift: Avoid calling
 `ensure_struct_return_pointer_is_returned` and cloning sigs for every call
 (#5113)

* Cranelift: pass iterators to `ABIMachineSpec::compute_arg_locs`

Instead of slices. This gives us more flexibility to pass custom sequences
without needing to allocate a `Vec` to hold them and pass in as a slice.

* Cranelift: Avoid cloning `ir::Signature`s in `SigData::from_func_sig`

This avoids two heap allocations per signature that are unnecessary 99% of the
time.

* fix typo

* Simplify condition in `missing_struct_return`
---
 cranelift/codegen/src/machinst/abi.rs | 35 ++++++++++++++++-----------
 1 file changed, 21 insertions(+), 14 deletions(-)

diff --git a/cranelift/codegen/src/machinst/abi.rs b/cranelift/codegen/src/machinst/abi.rs
index d6d129cc77..2ff928b81b 100644
--- a/cranelift/codegen/src/machinst/abi.rs
+++ b/cranelift/codegen/src/machinst/abi.rs
@@ -586,14 +586,15 @@ impl SigData {
         sig: &ir::Signature,
         flags: &settings::Flags,
     ) -> CodegenResult<SigData> {
-        let sig = ensure_struct_return_ptr_is_returned(sig);
+        let sret = missing_struct_return(sig);
+        let returns = sret.as_ref().into_iter().chain(&sig.returns);
 
         // Compute args and retvals from signature. Handle retvals first,
         // because we may need to add a return-area arg to the args.
         let (rets, sized_stack_ret_space, _) = M::compute_arg_locs(
             sig.call_conv,
             flags,
-            &sig.returns,
+            returns,
             ArgsOrRets::Rets,
             /* extra ret-area ptr = */ false,
         )?;
@@ -1175,19 +1176,21 @@ fn gen_store_stack_multi<M: ABIMachineSpec>(
     ret
 }
 
-pub(crate) fn ensure_struct_return_ptr_is_returned(sig: &ir::Signature) -> ir::Signature {
-    let params_structret = sig
-        .params
-        .iter()
-        .find(|p| p.purpose == ArgumentPurpose::StructReturn);
-    let rets_have_structret = sig.returns.len() > 0
-        && sig
-            .returns
-            .iter()
-            .any(|arg| arg.purpose == ArgumentPurpose::StructReturn);
+/// If the signature needs to be legalized, then return the struct-return
+/// parameter that should be prepended to its returns. Otherwise, return `None`.
+fn missing_struct_return(sig: &ir::Signature) -> Option<ir::AbiParam> {
+    let struct_ret_index = sig.special_param_index(ArgumentPurpose::StructReturn)?;
+    if !sig.uses_special_return(ArgumentPurpose::StructReturn) {
+        return Some(sig.params[struct_ret_index]);
+    }
+
+    None
+}
+
+fn ensure_struct_return_ptr_is_returned(sig: &ir::Signature) -> ir::Signature {
     let mut sig = sig.clone();
-    if params_structret.is_some() && !rets_have_structret {
-        sig.returns.insert(0, params_structret.unwrap().clone());
+    if let Some(sret) = missing_struct_return(&sig) {
+        sig.returns.insert(0, sret);
     }
     sig
 }
@@ -1198,6 +1201,10 @@ pub(crate) fn ensure_struct_return_ptr_is_returned(sig: &ir::Signature) -> ir::S
 impl<M: ABIMachineSpec> Callee<M> {
     /// Access the (possibly legalized) signature.
     pub fn signature(&self) -> &ir::Signature {
+        debug_assert!(
+            missing_struct_return(&self.ir_sig).is_none(),
+            "`Callee::ir_sig` is always legalized"
+        );
         &self.ir_sig
     }