Refactor RegTraversalIter

src/ion/moves.rs

@@ -892,14 +892,8 @@ impl<'a, F: Function> Env<'a, F> {
                 }
 
                 let resolved = parallel_moves.resolve();
-                let mut scratch_iter = RegTraversalIter::new(
-                    self.env,
-                    regclass,
-                    PReg::invalid(),
-                    PReg::invalid(),
-                    0,
-                    None,
-                );
+                let mut scratch_iter =
+                    RegTraversalIter::new(self.env, regclass, None, PReg::invalid(), 0);
                 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,14 +1078,7 @@ 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,
-                hint_reg,
-                PReg::invalid(),
-                scan_offset,
-                fixed_preg,
-            ) {
+            for preg in RegTraversalIter::new(self.env, class, fixed_preg, hint_reg, scan_offset) {
                 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

@@ -1,79 +1,104 @@
+//! Iterate over available registers.
+
 use crate::{MachineEnv, PReg, RegClass};
 
-/// This iterator represents a traversal through all allocatable
-/// registers of a given class, in a certain order designed to
-/// minimize allocation contention.
+/// Keep track of where we are in the register traversal.
+struct Cursor<'a> {
+    registers: &'a [PReg],
+    index: usize,
+    offset: usize,
+}
+
+impl<'a> Cursor<'a> {
+    #[inline]
+    fn new(registers: &'a [PReg], offset_hint: usize) -> Self {
+        let offset = if registers.len() > 0 {
+            offset_hint % registers.len()
+        } else {
+            0
+        };
+        Self {
+            registers,
+            index: 0,
+            offset,
+        }
+    }
+
+    /// Wrap around the end of the register list; [`Cursor::done`] guarantees we
+    /// do not see the same register twice.
+    #[inline]
+    fn wrap(index: usize, end: usize) -> usize {
+        if index >= end {
+            index - end
+        } else {
+            index
+        }
+    }
+
+    /// Advance to the next register and return it.
+    #[inline]
+    fn advance(&mut self) -> PReg {
+        let loc = Self::wrap(self.index + self.offset, self.registers.len());
+        let reg = self.registers[loc];
+        self.index += 1;
+        reg
+    }
+
+    /// Return `true` if we have seen all registers.
+    #[inline]
+    fn done(&self) -> bool {
+        self.index >= self.registers.len()
+    }
+}
+
+/// This iterator represents a traversal through all allocatable registers of a
+/// given class, in a certain order designed to minimize allocation contention.
 ///
 /// The order in which we try registers is somewhat complex:
-/// - First, if there is a hint, we try that.
-/// - Then, we try registers in a traversal order that is based on an
-///   "offset" (usually the bundle index) spreading pressure evenly
-///   among registers to reduce commitment-map contention.
-/// - Within that scan, we try registers in two groups: first,
-///   prferred registers; then, non-preferred registers. (In normal
-///   usage, these consist of caller-save and callee-save registers
-///   respectively, to minimize clobber-saves; but they need not.)
-
+/// - First, if the register is fixed (i.e., pre-assigned), return that and stop
+///   iteration.
+/// - Then, if there is a hint, try that one.
+/// - Next we try registers in a traversal order that is based on an "offset"
+///   (usually the bundle index) spreading pressure evenly among registers to
+///   reduce commitment-map contention.
+/// - Within that scan, we try registers in two groups: first, preferred
+///   registers; then, non-preferred registers. (In normal usage, these consist
+///   of caller-save and callee-save registers respectively, to minimize
+///   clobber-saves; but they need not.)
 pub struct RegTraversalIter<'a> {
-    env: &'a MachineEnv,
-    class: usize,
-    hints: [Option<PReg>; 2],
-    hint_idx: usize,
-    pref_idx: usize,
-    non_pref_idx: usize,
-    offset_pref: usize,
-    offset_non_pref: usize,
     is_fixed: bool,
     fixed: Option<PReg>,
+    use_hint: bool,
+    hint: Option<PReg>,
+    preferred: Cursor<'a>,
+    non_preferred: Cursor<'a>,
 }
 
 impl<'a> RegTraversalIter<'a> {
     pub fn new(
         env: &'a MachineEnv,
         class: RegClass,
-        hint_reg: PReg,
-        hint2_reg: PReg,
-        offset: usize,
         fixed: Option<PReg>,
+        hint: PReg,
+        offset: usize,
     ) -> Self {
-        let mut hint_reg = if hint_reg != PReg::invalid() {
-            Some(hint_reg)
-        } else {
-            None
-        };
-        let mut hint2_reg = if hint2_reg != PReg::invalid() {
-            Some(hint2_reg)
+        let hint = if hint != PReg::invalid() {
+            Some(hint)
         } else {
             None
         };
 
-        if hint_reg.is_none() {
-            hint_reg = hint2_reg;
-            hint2_reg = None;
-        }
-        let hints = [hint_reg, hint2_reg];
         let class = class as u8 as usize;
-        let offset_pref = if env.preferred_regs_by_class[class].len() > 0 {
-            offset % env.preferred_regs_by_class[class].len()
-        } else {
-            0
-        };
-        let offset_non_pref = if env.non_preferred_regs_by_class[class].len() > 0 {
-            offset % env.non_preferred_regs_by_class[class].len()
-        } else {
-            0
-        };
+        let preferred = Cursor::new(&env.preferred_regs_by_class[class], offset);
+        let non_preferred = Cursor::new(&env.non_preferred_regs_by_class[class], offset);
+
         Self {
-            env,
-            class,
-            hints,
-            hint_idx: 0,
-            pref_idx: 0,
-            non_pref_idx: 0,
-            offset_pref,
-            offset_non_pref,
             is_fixed: fixed.is_some(),
             fixed,
+            use_hint: hint.is_some(),
+            hint,
+            preferred,
+            non_preferred,
         }
     }
 }
@@ -83,41 +108,30 @@ impl<'a> core::iter::Iterator for RegTraversalIter<'a> {
 
     fn next(&mut self) -> Option<PReg> {
         if self.is_fixed {
-            let ret = self.fixed;
-            self.fixed = None;
-            return ret;
+            return self.fixed.take();
         }
 
-        fn wrap(idx: usize, limit: usize) -> usize {
-            if idx >= limit {
-                idx - limit
-            } else {
-                idx
-            }
+        if self.use_hint {
+            self.use_hint = false;
+            return self.hint;
         }
-        if self.hint_idx < 2 && self.hints[self.hint_idx].is_some() {
-            let h = self.hints[self.hint_idx];
-            self.hint_idx += 1;
-            return h;
-        }
-        while self.pref_idx < self.env.preferred_regs_by_class[self.class].len() {
-            let arr = &self.env.preferred_regs_by_class[self.class][..];
-            let r = arr[wrap(self.pref_idx + self.offset_pref, arr.len())];
-            self.pref_idx += 1;
-            if Some(r) == self.hints[0] || Some(r) == self.hints[1] {
-                continue;
+
+        while !self.preferred.done() {
+            let reg = self.preferred.advance();
+            if Some(reg) == self.hint {
+                continue; // Try again; we already tried the hint.
             }
-            return Some(r);
+            return Some(reg);
         }
-        while self.non_pref_idx < self.env.non_preferred_regs_by_class[self.class].len() {
-            let arr = &self.env.non_preferred_regs_by_class[self.class][..];
-            let r = arr[wrap(self.non_pref_idx + self.offset_non_pref, arr.len())];
-            self.non_pref_idx += 1;
-            if Some(r) == self.hints[0] || Some(r) == self.hints[1] {
-                continue;
+
+        while !self.non_preferred.done() {
+            let reg = self.non_preferred.advance();
+            if Some(reg) == self.hint {
+                continue; // Try again; we already tried the hint.
             }
-            return Some(r);
+            return Some(reg);
         }
+
         None
     }
 }

src/ion/spill.rs

@@ -13,8 +13,8 @@
 //! Spillslot allocation.
 
 use super::{
-    AllocRegResult, Env, LiveRangeKey, PReg, PRegIndex, RegTraversalIter, SpillSetIndex,
-    SpillSlotData, SpillSlotIndex,
+    AllocRegResult, Env, LiveRangeKey, PRegIndex, RegTraversalIter, SpillSetIndex, SpillSlotData,
+    SpillSlotIndex,
 };
 use crate::{Allocation, Function, SpillSlot};
 
@@ -40,9 +40,7 @@ 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, hint, PReg::invalid(), bundle.index(), None)
-            {
+            for preg in RegTraversalIter::new(self.env, class, None, hint, bundle.index()) {
                 trace!("trying bundle {:?} to preg {:?}", bundle, preg);
                 let preg_idx = PRegIndex::new(preg.index());
                 if let AllocRegResult::Allocated(_) =