PregSet machine env

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: [Vec<PReg>; 3] = [
-        regs(0..24, RegClass::Int),
-        regs(0..24, RegClass::Float),
-        regs(0..24, RegClass::Vector),
-    ];
-    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 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,10 +113,12 @@ 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())
-                .cloned()
+            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());
         }
         // Create VRegs from the vreg count.

src/ion/process.rs

@@ -1221,9 +1221,17 @@ 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]
-                        .iter()
-                        .chain(self.env.non_preferred_regs_by_class[class as u8 as usize].iter())
+                    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
+                                .to_preg_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};
-
+use crate::{MachineEnv, PReg, PRegClass, RegClass};
 /// This iterator represents a traversal through all allocatable
 /// registers of a given class, in a certain order designed to
 /// minimize allocation contention.
@@ -14,9 +13,9 @@ use crate::{MachineEnv, PReg, RegClass};
 ///   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,
+pub struct RegTraversalIter {
+    pref_regs_by_class: PRegClass,
+    non_pref_regs_by_class: PRegClass,
     hints: [Option<PReg>; 2],
     hint_idx: usize,
     pref_idx: usize,
@@ -27,9 +26,9 @@ pub struct RegTraversalIter<'a> {
     fixed: Option<PReg>,
 }
 
-impl<'a> RegTraversalIter<'a> {
+impl RegTraversalIter {
     pub fn new(
-        env: &'a MachineEnv,
+        env: &MachineEnv,
         class: RegClass,
         hint_reg: PReg,
         hint2_reg: PReg,
@@ -53,19 +52,23 @@ 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 pref_regs_by_class = env.preferred_regs_by_class.to_preg_class(class);
+        let non_pref_regs_by_class = env.non_preferred_regs_by_class.to_preg_class(class);
+
+        let offset_pref = if pref_regs_by_class.len() > 0 {
+            offset % pref_regs_by_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()
+        let offset_non_pref = if non_pref_regs_by_class.len() > 0 {
+            offset % non_pref_regs_by_class.len()
         } else {
             0
         };
         Self {
-            env,
-            class,
+            pref_regs_by_class,
+            non_pref_regs_by_class,
             hints,
             hint_idx: 0,
             pref_idx: 0,
@@ -78,46 +81,61 @@ 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> {
+        // only take the fixed register if it exists
         if self.is_fixed {
             let ret = self.fixed;
             self.fixed = None;
             return ret;
         }
 
-        fn wrap(idx: usize, limit: usize) -> usize {
-            if idx >= limit {
-                idx - limit
-            } else {
-                idx
-            }
-        }
+        // if there are hints, return them first
         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())];
+
+        // iterate over the preferred register rotated by offset
+        // ignoring hint register
+        let n_pref_regs = self.pref_regs_by_class.len();
+        while self.pref_idx < n_pref_regs {
+            let mut arr = self.pref_regs_by_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;
         }
-        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())];
+
+        // iterate over the nonpreferred register rotated by offset
+        // ignoring hint register
+        let n_non_pref_regs = self.non_pref_regs_by_class.len();
+        while self.non_pref_idx < n_non_pref_regs {
+            let mut arr = self.non_pref_regs_by_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
     }
 }
+
+/// Wrapping function to wrap around the index for an iterator
+fn wrap(idx: usize, limit: usize) -> usize {
+    if idx >= limit {
+        idx - limit
+    } else {
+        idx
+    }
+}