Store PRegSet in MachineEnv

src/fastalloc/lru.rs

@@ -28,7 +28,8 @@ pub struct LruNode {
 }
 
 impl Lru {
-    pub fn new(regclass: RegClass, regs: &[PReg]) -> Self {
+    pub fn new(regclass: RegClass, regs: &PRegSet) -> Self {
+        let regs = regs.into_iter().collect::<Vec<_>>();
         let mut data = vec![
             LruNode {
                 prev: u8::MAX,
@@ -248,7 +249,7 @@ impl fmt::Debug for Lru {
             while node != self.head {
                 if seen.contains(&node) {
                     panic!(
-                        "The {:?} LRU is messed up: 
+                        "The {:?} LRU is messed up:
                        head: {:?}, {:?} -> p{node}, actual data: {:?}",
                         self.regclass, self.head, data_str, self.data
                     );
@@ -298,7 +299,7 @@ impl<T: PartialEq> PartialEq for PartedByRegClass<T> {
 pub type Lrus = PartedByRegClass<Lru>;
 
 impl Lrus {
-    pub fn new(int_regs: &[PReg], float_regs: &[PReg], vec_regs: &[PReg]) -> Self {
+    pub fn new(int_regs: &PRegSet, float_regs: &PRegSet, vec_regs: &PRegSet) -> Self {
         Self {
             items: [
                 Lru::new(RegClass::Int, int_regs),

src/fastalloc/mod.rs

@@ -444,21 +444,9 @@ impl<'a, F: Function> Env<'a, F> {
             env.preferred_regs_by_class[RegClass::Float as usize].clone(),
             env.preferred_regs_by_class[RegClass::Vector as usize].clone(),
         ];
-        regs[0].extend(
-            env.non_preferred_regs_by_class[RegClass::Int as usize]
-                .iter()
-                .cloned(),
-        );
-        regs[1].extend(
-            env.non_preferred_regs_by_class[RegClass::Float as usize]
-                .iter()
-                .cloned(),
-        );
-        regs[2].extend(
-            env.non_preferred_regs_by_class[RegClass::Vector as usize]
-                .iter()
-                .cloned(),
-        );
+        regs[0].union_from(env.non_preferred_regs_by_class[RegClass::Int as usize]);
+        regs[1].union_from(env.non_preferred_regs_by_class[RegClass::Float as usize]);
+        regs[2].union_from(env.non_preferred_regs_by_class[RegClass::Vector as usize]);
         let allocatable_regs = PRegSet::from(env);
         let num_available_pregs: PartedByRegClass<i16> = PartedByRegClass {
             items: [
@@ -508,9 +496,9 @@ impl<'a, F: Function> Env<'a, F> {
             ],
             preferred_victim: PartedByRegClass {
                 items: [
-                    regs[0].last().cloned().unwrap_or(PReg::invalid()),
-                    regs[1].last().cloned().unwrap_or(PReg::invalid()),
-                    regs[2].last().cloned().unwrap_or(PReg::invalid()),
+                    regs[0].max_preg().unwrap_or(PReg::invalid()),
+                    regs[1].max_preg().unwrap_or(PReg::invalid()),
+                    regs[2].max_preg().unwrap_or(PReg::invalid()),
                 ],
             },
             reused_input_to_reuse_op: vec![usize::MAX; max_operand_len as usize],

src/fastalloc/tests.rs

@@ -146,10 +146,10 @@ fn mach_env(no_of_regs: usize) -> MachineEnv {
             (0..no_of_regs)
                 .map(|no| PReg::new(no, RegClass::Int))
                 .collect(),
-            vec![],
-            vec![],
+            PRegSet::empty(),
+            PRegSet::empty(),
         ],
-        non_preferred_regs_by_class: [vec![], vec![], vec![]],
+        non_preferred_regs_by_class: [PRegSet::empty(); 3],
         scratch_by_class: [None, None, None],
         fixed_stack_slots: vec![],
     }

src/fuzzing/func.rs

@@ -696,20 +696,20 @@ impl core::fmt::Debug for Func {
 }
 
 pub fn machine_env() -> MachineEnv {
-    fn regs(r: core::ops::Range<usize>, c: RegClass) -> Vec<PReg> {
+    fn regs(r: core::ops::Range<usize>, c: RegClass) -> PRegSet {
         r.map(|i| PReg::new(i, c)).collect()
     }
-    let preferred_regs_by_class: [Vec<PReg>; 3] = [
+    let preferred_regs_by_class = [
         regs(0..24, RegClass::Int),
         regs(0..24, RegClass::Float),
         regs(0..24, RegClass::Vector),
     ];
-    let non_preferred_regs_by_class: [Vec<PReg>; 3] = [
+    let non_preferred_regs_by_class = [
         regs(24..32, RegClass::Int),
         regs(24..32, RegClass::Float),
         regs(24..32, RegClass::Vector),
     ];
-    let scratch_by_class: [Option<PReg>; 3] = [None, None, None];
+    let scratch_by_class = [None, None, None];
     let fixed_stack_slots = (32..63)
         .flat_map(|i| {
             [

src/ion/liveranges.rs

@@ -115,9 +115,8 @@ impl<'a, F: Function> Env<'a, F> {
         }
         for class in 0..self.preferred_victim_by_class.len() {
             self.preferred_victim_by_class[class] = self.env.non_preferred_regs_by_class[class]
-                .last()
-                .or(self.env.preferred_regs_by_class[class].last())
-                .cloned()
+                .max_preg()
+                .or(self.env.preferred_regs_by_class[class].max_preg())
                 .unwrap_or(PReg::invalid());
         }
         // Create VRegs from the vreg count.

src/ion/process.rs

@@ -1219,8 +1219,8 @@ impl<'a, F: Function> Env<'a, F> {
                     let mut fixed_assigned = 0;
                     let mut total_regs = 0;
                     for preg in self.env.preferred_regs_by_class[class as u8 as usize]
-                        .iter()
-                        .chain(self.env.non_preferred_regs_by_class[class as u8 as usize].iter())
+                        .into_iter()
+                        .chain(self.env.non_preferred_regs_by_class[class as u8 as usize])
                     {
                         trace!(" -> PR {:?}", preg);
                         let start = LiveRangeKey::from_range(&CodeRange {

src/ion/reg_traversal.rs

@@ -1,53 +1,28 @@
 //! Iterate over available registers.
 
-use crate::{MachineEnv, PReg, RegClass};
+use crate::{MachineEnv, PReg, PRegSet, PRegSetIter, RegClass};
 
 /// Keep track of where we are in the register traversal.
-struct Cursor<'a> {
-    registers: &'a [PReg],
-    index: usize,
-    offset: usize,
+struct Cursor {
+    first: PRegSetIter,
+    second: PRegSetIter,
 }
 
-impl<'a> Cursor<'a> {
+impl Cursor {
     #[inline]
-    fn new(registers: &'a [PReg], offset_hint: usize) -> Self {
-        let offset = if registers.len() > 0 {
-            offset_hint % registers.len()
-        } else {
-            0
-        };
+    fn new(registers: &PRegSet, class: RegClass, offset_hint: usize) -> Self {
+        let mut mask = PRegSet::empty();
+        mask.add_up_to(PReg::new(offset_hint % PReg::MAX, class));
+        let first = *registers & mask.invert();
+        let second = *registers & mask;
         Self {
-            registers,
-            index: 0,
-            offset,
+            first: first.into_iter(),
+            second: second.into_iter(),
         }
     }
 
-    /// 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()
+    fn next(&mut self) -> Option<PReg> {
+        self.first.next().or_else(|| self.second.next())
     }
 }
 
@@ -65,19 +40,19 @@ impl<'a> Cursor<'a> {
 ///   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> {
+pub struct RegTraversalIter {
     is_fixed: bool,
     fixed: Option<PReg>,
     use_hint: bool,
     hint: Option<PReg>,
-    preferred: Cursor<'a>,
-    non_preferred: Cursor<'a>,
+    preferred: Cursor,
+    non_preferred: Cursor,
     limit: Option<usize>,
 }
 
-impl<'a> RegTraversalIter<'a> {
+impl RegTraversalIter {
     pub fn new(
-        env: &'a MachineEnv,
+        env: &MachineEnv,
         class: RegClass,
         fixed: Option<PReg>,
         hint: Option<PReg>,
@@ -87,9 +62,10 @@ impl<'a> RegTraversalIter<'a> {
         debug_assert!(fixed != Some(PReg::invalid()));
         debug_assert!(hint != Some(PReg::invalid()));
 
-        let class = class as u8 as usize;
-        let preferred = Cursor::new(&env.preferred_regs_by_class[class], offset);
-        let non_preferred = Cursor::new(&env.non_preferred_regs_by_class[class], offset);
+        let class_index = class as u8 as usize;
+        let preferred = Cursor::new(&env.preferred_regs_by_class[class_index], class, offset);
+        let non_preferred =
+            Cursor::new(&env.non_preferred_regs_by_class[class_index], class, offset);
 
         Self {
             is_fixed: fixed.is_some(),
@@ -103,7 +79,7 @@ impl<'a> RegTraversalIter<'a> {
     }
 }
 
-impl<'a> core::iter::Iterator for RegTraversalIter<'a> {
+impl core::iter::Iterator for RegTraversalIter {
     type Item = PReg;
 
     fn next(&mut self) -> Option<PReg> {
@@ -118,16 +94,14 @@ impl<'a> core::iter::Iterator for RegTraversalIter<'a> {
             }
         }
 
-        while !self.preferred.done() {
-            let reg = self.preferred.advance();
+        while let Some(reg) = self.preferred.next() {
             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();
+        while let Some(reg) = self.non_preferred.next() {
             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.
             }