Cranelift: Use a single, shared vector allocation for all ABIArgs (#5127)

* Cranelift: Use a single, shared vector allocation for all `ABIArg`s

Instead of two `SmallVec`s per `SigData`.

* Remove `Deref` and `DerefMut` impls for `ArgsAccumulator`

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

@@ -93,7 +93,8 @@ impl ABIMachineSpec for AArch64MachineDeps {
         params: I,
         args_or_rets: ArgsOrRets,
         add_ret_area_ptr: bool,
-    ) -> CodegenResult<(ABIArgVec, i64, Option<usize>)>
+        mut args: ArgsAccumulator<'_>,
+    ) -> CodegenResult<(i64, Option<usize>)>
     where
         I: IntoIterator<Item = &'a ir::AbiParam>,
     {
@@ -116,7 +117,6 @@ impl ABIMachineSpec for AArch64MachineDeps {
         let mut next_xreg = 0;
         let mut next_vreg = 0;
         let mut next_stack: u64 = 0;
-        let mut ret = ABIArgVec::new();
 
         let (max_per_class_reg_vals, mut remaining_reg_vals) = match args_or_rets {
             ArgsOrRets::Args => (8, 16), // x0-x7 and v0-v7
@@ -155,7 +155,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
                 let size = size as u64;
                 assert!(size % 8 == 0, "StructArgument size is not properly aligned");
                 next_stack += size;
-                ret.push(ABIArg::StructArg {
+                args.push(ABIArg::StructArg {
                     pointer: None,
                     offset,
                     size,
@@ -171,7 +171,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
                     param.value_type == types::I64,
                     "StructReturn must be a pointer sized integer"
                 );
-                ret.push(ABIArg::Slots {
+                args.push(ABIArg::Slots {
                     slots: smallvec![ABIArgSlot::Reg {
                         reg: xreg(8).to_real_reg().unwrap(),
                         ty: types::I64,
@@ -228,7 +228,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
                     let lower_reg = xreg(next_xreg);
                     let upper_reg = xreg(next_xreg + 1);
 
-                    ret.push(ABIArg::Slots {
+                    args.push(ABIArg::Slots {
                         slots: smallvec![
                             ABIArgSlot::Reg {
                                 reg: lower_reg.to_real_reg().unwrap(),
@@ -267,7 +267,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
                     } else {
                         param.value_type
                     };
-                    ret.push(ABIArg::reg(
+                    args.push(ABIArg::reg(
                         reg.to_real_reg().unwrap(),
                         ty,
                         param.extension,
@@ -319,7 +319,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
                 })
                 .collect();
 
-            ret.push(ABIArg::Slots {
+            args.push(ABIArg::Slots {
                 slots,
                 purpose: param.purpose,
             });
@@ -330,14 +330,14 @@ impl ABIMachineSpec for AArch64MachineDeps {
         let extra_arg = if add_ret_area_ptr {
             debug_assert!(args_or_rets == ArgsOrRets::Args);
             if next_xreg < max_per_class_reg_vals && remaining_reg_vals > 0 {
-                ret.push(ABIArg::reg(
+                args.push(ABIArg::reg(
                     xreg(next_xreg).to_real_reg().unwrap(),
                     I64,
                     ir::ArgumentExtension::None,
                     ir::ArgumentPurpose::Normal,
                 ));
             } else {
-                ret.push(ABIArg::stack(
+                args.push(ABIArg::stack(
                     next_stack as i64,
                     I64,
                     ir::ArgumentExtension::None,
@@ -345,7 +345,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
                 ));
                 next_stack += 8;
             }
-            Some(ret.len() - 1)
+            Some(args.args().len() - 1)
         } else {
             None
         };
@@ -358,7 +358,7 @@ impl ABIMachineSpec for AArch64MachineDeps {
             return Err(CodegenError::ImplLimitExceeded);
         }
 
-        Ok((ret, next_stack as i64, extra_arg))
+        Ok((next_stack as i64, extra_arg))
     }
 
     fn fp_to_arg_offset(_call_conv: isa::CallConv, _flags: &settings::Flags) -> i64 {

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

@@ -62,7 +62,8 @@ impl ABIMachineSpec for Riscv64MachineDeps {
         params: I,
         args_or_rets: ArgsOrRets,
         add_ret_area_ptr: bool,
-    ) -> CodegenResult<(ABIArgVec, i64, Option<usize>)>
+        mut args: ArgsAccumulator<'_>,
+    ) -> CodegenResult<(i64, Option<usize>)>
     where
         I: IntoIterator<Item = &'a ir::AbiParam>,
     {
@@ -78,7 +79,6 @@ impl ABIMachineSpec for Riscv64MachineDeps {
         let mut next_f_reg = f_start;
         // Stack space.
         let mut next_stack: u64 = 0;
-        let mut ret = smallvec![];
         let mut return_one_register_used = false;
 
         for param in params {
@@ -86,7 +86,7 @@ impl ABIMachineSpec for Riscv64MachineDeps {
                 let offset = next_stack;
                 assert!(size % 8 == 0, "StructArgument size is not properly aligned");
                 next_stack += size as u64;
-                ret.push(ABIArg::StructArg {
+                args.push(ABIArg::StructArg {
                     pointer: None,
                     offset: offset as i64,
                     size: size as u64,
@@ -152,7 +152,7 @@ impl ABIMachineSpec for Riscv64MachineDeps {
                     next_stack += size;
                 }
             }
-            ret.push(ABIArg::Slots {
+            args.push(ABIArg::Slots {
                 slots,
                 purpose: param.purpose,
             });
@@ -166,7 +166,7 @@ impl ABIMachineSpec for Riscv64MachineDeps {
                     ir::ArgumentExtension::None,
                     ir::ArgumentPurpose::Normal,
                 );
-                ret.push(arg);
+                args.push(arg);
             } else {
                 let arg = ABIArg::stack(
                     next_stack as i64,
@@ -174,10 +174,10 @@ impl ABIMachineSpec for Riscv64MachineDeps {
                     ir::ArgumentExtension::None,
                     ir::ArgumentPurpose::Normal,
                 );
-                ret.push(arg);
+                args.push(arg);
                 next_stack += 8;
             }
-            Some(ret.len() - 1)
+            Some(args.args().len() - 1)
         } else {
             None
         };
@@ -187,7 +187,7 @@ impl ABIMachineSpec for Riscv64MachineDeps {
         if next_stack > STACK_ARG_RET_SIZE_LIMIT {
             return Err(CodegenError::ImplLimitExceeded);
         }
-        CodegenResult::Ok((ret, next_stack as i64, pos))
+        CodegenResult::Ok((next_stack as i64, pos))
     }
 
     fn fp_to_arg_offset(_call_conv: isa::CallConv, _flags: &settings::Flags) -> i64 {

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

@@ -227,7 +227,8 @@ impl ABIMachineSpec for S390xMachineDeps {
         params: I,
         args_or_rets: ArgsOrRets,
         add_ret_area_ptr: bool,
-    ) -> CodegenResult<(ABIArgVec, i64, Option<usize>)>
+        mut args: ArgsAccumulator<'_>,
+    ) -> CodegenResult<(i64, Option<usize>)>
     where
         I: IntoIterator<Item = &'a ir::AbiParam>,
     {
@@ -235,7 +236,6 @@ impl ABIMachineSpec for S390xMachineDeps {
         let mut next_fpr = 0;
         let mut next_vr = 0;
         let mut next_stack: u64 = 0;
-        let mut ret = ABIArgVec::new();
 
         if args_or_rets == ArgsOrRets::Args {
             next_stack = REG_SAVE_AREA_SIZE as u64;
@@ -338,7 +338,7 @@ impl ABIMachineSpec for S390xMachineDeps {
 
             if let ir::ArgumentPurpose::StructArgument(size) = param.purpose {
                 assert!(size % 8 == 0, "StructArgument size is not properly aligned");
-                ret.push(ABIArg::StructArg {
+                args.push(ABIArg::StructArg {
                     pointer: Some(slot),
                     offset: 0,
                     size: size as u64,
@@ -349,14 +349,14 @@ impl ABIMachineSpec for S390xMachineDeps {
                     (ty_bits(ty) / 8) % 8 == 0,
                     "implicit argument size is not properly aligned"
                 );
-                ret.push(ABIArg::ImplicitPtrArg {
+                args.push(ABIArg::ImplicitPtrArg {
                     pointer: slot,
                     offset: 0,
                     ty,
                     purpose: param.purpose,
                 });
             } else {
-                ret.push(ABIArg::Slots {
+                args.push(ABIArg::Slots {
                     slots: smallvec![slot],
                     purpose: param.purpose,
                 });
@@ -375,14 +375,14 @@ impl ABIMachineSpec for S390xMachineDeps {
                 0
             };
             if let Some(reg) = get_intreg_for_arg(next_gpr) {
-                ret.push(ABIArg::reg(
+                args.push(ABIArg::reg(
                     reg.to_real_reg().unwrap(),
                     types::I64,
                     ir::ArgumentExtension::None,
                     ir::ArgumentPurpose::Normal,
                 ));
             } else {
-                ret.push(ABIArg::stack(
+                args.push(ABIArg::stack(
                     next_stack as i64,
                     types::I64,
                     ir::ArgumentExtension::None,
@@ -390,15 +390,15 @@ impl ABIMachineSpec for S390xMachineDeps {
                 ));
                 next_stack += 8;
             }
-            Some(ret.len() - 1)
+            Some(args.args().len() - 1)
         } else {
             None
         };
 
         // After all arguments are in their well-defined location,
         // allocate buffers for all StructArg or ImplicitPtrArg arguments.
-        for i in 0..ret.len() {
+        for i in 0..args.args().len() {
-            match &mut ret[i] {
+            match &mut args.args_mut()[i] {
                 &mut ABIArg::StructArg {
                     ref mut offset,
                     size,
@@ -423,7 +423,7 @@ impl ABIMachineSpec for S390xMachineDeps {
             return Err(CodegenError::ImplLimitExceeded);
         }
 
-        Ok((ret, next_stack as i64, extra_arg))
+        Ok((next_stack as i64, extra_arg))
     }
 
     fn fp_to_arg_offset(_call_conv: isa::CallConv, _flags: &settings::Flags) -> i64 {

cranelift/codegen/src/isa/s390x/lower/isle.rs

@@ -129,7 +129,7 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Flags, IsaFlags, 6>
         for i in 0..self.lower_ctx.sigs()[*abi].num_rets() {
             if let &ABIArg::Slots {
                 ref slots, purpose, ..
-            } = &self.lower_ctx.sigs()[*abi].get_ret(i)
+            } = &self.lower_ctx.sigs()[*abi].get_ret(self.lower_ctx.sigs(), i)
             {
                 if purpose == ArgumentPurpose::StructReturn {
                     continue;
@@ -192,7 +192,8 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Flags, IsaFlags, 6>
         defs: &CallRetList,
         opcode: &Opcode,
     ) -> BoxCallInfo {
-        let clobbers = self.lower_ctx.sigs()[*abi].call_clobbers::<S390xMachineDeps>();
+        let clobbers =
+            self.lower_ctx.sigs()[*abi].call_clobbers::<S390xMachineDeps>(self.lower_ctx.sigs());
         Box::new(CallInfo {
             dest: name.clone(),
             uses: uses.clone(),
@@ -213,7 +214,8 @@ impl generated_code::Context for IsleContext<'_, '_, MInst, Flags, IsaFlags, 6>
         defs: &CallRetList,
         opcode: &Opcode,
     ) -> BoxCallIndInfo {
-        let clobbers = self.lower_ctx.sigs()[*abi].call_clobbers::<S390xMachineDeps>();
+        let clobbers =
+            self.lower_ctx.sigs()[*abi].call_clobbers::<S390xMachineDeps>(self.lower_ctx.sigs());
         Box::new(CallIndInfo {
             rn: target,
             uses: uses.clone(),

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

@@ -87,7 +87,8 @@ impl ABIMachineSpec for X64ABIMachineSpec {
         params: I,
         args_or_rets: ArgsOrRets,
         add_ret_area_ptr: bool,
-    ) -> CodegenResult<(ABIArgVec, i64, Option<usize>)>
+        mut args: ArgsAccumulator<'_>,
+    ) -> CodegenResult<(i64, Option<usize>)>
     where
         I: IntoIterator<Item = &'a ir::AbiParam>,
     {
@@ -97,7 +98,6 @@ impl ABIMachineSpec for X64ABIMachineSpec {
         let mut next_vreg = 0;
         let mut next_stack: u64 = 0;
         let mut next_param_idx = 0; // Fastcall cares about overall param index
-        let mut ret = ABIArgVec::new();
 
         if args_or_rets == ArgsOrRets::Args && is_fastcall {
             // Fastcall always reserves 32 bytes of shadow space corresponding to
@@ -114,7 +114,7 @@ impl ABIMachineSpec for X64ABIMachineSpec {
                 let size = size as u64;
                 assert!(size % 8 == 0, "StructArgument size is not properly aligned");
                 next_stack += size;
-                ret.push(ABIArg::StructArg {
+                args.push(ABIArg::StructArg {
                     pointer: None,
                     offset,
                     size,
@@ -216,7 +216,7 @@ impl ABIMachineSpec for X64ABIMachineSpec {
                 }
             }
 
-            ret.push(ABIArg::Slots {
+            args.push(ABIArg::Slots {
                 slots,
                 purpose: param.purpose,
             });
@@ -225,14 +225,14 @@ impl ABIMachineSpec for X64ABIMachineSpec {
         let extra_arg = if add_ret_area_ptr {
             debug_assert!(args_or_rets == ArgsOrRets::Args);
             if let Some(reg) = get_intreg_for_arg(&call_conv, next_gpr, next_param_idx) {
-                ret.push(ABIArg::reg(
+                args.push(ABIArg::reg(
                     reg.to_real_reg().unwrap(),
                     types::I64,
                     ir::ArgumentExtension::None,
                     ir::ArgumentPurpose::Normal,
                 ));
             } else {
-                ret.push(ABIArg::stack(
+                args.push(ABIArg::stack(
                     next_stack as i64,
                     types::I64,
                     ir::ArgumentExtension::None,
@@ -240,7 +240,7 @@ impl ABIMachineSpec for X64ABIMachineSpec {
                 ));
                 next_stack += 8;
             }
-            Some(ret.len() - 1)
+            Some(args.args().len() - 1)
         } else {
             None
         };
@@ -252,7 +252,7 @@ impl ABIMachineSpec for X64ABIMachineSpec {
             return Err(CodegenError::ImplLimitExceeded);
         }
 
-        Ok((ret, next_stack as i64, extra_arg))
+        Ok((next_stack as i64, extra_arg))
     }
 
     fn fp_to_arg_offset(_call_conv: isa::CallConv, _flags: &settings::Flags) -> i64 {

cranelift/codegen/src/machinst/abi.rs

@@ -119,6 +119,7 @@ use std::collections::HashMap;
 use std::convert::TryFrom;
 use std::marker::PhantomData;
 use std::mem;
+use std::ops::Range;
 
 /// A small vector of instructions (with some reasonable size); appropriate for
 /// a small fixed sequence implementing one operation.
@@ -304,6 +305,40 @@ pub trait IsaFlags: Clone {
     }
 }
 
+/// Used as an out-parameter to accumulate a sequence of `ABIArg`s in
+/// `ABIMachineSpec::compute_arg_locs`. Wraps the shared allocation for all
+/// `ABIArg`s in `SigSet` and exposes just the args for the current
+/// `compute_arg_locs` call.
+pub struct ArgsAccumulator<'a> {
+    sig_set_abi_args: &'a mut Vec<ABIArg>,
+    start: usize,
+}
+
+impl<'a> ArgsAccumulator<'a> {
+    fn new(sig_set_abi_args: &'a mut Vec<ABIArg>) -> Self {
+        let start = sig_set_abi_args.len();
+        ArgsAccumulator {
+            sig_set_abi_args,
+            start,
+        }
+    }
+
+    #[inline]
+    pub fn push(&mut self, arg: ABIArg) {
+        self.sig_set_abi_args.push(arg)
+    }
+
+    #[inline]
+    pub fn args(&self) -> &[ABIArg] {
+        &self.sig_set_abi_args[self.start..]
+    }
+
+    #[inline]
+    pub fn args_mut(&mut self) -> &mut [ABIArg] {
+        &mut self.sig_set_abi_args[self.start..]
+    }
+}
+
 /// Trait implemented by machine-specific backend to provide information about
 /// register assignments and to allow generating the specific instructions for
 /// stack loads/saves, prologues/epilogues, etc.
@@ -342,16 +377,20 @@ pub trait ABIMachineSpec {
     /// Process a list of parameters or return values and allocate them to registers
     /// and stack slots.
     ///
-    /// Returns the list of argument locations, the stack-space used (rounded up
+    /// The argument locations should be pushed onto the given `ArgsAccumulator`
-    /// to as alignment requires), and if `add_ret_area_ptr` was passed, the
+    /// in order.
-    /// index of the extra synthetic arg that was added.
+    ///
+    /// Returns the stack-space used (rounded up to as alignment requires), and
+    /// if `add_ret_area_ptr` was passed, the index of the extra synthetic arg
+    /// that was added.
     fn compute_arg_locs<'a, I>(
         call_conv: isa::CallConv,
         flags: &settings::Flags,
         params: I,
         args_or_rets: ArgsOrRets,
         add_ret_area_ptr: bool,
-    ) -> CodegenResult<(ABIArgVec, i64, Option<usize>)>
+        args: ArgsAccumulator<'_>,
+    ) -> CodegenResult<(i64, Option<usize>)>
     where
         I: IntoIterator<Item = &'a ir::AbiParam>;
 
@@ -554,9 +593,6 @@ pub trait ABIMachineSpec {
     ) -> ir::ArgumentExtension;
 }
 
-// A vector of `ABIArg`s with inline capacity, since they are typically small.
-pub type ABIArgVec = SmallVec<[ABIArg; 6]>;
-
 /// The id of an ABI signature within the `SigSet`.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct Sig(u32);
@@ -567,22 +603,32 @@ cranelift_entity::entity_impl!(Sig);
 pub struct SigData {
     /// Argument locations (regs or stack slots). Stack offsets are relative to
     /// SP on entry to function.
-    args: ABIArgVec,
+    ///
+    /// These are indices into the `SigSet::abi_args`.
+    arg_indices: Range<u32>,
+
     /// Return-value locations. Stack offsets are relative to the return-area
     /// pointer.
-    rets: ABIArgVec,
+    ///
+    /// These are indices into the `SigSet::abi_args`.
+    ret_indices: Range<u32>,
+
     /// Space on stack used to store arguments.
     sized_stack_arg_space: i64,
+
     /// Space on stack used to store return values.
     sized_stack_ret_space: i64,
+
     /// Index in `args` of the stack-return-value-area argument.
     stack_ret_arg: Option<usize>,
+
     /// Calling convention used.
     call_conv: isa::CallConv,
 }
 
 impl SigData {
     pub fn from_func_sig<M: ABIMachineSpec>(
+        sigs: &mut SigSet,
         sig: &ir::Signature,
         flags: &settings::Flags,
     ) -> CodegenResult<SigData> {
@@ -591,35 +637,45 @@ impl SigData {
 
         // 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(
+        let rets_start = u32::try_from(sigs.abi_args.len()).unwrap();
+        let (sized_stack_ret_space, _) = M::compute_arg_locs(
             sig.call_conv,
             flags,
             returns,
             ArgsOrRets::Rets,
             /* extra ret-area ptr = */ false,
+            ArgsAccumulator::new(&mut sigs.abi_args),
         )?;
+        let rets_end = u32::try_from(sigs.abi_args.len()).unwrap();
+
         let need_stack_return_area = sized_stack_ret_space > 0;
-        let (args, sized_stack_arg_space, stack_ret_arg) = M::compute_arg_locs(
+        let args_start = u32::try_from(sigs.abi_args.len()).unwrap();
+        let (sized_stack_arg_space, stack_ret_arg) = M::compute_arg_locs(
             sig.call_conv,
             flags,
             &sig.params,
             ArgsOrRets::Args,
             need_stack_return_area,
+            ArgsAccumulator::new(&mut sigs.abi_args),
         )?;
+        let args_end = u32::try_from(sigs.abi_args.len()).unwrap();
+
+        let arg_indices = args_start..args_end;
+        let ret_indices = rets_start..rets_end;
 
         trace!(
             "ABISig: sig {:?} => args = {:?} rets = {:?} arg stack = {} ret stack = {} stack_ret_arg = {:?}",
             sig,
-            args,
+            arg_indices,
-            rets,
+            ret_indices,
             sized_stack_arg_space,
             sized_stack_ret_space,
             stack_ret_arg,
         );
 
         Ok(SigData {
-            args,
+            arg_indices,
-            rets,
+            ret_indices,
             sized_stack_arg_space,
             sized_stack_ret_space,
             stack_ret_arg,
@@ -627,8 +683,22 @@ impl SigData {
         })
     }
 
+    /// Get this signature's ABI arguments.
+    pub fn args<'a>(&self, sigs: &'a SigSet) -> &'a [ABIArg] {
+        let start = usize::try_from(self.arg_indices.start).unwrap();
+        let end = usize::try_from(self.arg_indices.end).unwrap();
+        &sigs.abi_args[start..end]
+    }
+
+    /// Get this signature's ABI returns.
+    pub fn rets<'a>(&self, sigs: &'a SigSet) -> &'a [ABIArg] {
+        let start = usize::try_from(self.ret_indices.start).unwrap();
+        let end = usize::try_from(self.ret_indices.end).unwrap();
+        &sigs.abi_args[start..end]
+    }
+
     /// Return all clobbers for the callsite.
-    pub fn call_clobbers<M: ABIMachineSpec>(&self) -> PRegSet {
+    pub fn call_clobbers<M: ABIMachineSpec>(&self, sigs: &SigSet) -> PRegSet {
         // Get clobbers: all caller-saves. These may include return value
         // regs, which we will remove from the clobber set below.
         let mut clobbers = M::get_regs_clobbered_by_call(self.call_conv);
@@ -636,7 +706,7 @@ impl SigData {
         // 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 {
+        for ret in self.rets(sigs) {
             if let &ABIArg::Slots {
                 ref slots, purpose, ..
             } = ret
@@ -661,16 +731,18 @@ impl SigData {
 
     /// Get the number of arguments expected.
     pub fn num_args(&self) -> usize {
+        let len = self.arg_indices.end - self.arg_indices.start;
+        let len = usize::try_from(len).unwrap();
         if self.stack_ret_arg.is_some() {
-            self.args.len() - 1
+            len - 1
         } else {
-            self.args.len()
+            len
         }
     }
 
     /// Get information specifying how to pass one argument.
-    pub fn get_arg(&self, idx: usize) -> ABIArg {
+    pub fn get_arg(&self, sigs: &SigSet, idx: usize) -> ABIArg {
-        self.args[idx].clone()
+        self.args(sigs)[idx].clone()
     }
 
     /// Get total stack space required for arguments.
@@ -680,12 +752,13 @@ impl SigData {
 
     /// Get the number of return values expected.
     pub fn num_rets(&self) -> usize {
-        self.rets.len()
+        let len = self.ret_indices.end - self.ret_indices.start;
+        usize::try_from(len).unwrap()
     }
 
     /// Get information specifying how to pass one return value.
-    pub fn get_ret(&self, idx: usize) -> ABIArg {
+    pub fn get_ret(&self, sigs: &SigSet, idx: usize) -> ABIArg {
-        self.rets[idx].clone()
+        self.rets(sigs)[idx].clone()
     }
 
     /// Get total stack space required for return values.
@@ -695,9 +768,9 @@ impl SigData {
 
     /// Get information specifying how to pass the implicit pointer
     /// to the return-value area on the stack, if required.
-    pub fn get_ret_arg(&self) -> Option<ABIArg> {
+    pub fn get_ret_arg(&self, sigs: &SigSet) -> Option<ABIArg> {
         let ret_arg = self.stack_ret_arg?;
-        Some(self.args[ret_arg].clone())
+        Some(self.args(sigs)[ret_arg].clone())
     }
 
     /// Get calling convention used.
@@ -728,6 +801,11 @@ pub struct SigSet {
     /// Interned `ir::SigRef`s that we already have an ABI signature for.
     ir_sig_ref_to_abi_sig: SecondaryMap<ir::SigRef, Option<Sig>>,
 
+    /// A single, shared allocation for all `ABIArg`s used by all
+    /// `SigData`s. Each `SigData` references its args/rets via indices into
+    /// this allocation.
+    abi_args: Vec<ABIArg>,
+
     /// The actual ABI signatures, keyed by `Sig`.
     sigs: PrimaryMap<Sig, SigData>,
 }
@@ -739,9 +817,12 @@ impl SigSet {
     where
         M: ABIMachineSpec,
     {
+        let arg_estimate = func.dfg.signatures.len() * 6;
+
         let mut sigs = SigSet {
             ir_signature_to_abi_sig: FxHashMap::default(),
             ir_sig_ref_to_abi_sig: SecondaryMap::with_capacity(func.dfg.signatures.len()),
+            abi_args: Vec::with_capacity(arg_estimate),
             sigs: PrimaryMap::with_capacity(1 + func.dfg.signatures.len()),
         };
 
@@ -776,7 +857,7 @@ impl SigSet {
         // `ir::Signature`.
         debug_assert!(!self.have_abi_sig_for_signature(&signature));
 
-        let sig_data = SigData::from_func_sig::<M>(&signature, flags)?;
+        let sig_data = SigData::from_func_sig::<M>(self, &signature, flags)?;
         let sig = self.sigs.push(sig_data);
         self.ir_signature_to_abi_sig.insert(signature, sig);
         Ok(sig)
@@ -795,7 +876,7 @@ impl SigSet {
             return Ok(sig);
         }
         let signature = &dfg.signatures[sig_ref];
-        let sig_data = SigData::from_func_sig::<M>(signature, flags)?;
+        let sig_data = SigData::from_func_sig::<M>(self, signature, flags)?;
         let sig = self.sigs.push(sig_data);
         self.ir_sig_ref_to_abi_sig[sig_ref] = Some(sig);
         Ok(sig)
@@ -896,11 +977,12 @@ pub struct Callee<M: ABIMachineSpec> {
 
 fn get_special_purpose_param_register(
     f: &ir::Function,
+    sigs: &SigSet,
     abi: &SigData,
     purpose: ir::ArgumentPurpose,
 ) -> Option<Reg> {
     let idx = f.signature.special_param_index(purpose)?;
-    match &abi.args[idx] {
+    match &abi.args(sigs)[idx] {
         &ABIArg::Slots { ref slots, .. } => match &slots[0] {
             &ABIArgSlot::Reg { reg, .. } => Some(reg.into()),
             _ => None,
@@ -977,13 +1059,17 @@ impl<M: ABIMachineSpec> Callee<M> {
         // stack limit. This can either be specified as a special-purpose
         // argument or as a global value which often calculates the stack limit
         // from the arguments.
-        let stack_limit =
+        let stack_limit = get_special_purpose_param_register(
-            get_special_purpose_param_register(f, &sigs[sig], ir::ArgumentPurpose::StackLimit)
+            f,
-                .map(|reg| (reg, smallvec![]))
+            sigs,
-                .or_else(|| {
+            &sigs[sig],
-                    f.stack_limit
+            ir::ArgumentPurpose::StackLimit,
-                        .map(|gv| gen_stack_limit::<M>(f, &sigs[sig], gv))
+        )
-                });
+        .map(|reg| (reg, smallvec![]))
+        .or_else(|| {
+            f.stack_limit
+                .map(|gv| gen_stack_limit::<M>(f, sigs, &sigs[sig], gv))
+        });
 
         // Determine whether a probestack call is required for large enough
         // frames (and the minimum frame size if so).
@@ -1102,16 +1188,18 @@ impl<M: ABIMachineSpec> Callee<M> {
 /// it's used, because we're not participating in register allocation anyway!
 fn gen_stack_limit<M: ABIMachineSpec>(
     f: &ir::Function,
+    sigs: &SigSet,
     abi: &SigData,
     gv: ir::GlobalValue,
 ) -> (Reg, SmallInstVec<M::I>) {
     let mut insts = smallvec![];
-    let reg = generate_gv::<M>(f, abi, gv, &mut insts);
+    let reg = generate_gv::<M>(f, sigs, abi, gv, &mut insts);
     return (reg, insts);
 }
 
 fn generate_gv<M: ABIMachineSpec>(
     f: &ir::Function,
+    sigs: &SigSet,
     abi: &SigData,
     gv: ir::GlobalValue,
     insts: &mut SmallInstVec<M::I>,
@@ -1119,7 +1207,7 @@ fn generate_gv<M: ABIMachineSpec>(
     match f.global_values[gv] {
         // Return the direct register the vmcontext is in
         ir::GlobalValueData::VMContext => {
-            get_special_purpose_param_register(f, abi, ir::ArgumentPurpose::VMContext)
+            get_special_purpose_param_register(f, sigs, abi, ir::ArgumentPurpose::VMContext)
                 .expect("no vmcontext parameter found")
         }
         // Load our base value into a register, then load from that register
@@ -1130,7 +1218,7 @@ fn generate_gv<M: ABIMachineSpec>(
             global_type: _,
             readonly: _,
         } => {
-            let base = generate_gv::<M>(f, abi, base, insts);
+            let base = generate_gv::<M>(f, sigs, abi, base, insts);
             let into_reg = Writable::from_reg(M::get_stacklimit_reg());
             insts.push(M::gen_load_base_offset(
                 into_reg,
@@ -1212,7 +1300,7 @@ impl<M: ABIMachineSpec> Callee<M> {
     /// They will be provided to `init()` as the `temps` arg if so.
     pub fn temps_needed(&self, sigs: &SigSet) -> Vec<Type> {
         let mut temp_tys = vec![];
-        for arg in &sigs[self.sig].args {
+        for arg in sigs[self.sig].args(sigs) {
             match arg {
                 &ABIArg::ImplicitPtrArg { pointer, .. } => match &pointer {
                     &ABIArgSlot::Reg { .. } => {}
@@ -1234,7 +1322,7 @@ impl<M: ABIMachineSpec> Callee<M> {
     /// once the lowering context exists.
     pub fn init(&mut self, sigs: &SigSet, temps: Vec<Writable<Reg>>) {
         let mut temps_iter = temps.into_iter();
-        for arg in &sigs[self.sig].args {
+        for arg in sigs[self.sig].args(sigs) {
             let temp = match arg {
                 &ABIArg::ImplicitPtrArg { pointer, .. } => match &pointer {
                     &ABIArgSlot::Reg { .. } => None,
@@ -1331,7 +1419,7 @@ impl<M: ABIMachineSpec> Callee<M> {
             }
         };
 
-        match &sigs[self.sig].args[idx] {
+        match &sigs[self.sig].args(sigs)[idx] {
             &ABIArg::Slots { ref slots, .. } => {
                 assert_eq!(into_regs.len(), slots.len());
                 for (slot, into_reg) in slots.iter().zip(into_regs.regs().iter()) {
@@ -1399,7 +1487,7 @@ impl<M: ABIMachineSpec> Callee<M> {
     ) -> SmallInstVec<M::I> {
         let mut ret = smallvec![];
         let word_bits = M::word_bits() as u8;
-        match &sigs[self.sig].rets[idx] {
+        match &sigs[self.sig].rets(sigs)[idx] {
             &ABIArg::Slots { ref slots, .. } => {
                 assert_eq!(from_regs.len(), slots.len());
                 for (slot, &from_reg) in slots.iter().zip(from_regs.regs().iter()) {
@@ -1508,7 +1596,7 @@ impl<M: ABIMachineSpec> Callee<M> {
     /// Generate a return instruction.
     pub fn gen_ret(&self, sigs: &SigSet) -> M::I {
         let mut rets = vec![];
-        for ret in &sigs[self.sig].rets {
+        for ret in sigs[self.sig].rets(sigs) {
             match ret {
                 ABIArg::Slots { slots, .. } => {
                     for slot in slots {
@@ -1895,7 +1983,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         flags: settings::Flags,
     ) -> CodegenResult<Caller<M>> {
         let sig = sigs.abi_sig_for_sig_ref(sig_ref);
-        let clobbers = sigs[sig].call_clobbers::<M>();
+        let clobbers = sigs[sig].call_clobbers::<M>(sigs);
         Ok(Caller {
             sig,
             uses: smallvec![],
@@ -1920,7 +2008,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         flags: settings::Flags,
     ) -> CodegenResult<Caller<M>> {
         let sig = sigs.abi_sig_for_signature(sig);
-        let clobbers = sigs[sig].call_clobbers::<M>();
+        let clobbers = sigs[sig].call_clobbers::<M>(sigs);
         Ok(Caller {
             sig,
             uses: smallvec![],
@@ -1945,7 +2033,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         flags: settings::Flags,
     ) -> CodegenResult<Caller<M>> {
         let sig = sigs.abi_sig_for_sig_ref(sig_ref);
-        let clobbers = sigs[sig].call_clobbers::<M>();
+        let clobbers = sigs[sig].call_clobbers::<M>(sigs);
         Ok(Caller {
             sig,
             uses: smallvec![],
@@ -1975,10 +2063,11 @@ impl<M: ABIMachineSpec> Caller<M> {
     /// Get the number of arguments expected.
     pub fn num_args(&self, sigs: &SigSet) -> usize {
         let data = &sigs[self.sig];
+        let len = data.args(sigs).len();
         if data.stack_ret_arg.is_some() {
-            data.args.len() - 1
+            len - 1
         } else {
-            data.args.len()
+            len
         }
     }
 
@@ -2007,7 +2096,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         idx: usize,
         from_regs: ValueRegs<Reg>,
     ) {
-        match &ctx.sigs()[self.sig].args[idx] {
+        match &ctx.sigs()[self.sig].args(ctx.sigs())[idx] {
             &ABIArg::Slots { .. } => {}
             &ABIArg::StructArg { offset, size, .. } => {
                 let src_ptr = from_regs.only_reg().unwrap();
@@ -2059,7 +2148,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         // How many temps do we need for extends? Allocate them ahead
         // of time, since we can't do it while we're iterating over
         // the sig and immutably borrowing `ctx`.
-        let needed_tmps = match &ctx.sigs()[self.sig].args[idx] {
+        let needed_tmps = match &ctx.sigs()[self.sig].args(ctx.sigs())[idx] {
             &ABIArg::Slots { ref slots, .. } => slots
                 .iter()
                 .map(|slot| match slot {
@@ -2084,7 +2173,7 @@ impl<M: ABIMachineSpec> Caller<M> {
             .map(|_| ctx.alloc_tmp(M::word_type()).only_reg().unwrap())
             .collect();
 
-        match &ctx.sigs()[self.sig].args[idx] {
+        match &ctx.sigs()[self.sig].args(ctx.sigs())[idx] {
             &ABIArg::Slots { ref slots, .. } => {
                 assert_eq!(from_regs.len(), slots.len());
                 for (slot, from_reg) in slots.iter().zip(from_regs.regs().iter()) {
@@ -2186,7 +2275,7 @@ impl<M: ABIMachineSpec> Caller<M> {
         into_regs: ValueRegs<Writable<Reg>>,
     ) -> SmallInstVec<M::I> {
         let mut insts = smallvec![];
-        match &ctx.sigs()[self.sig].rets[idx] {
+        match &ctx.sigs()[self.sig].rets(ctx.sigs())[idx] {
             &ABIArg::Slots { ref slots, .. } => {
                 assert_eq!(into_regs.len(), slots.len());
                 for (slot, into_reg) in slots.iter().zip(into_regs.regs().iter()) {

cranelift/codegen/src/machinst/isle.rs

@@ -439,7 +439,7 @@ macro_rules! isle_lower_prelude_methods {
         }
 
         fn abi_get_arg(&mut self, abi: &Sig, idx: usize) -> ABIArg {
-            self.lower_ctx.sigs()[*abi].get_arg(idx)
+            self.lower_ctx.sigs()[*abi].get_arg(self.lower_ctx.sigs(), idx)
         }
 
         fn abi_num_rets(&mut self, abi: &Sig) -> usize {
@@ -447,15 +447,15 @@ macro_rules! isle_lower_prelude_methods {
         }
 
         fn abi_get_ret(&mut self, abi: &Sig, idx: usize) -> ABIArg {
-            self.lower_ctx.sigs()[*abi].get_ret(idx)
+            self.lower_ctx.sigs()[*abi].get_ret(self.lower_ctx.sigs(), idx)
         }
 
         fn abi_ret_arg(&mut self, abi: &Sig) -> Option<ABIArg> {
-            self.lower_ctx.sigs()[*abi].get_ret_arg()
+            self.lower_ctx.sigs()[*abi].get_ret_arg(self.lower_ctx.sigs())
         }
 
         fn abi_no_ret_arg(&mut self, abi: &Sig) -> Option<()> {
-            if let Some(_) = self.lower_ctx.sigs()[*abi].get_ret_arg() {
+            if let Some(_) = self.lower_ctx.sigs()[*abi].get_ret_arg(self.lower_ctx.sigs()) {
                 None
             } else {
                 Some(())
@@ -709,7 +709,7 @@ macro_rules! isle_prelude_method_helpers {
                 // 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 ret = sigdata.get_ret(self.lower_ctx.sigs(), i);
                 let retval_regs = self.abi_arg_slot_regs(&ret).unwrap();
                 retval_insts.extend(
                     caller