Constrain iteration of registers with a Limit

During allocation, we iterate over the available registers and pick one
within all the constraints (i.e., fixed, hints, preferred/non-preferred
lists). This adds a limit to all of those constraints: unless the
register is fixed, we only pick a register with a HW encoding below the
limit, if one is given.

Author: abrown

src/ion/moves.rs

@@ -892,8 +892,14 @@ impl<'a, F: Function> Env<'a, F> {
                 }
 
                 let resolved = parallel_moves.resolve();
-                let mut scratch_iter =
-                    RegTraversalIter::new(self.env, regclass, None, PReg::invalid(), 0);
+                let mut scratch_iter = RegTraversalIter::new(
+                    self.env,
+                    regclass,
+                    None,
+                    PReg::invalid(),
+                    0,
+                    None, // We assume there is no limit on registers available for moves.
+                );
                 let mut dedicated_scratch = self.env.scratch_by_class[regclass as usize];
                 let key = LiveRangeKey::from_range(&CodeRange {
                     from: pos_prio.pos,

src/ion/process.rs

@@ -1078,7 +1078,10 @@ impl<'a, F: Function> Env<'a, F> {
                 + bundle.index();
 
             self.ctx.output.stats.process_bundle_reg_probe_start_any += 1;
-            for preg in RegTraversalIter::new(self.env, class, fixed_preg, hint_reg, scan_offset) {
+            let limit = self.bundles[bundle].limit;
+            for preg in
+                RegTraversalIter::new(self.env, class, fixed_preg, hint_reg, scan_offset, limit)
+            {
                 self.ctx.output.stats.process_bundle_reg_probes_any += 1;
                 let preg_idx = PRegIndex::new(preg.index());
                 trace!("trying preg {:?}", preg_idx);

src/ion/reg_traversal.rs

@@ -72,6 +72,7 @@ pub struct RegTraversalIter<'a> {
     hint: Option<PReg>,
     preferred: Cursor<'a>,
     non_preferred: Cursor<'a>,
+    limit: Option<usize>,
 }
 
 impl<'a> RegTraversalIter<'a> {
@@ -81,6 +82,7 @@ impl<'a> RegTraversalIter<'a> {
         fixed: Option<PReg>,
         hint: PReg,
         offset: usize,
+        limit: Option<usize>,
     ) -> Self {
         let hint = if hint != PReg::invalid() {
             Some(hint)
@@ -99,6 +101,7 @@ impl<'a> RegTraversalIter<'a> {
             hint,
             preferred,
             non_preferred,
+            limit,
         }
     }
 }
@@ -113,21 +116,23 @@ impl<'a> core::iter::Iterator for RegTraversalIter<'a> {
 
         if self.use_hint {
             self.use_hint = false;
-            return self.hint;
+            if self.hint.unwrap().hw_enc() < self.limit.unwrap_or(usize::MAX) {
+                return self.hint;
+            }
         }
 
         while !self.preferred.done() {
             let reg = self.preferred.advance();
-            if Some(reg) == self.hint {
-                continue; // Try again; we already tried the hint.
+            if Some(reg) == self.hint || reg.hw_enc() >= self.limit.unwrap_or(usize::MAX) {
+                continue; // Try again; we already tried the hint or we are outside of the register range limit.
             }
             return Some(reg);
         }
 
         while !self.non_preferred.done() {
             let reg = self.non_preferred.advance();
-            if Some(reg) == self.hint {
-                continue; // Try again; we already tried the hint.
+            if Some(reg) == self.hint || reg.hw_enc() >= self.limit.unwrap_or(usize::MAX) {
+                continue; // Try again; we already tried the hint or we are outside of the register range limit.
             }
             return Some(reg);
         }

src/ion/spill.rs

@@ -40,7 +40,8 @@ impl<'a, F: Function> Env<'a, F> {
 
             let mut success = false;
             self.ctx.output.stats.spill_bundle_reg_probes += 1;
-            for preg in RegTraversalIter::new(self.env, class, None, hint, bundle.index()) {
+            let limit = self.bundles[bundle].limit;
+            for preg in RegTraversalIter::new(self.env, class, None, hint, bundle.index(), limit) {
                 trace!("trying bundle {:?} to preg {:?}", bundle, preg);
                 let preg_idx = PRegIndex::new(preg.index());
                 if let AllocRegResult::Allocated(_) =
@@ -51,6 +52,7 @@ impl<'a, F: Function> Env<'a, F> {
                     break;
                 }
             }
+
             if !success {
                 trace!(
                     "spilling bundle {:?}: marking spillset {:?} as required",