Store PRegSet in MachineEnv

Helps to make the `MachineEnv`s in Wasmtime `const`-allocatable.

Author: alexcrichton

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,31 @@
 //! 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
-        };
-        Self {
-            registers,
-            index: 0,
-            offset,
-        }
-    }
-
-    /// 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
-        }
+    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 }
     }
 
-    /// 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 +43,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>,
@@ -103,7 +81,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 +96,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.
             }

src/lib.rs

@@ -293,6 +293,16 @@ impl PRegSet {
     pub fn is_empty(&self, regclass: RegClass) -> bool {
         self.bits[regclass as usize] == 0
     }
+
+    /// Returns the number of register in this set.
+    pub fn len(&self) -> u32 {
+        self.bits.iter().map(|s| s.count_ones()).sum()
+    }
+
+    /// Returns the maximum register in this set, with the highest hw_enc value.
+    pub fn max_preg(&self) -> Option<PReg> {
+        self.into_iter().last()
+    }
 }
 
 impl core::ops::BitAnd<PRegSet> for PRegSet {
@@ -315,6 +325,14 @@ impl core::ops::BitOr<PRegSet> for PRegSet {
     }
 }
 
+impl IntoIterator for &PRegSet {
+    type Item = PReg;
+    type IntoIter = PRegSetIter;
+    fn into_iter(self) -> PRegSetIter {
+        (*self).into_iter()
+    }
+}
+
 impl IntoIterator for PRegSet {
     type Item = PReg;
     type IntoIter = PRegSetIter;
@@ -352,15 +370,11 @@ impl From<&MachineEnv> for PRegSet {
         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
@@ -1483,7 +1497,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
@@ -1492,7 +1506,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
@@ -1774,3 +1788,37 @@ unsafe impl allocator_api2::alloc::Allocator for Bump {
         self.0.deref().shrink(ptr, old_layout, new_layout)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{PReg, PRegSet, RegClass::Int};
+
+    #[test]
+    fn preg_set_len() {
+        let mut set = PRegSet::empty();
+        assert_eq!(set.len(), 0);
+
+        set.add(PReg::new(3, Int));
+        assert_eq!(set.len(), 1);
+        set.add(PReg::new(3, Int));
+        assert_eq!(set.len(), 1);
+
+        set.add(PReg::new(4, Int));
+        assert_eq!(set.len(), 2);
+    }
+
+    #[test]
+    fn preg_set_max_preg() {
+        let mut set = PRegSet::empty();
+        assert_eq!(set.max_preg(), None);
+
+        set.add(PReg::new(3, Int));
+        assert_eq!(set.max_preg(), Some(PReg::new(3, Int)));
+
+        set.add(PReg::new(4, Int));
+        assert_eq!(set.max_preg(), Some(PReg::new(4, Int)));
+
+        set.add(PReg::new(2, Int));
+        assert_eq!(set.max_preg(), Some(PReg::new(4, Int)));
+    }
+}