Add support for fixed stack slots

This works by allowing a PReg to be marked as being a stack location
instead of a physical register.

src/checker.rs

@@ -78,7 +78,7 @@
 
 use crate::{
     Allocation, AllocationKind, Block, Edit, Function, Inst, InstPosition, Operand,
-    OperandConstraint, OperandKind, OperandPos, Output, PReg, ProgPoint, VReg,
+    OperandConstraint, OperandKind, OperandPos, Output, PReg, ProgPoint, RegClass, VReg,
 };
 
 use std::collections::{HashMap, HashSet, VecDeque};

src/ion/data_structures.rs

@@ -266,6 +266,7 @@ pub struct VRegData {
 pub struct PRegData {
     pub reg: PReg,
     pub allocations: LiveRangeSet,
+    pub is_stack: bool,
 }
 
 #[derive(Clone, Debug)]

src/ion/liveranges.rs

@@ -102,6 +102,7 @@ impl<'a, F: Function> Env<'a, F> {
             PRegData {
                 reg: PReg::invalid(),
                 allocations: LiveRangeSet::new(),
+                is_stack: false,
             },
         );
         for i in 0..=PReg::MAX {
@@ -110,6 +111,9 @@ impl<'a, F: Function> Env<'a, F> {
             let preg_float = PReg::new(i, RegClass::Float);
             self.pregs[preg_float.index()].reg = preg_float;
         }
+        for &preg in &self.env.fixed_stack_slots {
+            self.pregs[preg.index()].is_stack = true;
+        }
         // Create VRegs from the vreg count.
         for idx in 0..self.func.num_vregs() {
             // We'll fill in the real details when we see the def.

src/ion/merge.rs

@@ -99,9 +99,9 @@ impl<'a, F: Function> Env<'a, F> {
             || self.bundles[to.index()].cached_stack()
             || self.bundles[to.index()].cached_fixed()
         {
-            let req = self
+            let req_from = self.compute_requirement(from);
-                .compute_requirement(from)
+            let req_to = self.compute_requirement(to);
-                .merge(self.compute_requirement(to));
+            let req = self.merge_requirement(req_from, req_to);
             if req == Requirement::Conflict {
                 log::trace!(" -> conflicting requirements; aborting merge");
                 return false;

src/ion/moves.rs

@@ -35,6 +35,16 @@ impl<'a, F: Function> Env<'a, F> {
         pos == self.cfginfo.block_exit[block.index()]
     }
 
+    fn allocation_is_stack(&self, alloc: Allocation) -> bool {
+        if alloc.is_stack() {
+            true
+        } else if let Some(preg) = alloc.as_reg() {
+            self.pregs[preg.index()].is_stack
+        } else {
+            false
+        }
+    }
+
     pub fn insert_move(
         &mut self,
         pos: ProgPoint,
@@ -968,9 +978,9 @@ impl<'a, F: Function> Env<'a, F> {
                 let scratch_used = resolved.iter().any(|&(src, dst, _)| {
                     src == Allocation::reg(scratch) || dst == Allocation::reg(scratch)
                 });
-                let stack_stack_move = resolved
+                let stack_stack_move = resolved.iter().any(|&(src, dst, _)| {
-                    .iter()
+                    self.allocation_is_stack(src) && self.allocation_is_stack(dst)
-                    .any(|&(src, dst, _)| src.is_stack() && dst.is_stack());
+                });
                 let extra_slot = if scratch_used && stack_stack_move {
                     if self.extra_spillslot[regclass as u8 as usize].is_none() {
                         let slot = self.allocate_spillslot(regclass);
@@ -989,7 +999,7 @@ impl<'a, F: Function> Env<'a, F> {
                         if dst == Allocation::reg(scratch) {
                             scratch_used_yet = true;
                         }
-                        if src.is_stack() && dst.is_stack() {
+                        if self.allocation_is_stack(src) && self.allocation_is_stack(dst) {
                             if !scratch_used_yet {
                                 self.add_edit(
                                     pos,

src/ion/process.rs

@@ -374,7 +374,7 @@ impl<'a, F: Function> Env<'a, F> {
         for entry in &self.bundles[bundle.index()].ranges {
             for u in &self.ranges[entry.index.index()].uses {
                 let this_req = Requirement::from_operand(u.operand);
-                req = req.merge(this_req);
+                req = self.merge_requirement(req, this_req);
                 if req == Requirement::Conflict {
                     return u.pos;
                 }
@@ -754,11 +754,14 @@ impl<'a, F: Function> Env<'a, F> {
         let class = self.spillsets[self.bundles[bundle.index()].spillset.index()].class;
         let req = self.compute_requirement(bundle);
         // Grab a hint from either the queue or our spillset, if any.
-        let hint_reg = if reg_hint != PReg::invalid() {
+        let mut hint_reg = if reg_hint != PReg::invalid() {
             reg_hint
         } else {
             self.spillsets[self.bundles[bundle.index()].spillset.index()].reg_hint
         };
+        if self.pregs[hint_reg.index()].is_stack {
+            hint_reg = PReg::invalid();
+        }
         log::trace!("process_bundle: bundle {:?} hint {:?}", bundle, hint_reg,);
 
         if let Requirement::Conflict = req {

src/ion/requirement.rs

@@ -25,20 +25,6 @@ pub enum Requirement {
     Conflict,
 }
 impl Requirement {
-    #[inline(always)]
-    pub fn merge(self, other: Requirement) -> Requirement {
-        match (self, other) {
-            (Requirement::Unknown, other) | (other, Requirement::Unknown) => other,
-            (Requirement::Conflict, _) | (_, Requirement::Conflict) => Requirement::Conflict,
-            (other, Requirement::Any) | (Requirement::Any, other) => other,
-            (Requirement::Stack, Requirement::Stack) => self,
-            (Requirement::Register, Requirement::Fixed(preg))
-            | (Requirement::Fixed(preg), Requirement::Register) => Requirement::Fixed(preg),
-            (Requirement::Register, Requirement::Register) => self,
-            (Requirement::Fixed(a), Requirement::Fixed(b)) if a == b => self,
-            _ => Requirement::Conflict,
-        }
-    }
     #[inline(always)]
     pub fn from_operand(op: Operand) -> Requirement {
         match op.constraint() {
@@ -51,6 +37,23 @@ impl Requirement {
 }
 
 impl<'a, F: Function> Env<'a, F> {
+    #[inline(always)]
+    pub fn merge_requirement(&self, a: Requirement, b: Requirement) -> Requirement {
+        match (a, b) {
+            (Requirement::Unknown, other) | (other, Requirement::Unknown) => other,
+            (Requirement::Conflict, _) | (_, Requirement::Conflict) => Requirement::Conflict,
+            (other, Requirement::Any) | (Requirement::Any, other) => other,
+            (Requirement::Stack, Requirement::Stack) => Requirement::Stack,
+            (Requirement::Register, Requirement::Register) => Requirement::Register,
+            (Requirement::Register, Requirement::Fixed(preg))
+            | (Requirement::Fixed(preg), Requirement::Register) if !self.pregs[preg.index()].is_stack => Requirement::Fixed(preg),
+            (Requirement::Stack, Requirement::Fixed(preg))
+            | (Requirement::Fixed(preg), Requirement::Stack) if self.pregs[preg.index()].is_stack => Requirement::Fixed(preg),
+            (Requirement::Fixed(a), Requirement::Fixed(b)) if a == b => Requirement::Fixed(a),
+            _ => Requirement::Conflict,
+        }
+    }
+
     pub fn compute_requirement(&self, bundle: LiveBundleIndex) -> Requirement {
         let mut req = Requirement::Unknown;
         log::trace!("compute_requirement: {:?}", bundle);
@@ -59,7 +62,7 @@ impl<'a, F: Function> Env<'a, F> {
             for u in &self.ranges[entry.index.index()].uses {
                 log::trace!("  -> use {:?}", u);
                 let r = Requirement::from_operand(u.operand);
-                req = req.merge(r);
+                req = self.merge_requirement(req, r);
                 log::trace!("     -> req {:?}", req);
             }
         }

src/lib.rs

@@ -1154,6 +1154,14 @@ pub struct MachineEnv {
     /// the emission of two machine-code instructions, this lowering
     /// can use the scratch register between them.
     pub scratch_by_class: [PReg; 2],
+
+    /// Some `PReg`s can be designated as locations on the stack rather than
+    /// actual registers. These can be used to tell the register allocator about
+    /// pre-defined stack slots used for function arguments and return values.
+    ///
+    /// `PReg`s in this list cannot be used as a scratch register or as an
+    /// allocatable regsiter.
+    pub fixed_stack_slots: Vec<PReg>,
 }
 
 /// The output of the register allocator.