Store PRegSet in MachineEnv

src/ion/reg_traversal.rs

@@ -10,18 +10,15 @@ struct Cursor {
 
 impl Cursor {
     #[inline]
-    fn new(registers: &PRegSet, offset_hint: usize) -> Self {
-        let first = registers
-            .into_iter()
-            .skip(offset_hint)
-            .collect::<PRegSet>()
-            .into_iter();
-        let second = registers
-            .into_iter()
-            .take(offset_hint)
-            .collect::<PRegSet>()
-            .into_iter();
-        Self { first, second }
+    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 {
+            first: first.into_iter(),
+            second: second.into_iter(),
+        }
     }
 
     fn next(&mut self) -> Option<PReg> {
@@ -65,9 +62,10 @@ impl RegTraversalIter {
         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(),

src/lib.rs

@@ -303,6 +303,15 @@ impl PRegSet {
     pub fn max_preg(&self) -> Option<PReg> {
         self.into_iter().last()
     }
+
+    /// Add all registers from `0..reg` to this set, not including `reg` itself.
+    pub fn add_up_to(&mut self, reg: PReg) {
+        let (index, bit) = Self::split_index(reg);
+        for i in 0..index {
+            self.bits[i] = !0;
+        }
+        self.bits[index] = (1 << bit) - 1;
+    }
 }
 
 impl core::ops::BitAnd<PRegSet> for PRegSet {
@@ -1821,4 +1830,31 @@ mod tests {
         set.add(PReg::new(2, Int));
         assert_eq!(set.max_preg(), Some(PReg::new(4, Int)));
     }
+
+    #[test]
+    fn preg_set_new_up_to() {
+        let p0 = PReg::new(0, Int);
+        let p1 = PReg::new(1, Int);
+        let p2 = PReg::new(2, Int);
+        let p3 = PReg::new(3, Int);
+        {
+            let mut set = PRegSet::empty();
+            set.add_up_to(p1);
+            assert!(set.contains(p0));
+            assert!(!set.contains(p1));
+        }
+        {
+            let mut set = PRegSet::empty();
+            set.add_up_to(p0);
+            assert!(!set.contains(p0));
+        }
+        {
+            let mut set = PRegSet::empty();
+            set.add_up_to(p3);
+            assert!(set.contains(p0));
+            assert!(set.contains(p1));
+            assert!(set.contains(p2));
+            assert!(!set.contains(p3));
+        }
+    }
 }