PregSet machine env

src/fuzzing/func.rs

@@ -645,15 +645,15 @@ 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: [Vec<PReg>; 3] = [
-        regs(0..24, RegClass::Int),
-        regs(0..24, RegClass::Float),
-        regs(0..24, RegClass::Vector),
+    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: [Vec<PReg>; 3] = [
-        regs(24..32, RegClass::Int),
-        regs(24..32, RegClass::Float),
-        regs(24..32, RegClass::Vector),
+    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 scratch_by_class: [Option<PReg>; 3] = [None, None, None];
     let fixed_stack_slots = (32..63)

src/ion/liveranges.rs

@@ -116,7 +116,7 @@ impl<'a, F: Function> Env<'a, F> {
             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()
+                // .cloned()
                 .unwrap_or(PReg::invalid());
         }
         // Create VRegs from the vreg count.

src/ion/process.rs

@@ -1222,8 +1222,10 @@ 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].into_iter(),
+                        )
                     {
                         trace!(" -> PR {:?}", preg);
                         let start = LiveRangeKey::from_range(&CodeRange {

src/ion/reg_traversal.rs

@@ -1,5 +1,4 @@
 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.
@@ -100,23 +99,30 @@ impl<'a> core::iter::Iterator for RegTraversalIter<'a> {
             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())];
+
+        let n_pref_regs = self.env.preferred_regs_by_class[self.class].len();
+        while self.pref_idx < n_pref_regs {
+            let mut arr = self.env.preferred_regs_by_class[self.class].into_iter();
+            let r = arr.nth(wrap(self.pref_idx + self.offset_pref, n_pref_regs));
             self.pref_idx += 1;
-            if Some(r) == self.hints[0] || Some(r) == self.hints[1] {
+            if r == self.hints[0] || r == self.hints[1] {
                 continue;
             }
-            return Some(r);
+            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 arr = &self.env.non_preferred_regs_by_class[self.class][..];
-            let r = arr[wrap(self.non_pref_idx + self.offset_non_pref, arr.len())];
+            let mut arr = self.env.non_preferred_regs_by_class[self.class].into_iter();
+            let r = arr.nth(wrap(
+                self.non_pref_idx + self.offset_non_pref,
+                n_non_pref_regs,
+            ));
             self.non_pref_idx += 1;
-            if Some(r) == self.hints[0] || Some(r) == self.hints[1] {
+            if r == self.hints[0] || r == self.hints[1] {
                 continue;
             }
-            return Some(r);
+            return r;
         }
         None
     }

src/lib.rs

@@ -240,21 +240,14 @@ impl PRegSet {
         self.bits[0] |= other.bits[0];
         self.bits[1] |= other.bits[1];
     }
-}
-
-impl IntoIterator for PRegSet {
-    type Item = PReg;
-    type IntoIter = PRegSetIter;
-    fn into_iter(self) -> PRegSetIter {
-        PRegSetIter { bits: self.bits }
-    }
-}
 
-pub struct PRegSetIter {
-    bits: [u128; 2],
+    // Get the number of registers in the set
+    // pub fn n_regs(&self) -> usize {
+    //     (self.bits[0].count_ones() + self.bits[1].count_ones()) as usize
+    // }
 }
 
-impl Iterator for PRegSetIter {
+impl Iterator for PRegSet {
     type Item = PReg;
     fn next(&mut self) -> Option<PReg> {
         if self.bits[0] != 0 {
@@ -269,22 +262,49 @@ impl Iterator for PRegSetIter {
             None
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let len = (self.bits[0].count_ones() + self.bits[1].count_ones()) as usize;
+        (len, Some(len))
+    }
+
+    fn last(self) -> Option<PReg> {
+        if self.bits[1] != 0 {
+            let index = 127 - self.bits[1].leading_zeros();
+            Some(PReg::from_index(index as usize + 128))
+        } else if self.bits[0] != 0 {
+            let index = self.bits[0].leading_zeros();
+            Some(PReg::from_index(index as usize))
+        } else {
+            None
+        }
+    }
 }
 
+impl ExactSizeIterator for PRegSet {}
+
 impl From<&MachineEnv> for PRegSet {
     fn from(env: &MachineEnv) -> Self {
         let mut res = Self::default();
 
         for class in env.preferred_regs_by_class.iter() {
-            for preg in class {
-                res.add(*preg)
-            }
+            res.union_from(*class);
         }
 
         for class in env.non_preferred_regs_by_class.iter() {
-            for preg in class {
-                res.add(*preg)
-            }
+            res.union_from(*class);
+        }
+
+        res
+    }
+}
+
+impl From<Vec<PReg>> for PRegSet {
+    fn from(regs: Vec<PReg>) -> Self {
+        let mut res = Self::default();
+
+        for preg in regs {
+            res.add(preg);
         }
 
         res
@@ -1380,7 +1400,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: [Vec<PReg>; 3],
+    pub preferred_regs_by_class: [PRegSet; 3],
 
     /// Non-preferred physical registers for each class. These are the
     /// registers that will be allocated if a preferred register is
@@ -1389,7 +1409,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: [Vec<PReg>; 3],
+    pub non_preferred_regs_by_class: [PRegSet; 3],
 
     /// Optional dedicated scratch register per class. This is needed to perform
     /// moves between registers when cyclic move patterns occur. The
@@ -1546,3 +1566,15 @@ pub struct RegallocOptions {
     /// Run the SSA validator before allocating registers.
     pub validate_ssa: bool,
 }
+
+#[cfg(test)]
+mod test {
+    use super::PRegSet;
+
+    #[test]
+    fn test_set_bits_iter() {
+        let registers = PRegSet { bits: [112, 131] }.into_iter();
+        let last = registers.last().unwrap().bits;
+        assert_eq!(last, 135);
+    }
+}