Append link and sret arguments in legalize_signature.

These special-purpose arguments and return values are only relevant for
the function being compiled, so add a `current` flag to
legalize_signature().

- Add the necessary argument values to the entry block to represent
  the special-purpose arguments.
- Propagate the link and sret arguments to return instructions if the
  legalized signature asks for it.

lib/cretonne/src/ir/extfunc.rs

@@ -6,7 +6,7 @@
 //! This module declares the data types used to represent external functions and call signatures.
 
 use ir::{Type, FunctionName, SigRef, ArgumentLoc};
-use isa::RegInfo;
+use isa::{RegInfo, RegUnit};
 use std::cmp;
 use std::fmt;
 use std::str::FromStr;
@@ -133,6 +133,16 @@ impl ArgumentType {
         }
     }
 
+    /// Create an argument type for a special-purpose register.
+    pub fn special_reg(vt: Type, purpose: ArgumentPurpose, regunit: RegUnit) -> ArgumentType {
+        ArgumentType {
+            value_type: vt,
+            extension: ArgumentExtension::None,
+            purpose: purpose,
+            location: ArgumentLoc::Reg(regunit),
+        }
+    }
+
     /// Return an object that can display `self` with correct register names.
     pub fn display<'a, R: Into<Option<&'a RegInfo>>>(&'a self, regs: R) -> DisplayArgumentType<'a> {
         DisplayArgumentType(self, regs.into())

lib/cretonne/src/isa/mod.rs

@@ -166,7 +166,24 @@ pub trait TargetIsa {
     /// - Vector types can be bit-cast and broken down into smaller vectors or scalars.
     ///
     /// The legalizer will adapt argument and return values as necessary at all ABI boundaries.
-    fn legalize_signature(&self, _sig: &mut Signature) {
+    ///
+    /// When this function is called to legalize the signature of the function currently begin
+    /// compiler, `_current` is true. The legalized signature can then also contain special purpose
+    /// arguments and return values such as:
+    ///
+    /// - A `link` argument representing the link registers on RISC architectures that don't push
+    ///   the return address on the stack.
+    /// - A `link` return value which will receive the value that was passed to the `link`
+    ///   argument.
+    /// - An `sret` argument can be added if one wasn't present already. This is necessary if the
+    ///   signature returns more values than registers are available for returning values.
+    /// - An `sret` return value can be added if the ABI requires a function to return its `sret`
+    ///   argument in a register.
+    ///
+    /// Arguments and return values for the caller's frame pointer and other callee-saved registers
+    /// should not be added by this function. These arguments are not added until after register
+    /// allocation.
+    fn legalize_signature(&self, _sig: &mut Signature, _current: bool) {
         unimplemented!()
     }
 

lib/cretonne/src/isa/riscv/abi.rs

@@ -6,7 +6,7 @@
 //! This doesn't support the soft-float ABI at the moment.
 
 use abi::{ArgAction, ValueConversion, ArgAssigner, legalize_args};
-use ir::{Signature, Type, ArgumentType, ArgumentLoc, ArgumentExtension};
+use ir::{Signature, Type, ArgumentType, ArgumentLoc, ArgumentExtension, ArgumentPurpose};
 use isa::riscv::registers::{GPR, FPR};
 use settings as shared_settings;
 
@@ -80,7 +80,7 @@ impl ArgAssigner for Args {
 }
 
 /// Legalize `sig` for RISC-V.
-pub fn legalize_signature(sig: &mut Signature, flags: &shared_settings::Flags) {
+pub fn legalize_signature(sig: &mut Signature, flags: &shared_settings::Flags, current: bool) {
     let bits = if flags.is_64bit() { 64 } else { 32 };
 
     let mut args = Args::new(bits);
@@ -88,4 +88,17 @@ pub fn legalize_signature(sig: &mut Signature, flags: &shared_settings::Flags) {
 
     let mut rets = Args::new(bits);
     legalize_args(&mut sig.return_types, &mut rets);
+
+    if current {
+        let ptr = Type::int(bits).unwrap();
+
+        // Add the link register as an argument and return value.
+        //
+        // The `jalr` instruction implementing a return can technically accept the return address
+        // in any register, but a micro-architecture with a return address predictor will only
+        // recognize it as a return if the address is in `x1`.
+        let link = ArgumentType::special_reg(ptr, ArgumentPurpose::Link, GPR.unit(1));
+        sig.argument_types.push(link);
+        sig.return_types.push(link);
+    }
 }

lib/cretonne/src/isa/riscv/mod.rs

@@ -78,9 +78,9 @@ impl TargetIsa for Isa {
                 })
     }
 
-    fn legalize_signature(&self, sig: &mut Signature) {
+    fn legalize_signature(&self, sig: &mut Signature, current: bool) {
         // We can pass in `self.isa_flags` too, if we need it.
-        abi::legalize_signature(sig, &self.shared_flags)
+        abi::legalize_signature(sig, &self.shared_flags, current)
     }
 
     fn emit_inst(&self, func: &Function, inst: Inst, sink: &mut CodeSink) {

lib/cretonne/src/legalizer/boundary.rs

@@ -20,7 +20,7 @@
 use abi::{legalize_abi_value, ValueConversion};
 use flowgraph::ControlFlowGraph;
 use ir::{Function, Cursor, DataFlowGraph, Inst, InstBuilder, Ebb, Type, Value, Signature, SigRef,
-         ArgumentType};
+         ArgumentType, ArgumentPurpose};
 use ir::instructions::CallInfo;
 use isa::TargetIsa;
 use legalizer::split::{isplit, vsplit};
@@ -31,9 +31,9 @@ use legalizer::split::{isplit, vsplit};
 /// change the entry block arguments, calls, or return instructions, so this can leave the function
 /// in a state with type discrepancies.
 pub fn legalize_signatures(func: &mut Function, isa: &TargetIsa) {
-    isa.legalize_signature(&mut func.signature);
+    isa.legalize_signature(&mut func.signature, true);
     for sig in func.dfg.signatures.keys() {
-        isa.legalize_signature(&mut func.dfg.signatures[sig]);
+        isa.legalize_signature(&mut func.dfg.signatures[sig], false);
     }
 
     if let Some(entry) = func.layout.entry_block() {
@@ -50,6 +50,9 @@ pub fn legalize_signatures(func: &mut Function, isa: &TargetIsa) {
 /// The original entry EBB arguments are computed from the new ABI arguments by code inserted at
 /// the top of the entry block.
 fn legalize_entry_arguments(func: &mut Function, entry: Ebb) {
+    let mut has_sret = false;
+    let mut has_link = false;
+
     // Insert position for argument conversion code.
     // We want to insert instructions before the first instruction in the entry block.
     // If the entry block is empty, append instructions to it instead.
@@ -73,11 +76,22 @@ fn legalize_entry_arguments(func: &mut Function, entry: Ebb) {
             // No value translation is necessary, this argument matches the ABI type.
             // Just use the original EBB argument value. This is the most common case.
             func.dfg.attach_ebb_arg(entry, arg);
+            match abi_types[abi_arg].purpose {
+                ArgumentPurpose::Normal => {}
+                ArgumentPurpose::StructReturn => {
+                    assert!(!has_sret, "Multiple sret arguments found");
+                    has_sret = true;
+                }
+                _ => panic!("Unexpected special-purpose arg {}", abi_types[abi_arg]),
+            }
             abi_arg += 1;
         } else {
             // Compute the value we want for `arg` from the legalized ABI arguments.
             let mut get_arg = |dfg: &mut DataFlowGraph, ty| {
                 let abi_type = abi_types[abi_arg];
+                assert_eq!(abi_type.purpose,
+                           ArgumentPurpose::Normal,
+                           "Can't legalize special-purpose argument");
                 if ty == abi_type.value_type {
                     abi_arg += 1;
                     Ok(dfg.append_ebb_arg(entry, ty))
@@ -92,6 +106,35 @@ fn legalize_entry_arguments(func: &mut Function, entry: Ebb) {
             assert_eq!(func.dfg.resolve_aliases(arg), converted);
         }
     }
+
+    // The legalized signature may contain additional arguments representing special-purpose
+    // registers.
+    for &arg in &abi_types[abi_arg..] {
+        match arg.purpose {
+            // Any normal arguments should have been processed above.
+            ArgumentPurpose::Normal => {
+                panic!("Leftover arg: {}", arg);
+            }
+            // The callee-save arguments should not appear until after register allocation is
+            // done.
+            ArgumentPurpose::FramePointer |
+            ArgumentPurpose::CalleeSaved => {
+                panic!("Premature callee-saved arg {}", arg);
+            }
+            // These can be meaningfully added by `legalize_signature()`.
+            ArgumentPurpose::Link => {
+                assert!(!has_link, "Multiple link arguments found");
+                has_link = true;
+            }
+            ArgumentPurpose::StructReturn => {
+                assert!(!has_sret, "Multiple sret arguments found");
+                has_sret = true;
+            }
+        }
+        // Just create entry block values to match here. We will use them in `handle_return_abi()`
+        // below.
+        func.dfg.append_ebb_arg(entry, arg.value_type);
+    }
 }
 
 /// Legalize the results returned from a call instruction to match the ABI signature.
@@ -445,7 +488,7 @@ pub fn handle_call_abi(dfg: &mut DataFlowGraph, cfg: &ControlFlowGraph, pos: &mu
     true
 }
 
-/// Insert ABI conversion code before and after the call instruction at `pos`.
+/// Insert ABI conversion code before and after the return instruction at `pos`.
 ///
 /// Return `true` if any instructions were inserted.
 pub fn handle_return_abi(dfg: &mut DataFlowGraph,
@@ -461,15 +504,58 @@ pub fn handle_return_abi(dfg: &mut DataFlowGraph,
         return false;
     }
 
-    let abi_args = sig.return_types.len();
+    // Count the special-purpose return values (`link` and `sret`) that were appended to the
+    // legalized signature.
+    let special_args = sig.return_types
+        .iter()
+        .rev()
+        .take_while(|&rt| {
+                        rt.purpose == ArgumentPurpose::Link ||
+                        rt.purpose == ArgumentPurpose::StructReturn
+                    })
+        .count();
+
+    let abi_args = sig.return_types.len() - special_args;
     legalize_inst_arguments(dfg,
                             cfg,
                             pos,
                             abi_args,
                             |_, abi_arg| sig.return_types[abi_arg]);
+    assert_eq!(dfg.inst_variable_args(inst).len(), abi_args);
+
+    // Append special return arguments for any `sret` and `link` return values added to the
+    // legalized signature. These values should simply be propagated from the entry block
+    // arguments.
+    if special_args > 0 {
+        dbg!("Adding {} special-purpose arguments to {}",
+             special_args,
+             dfg.display_inst(inst));
+        let mut vlist = dfg[inst].take_value_list().unwrap();
+        for arg in &sig.return_types[abi_args..] {
+            match arg.purpose {
+                ArgumentPurpose::Link |
+                ArgumentPurpose::StructReturn => {}
+                ArgumentPurpose::Normal => panic!("unexpected return value {}", arg),
+                _ => panic!("Unsupported special purpose return value {}", arg),
+            }
+            // A `link` or `sret` return value can only appear in a signature that has a unique
+            // matching argument. They are appended at the end, so search the signature from the
+            // end.
+            let idx = sig.argument_types
+                .iter()
+                .rposition(|t| t.purpose == arg.purpose)
+                .expect("No matching special purpose argument.");
+            // Get the corresponding entry block value and add it to the return instruction's
+            // arguments.
+            let val = dfg.ebb_args(pos.layout.entry_block().unwrap())[idx];
+            debug_assert_eq!(dfg.value_type(val), arg.value_type);
+            vlist.push(val, &mut dfg.value_lists);
+        }
+        dfg[inst].put_value_list(vlist);
+    }
 
     debug_assert!(check_return_signature(dfg, inst, sig),
-                  "Signature still wrong: {}, sig{}",
+                  "Signature still wrong: {} / signature {}",
                   dfg.display_inst(inst),
                   sig);