only use one PRegSet

src/fuzzing/func.rs

@@ -645,16 +645,24 @@ pub fn machine_env() -> MachineEnv {
     fn regs(r: core::ops::Range<usize>, c: RegClass) -> Vec<PReg> {
         r.map(|i| PReg::new(i, c)).collect()
     }
-    let preferred_regs_by_class: [PRegSet; 3] = [
-        regs(0..24, RegClass::Int).into(),
-        regs(0..24, RegClass::Float).into(),
-        regs(0..24, RegClass::Vector).into(),
-    ];
-    let non_preferred_regs_by_class: [PRegSet; 3] = [
-        regs(24..32, RegClass::Int).into(),
-        regs(24..32, RegClass::Float).into(),
-        regs(24..32, RegClass::Vector).into(),
-    ];
+    let int_regs: PRegSet = regs(0..24, RegClass::Int).into();
+    let float_regs: PRegSet = regs(0..24, RegClass::Float).into();
+    let vector_regs: PRegSet = regs(0..24, RegClass::Vector).into();
+
+    let mut preferred_regs_by_class = PRegSet::default();
+    preferred_regs_by_class.union_from(int_regs);
+    preferred_regs_by_class.union_from(float_regs);
+    preferred_regs_by_class.union_from(vector_regs);
+
+    let int_regs: PRegSet = regs(24..32, RegClass::Int).into();
+    let float_regs: PRegSet = regs(24..32, RegClass::Float).into();
+    let vector_regs: PRegSet = regs(24..32, RegClass::Vector).into();
+
+    let mut non_preferred_regs_by_class = PRegSet::default();
+    non_preferred_regs_by_class.union_from(int_regs);
+    non_preferred_regs_by_class.union_from(float_regs);
+    non_preferred_regs_by_class.union_from(vector_regs);
+
     let scratch_by_class: [Option<PReg>; 3] = [None, None, None];
     let fixed_stack_slots = (32..63)
         .flat_map(|i| {

src/ion/liveranges.rs

@@ -113,9 +113,11 @@ impl<'a, F: Function> Env<'a, F> {
             self.pregs[preg.index()].is_stack = true;
         }
         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())
+            self.preferred_victim_by_class[class] = self
+                .env
+                .non_preferred_regs_by_class
+                .last_in_class(class)
+                .or(self.env.preferred_regs_by_class.last_in_class(class))
                 // .cloned()
                 .unwrap_or(PReg::invalid());
         }

src/ion/process.rs

@@ -1221,10 +1221,16 @@ impl<'a, F: Function> Env<'a, F> {
                     let mut min_bundles_assigned = 0;
                     let mut fixed_assigned = 0;
                     let mut total_regs = 0;
-                    for preg in self.env.preferred_regs_by_class[class as u8 as usize]
+                    for preg in self
+                        .env
+                        .preferred_regs_by_class
+                        .to_preg_class(class as u8 as usize)
                         .into_iter()
                         .chain(
-                            self.env.non_preferred_regs_by_class[class as u8 as usize].into_iter(),
+                            self.env
+                                .non_preferred_regs_by_class
+                                .to_preg_class(class as u8 as usize)
+                                .into_iter(),
                         )
                     {
                         trace!(" -> PR {:?}", preg);

src/ion/reg_traversal.rs

@@ -52,13 +52,13 @@ impl<'a> RegTraversalIter<'a> {
         }
         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()
+        let offset_pref = if env.preferred_regs_by_class.len_class(class) > 0 {
+            offset % env.preferred_regs_by_class.len_class(class)
         } 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()
+        let offset_non_pref = if env.non_preferred_regs_by_class.len_class(class) > 0 {
+            offset % env.non_preferred_regs_by_class.len_class(class)
         } else {
             0
         };
@@ -100,9 +100,13 @@ impl<'a> core::iter::Iterator for RegTraversalIter<'a> {
             return h;
         }
 
-        let n_pref_regs = self.env.preferred_regs_by_class[self.class].len();
+        let n_pref_regs = self.env.preferred_regs_by_class.len_class(self.class);
         while self.pref_idx < n_pref_regs {
-            let mut arr = self.env.preferred_regs_by_class[self.class].into_iter();
+            let mut arr = self
+                .env
+                .preferred_regs_by_class
+                .to_preg_class(self.class)
+                .into_iter();
             let r = arr.nth(wrap(self.pref_idx + self.offset_pref, n_pref_regs));
             self.pref_idx += 1;
             if r == self.hints[0] || r == self.hints[1] {
@@ -111,9 +115,13 @@ impl<'a> core::iter::Iterator for RegTraversalIter<'a> {
             return r;
         }
 
-        let n_non_pref_regs = self.env.non_preferred_regs_by_class[self.class].len();
-        while self.non_pref_idx < self.env.non_preferred_regs_by_class[self.class].len() {
-            let mut arr = self.env.non_preferred_regs_by_class[self.class].into_iter();
+        let n_non_pref_regs = self.env.non_preferred_regs_by_class.len_class(self.class);
+        while self.non_pref_idx < n_non_pref_regs {
+            let mut arr = self
+                .env
+                .non_preferred_regs_by_class
+                .to_preg_class(self.class)
+                .into_iter();
             let r = arr.nth(wrap(
                 self.non_pref_idx + self.offset_non_pref,
                 n_non_pref_regs,

src/lib.rs

@@ -248,10 +248,29 @@ impl PRegSet {
         self.bits[2] |= other.bits[2];
     }
 
-    pub fn last_in_class(&self, class: usize) -> PReg {
-        let i_in_class = 63 - self.bits[class].leading_zeros();
-        let i = i_in_class as usize | (class << 6);
-        PReg::from_index(i)
+    pub fn last_in_class(&self, class: usize) -> Option<PReg> {
+        if self.bits[class] == 0 {
+            None
+        } else {
+            // get the index in the class vec
+            let i_in_class = 63 - self.bits[class].leading_zeros();
+            // apply the class mask
+            let i = i_in_class as u8 | ((class as u8) << 6);
+            // return the PReg
+            Some(PReg::from(i))
+        }
+    }
+
+    pub fn len_class(&self, class: usize) -> usize {
+        self.bits[class].count_ones() as usize
+    }
+
+    /// Get a iterator over only one class in the set
+    fn to_preg_class(&self, class: usize) -> PRegClass {
+        PRegClass {
+            class_mask: (class as u8) << 6,
+            regs: self.bits[class],
+        }
     }
 }
 
@@ -287,10 +306,10 @@ impl Iterator for PRegSet {
             let index = 63 - self.bits[2].leading_zeros();
             Some(PReg::from_index(index as usize + 128))
         } else if self.bits[1] != 0 {
-            let index = self.bits[1].leading_zeros();
+            let index = 63 - self.bits[1].leading_zeros();
             Some(PReg::from_index(index as usize + 64))
         } else if self.bits[0] != 0 {
-            let index = self.bits[0].leading_zeros();
+            let index = 63 - self.bits[0].leading_zeros();
             Some(PReg::from_index(index as usize))
         } else {
             None
@@ -304,13 +323,9 @@ impl From<&MachineEnv> for PRegSet {
     fn from(env: &MachineEnv) -> Self {
         let mut res = Self::default();
 
-        for class in env.preferred_regs_by_class.iter() {
-            res.union_from(*class);
-        }
+        res.union_from(env.preferred_regs_by_class);
 
-        for class in env.non_preferred_regs_by_class.iter() {
-            res.union_from(*class);
-        }
+        res.union_from(env.non_preferred_regs_by_class);
 
         res
     }
@@ -328,6 +343,45 @@ impl From<Vec<PReg>> for PRegSet {
     }
 }
 
+/// A compact iterator over a single register class
+struct PRegClass {
+    // bit mask containing class data in UPPER two bits
+    class_mask: u8,
+    // bit packed vec of registers
+    regs: u64,
+}
+
+impl Iterator for PRegClass {
+    type Item = PReg;
+    fn next(&mut self) -> Option<PReg> {
+        if self.regs == 0 {
+            None
+        } else {
+            let index = self.regs.trailing_zeros() as u8;
+            self.regs &= !(1u64 << index);
+            let reg_index = index as u8 | self.class_mask;
+            Some(PReg::from_index(reg_index as usize))
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = (self.regs.count_ones()) as usize;
+        (len, Some(len))
+    }
+
+    fn last(self) -> Option<PReg> {
+        if self.regs == 0 {
+            None
+        } else {
+            let index = 63 - self.regs.leading_zeros();
+            let reg_index = index as u8 | self.class_mask;
+            Some(PReg::from_index(reg_index as usize))
+        }
+    }
+}
+
+impl ExactSizeIterator for PRegClass {}
+
 /// A virtual register. Contains a virtual register number and a
 /// class.
 ///
@@ -1417,7 +1471,7 @@ pub struct MachineEnv {
     ///
     /// If an explicit scratch register is provided in `scratch_by_class` then
     /// it must not appear in this list.
-    pub preferred_regs_by_class: [PRegSet; 3],
+    pub preferred_regs_by_class: PRegSet,
 
     /// Non-preferred physical registers for each class. These are the
     /// registers that will be allocated if a preferred register is
@@ -1426,7 +1480,7 @@ pub struct MachineEnv {
     ///
     /// If an explicit scratch register is provided in `scratch_by_class` then
     /// it must not appear in this list.
-    pub non_preferred_regs_by_class: [PRegSet; 3],
+    pub non_preferred_regs_by_class: PRegSet,
 
     /// Optional dedicated scratch register per class. This is needed to perform
     /// moves between registers when cyclic move patterns occur. The