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,7 +892,10 @@ impl<'a, F: Function> Env<'a, F> {
                 }
 
                 let resolved = parallel_moves.resolve();
-                let mut scratch_iter = RegTraversalIter::new(self.env, regclass, None, None, 0);
+                let mut scratch_iter = RegTraversalIter::new(
+                    self.env, regclass, None, None, 0,
+                    None, // We assume there is no limit on the set of 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

@@ -1072,9 +1072,15 @@ 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.as_valid(), scan_offset)
-            {
+            let limit = self.bundles[bundle].limit;
+            for preg in RegTraversalIter::new(
+                self.env,
+                class,
+                fixed_preg,
+                hint.as_valid(),
+                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: Option<PReg>,
         offset: usize,
+        limit: Option<usize>,
     ) -> Self {
         debug_assert!(fixed != Some(PReg::invalid()));
         debug_assert!(hint != Some(PReg::invalid()));
@@ -96,6 +98,7 @@ impl<'a> RegTraversalIter<'a> {
             hint,
             preferred,
             non_preferred,
+            limit,
         }
     }
 }
@@ -110,21 +113,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

@@ -42,7 +42,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(_) =
@@ -53,6 +54,7 @@ impl<'a, F: Function> Env<'a, F> {
                     break;
                 }
             }
+
             if !success {
                 trace!(
                     "spilling bundle {:?}: marking spillset {:?} as required",