xref: /wasmtime-44.0.1/winch/codegen/src/regset.rs (revision 14b39bc2)
1 use crate::isa::reg::Reg;
2 use regalloc2::RegClass;
3 
4 /// A bit set to track regiter availability.
5 pub(crate) struct RegSet {
6     /// Bitset to track general purpose register availability.
7     gpr: u32,
8     /// Bitset to track floating-point register availability.
9     fpr: u32,
10 }
11 
12 use std::ops::{Index, IndexMut};
13 
14 impl Index<RegClass> for RegSet {
15     type Output = u32;
16 
17     fn index(&self, class: RegClass) -> &Self::Output {
18         match class {
19             RegClass::Int => &self.gpr,
20             RegClass::Float => &self.fpr,
21             c => unreachable!("Unexpected register class {:?}", c),
22         }
23     }
24 }
25 
26 impl IndexMut<RegClass> for RegSet {
27     fn index_mut(&mut self, class: RegClass) -> &mut Self::Output {
28         match class {
29             RegClass::Int => &mut self.gpr,
30             RegClass::Float => &mut self.fpr,
31             c => unreachable!("Unexpected register class {:?}", c),
32         }
33     }
34 }
35 
36 impl RegSet {
37     /// Create a new register set.
38     pub fn new(gpr: u32, fpr: u32) -> Self {
39         Self { gpr, fpr }
40     }
41 
42     /// Allocate the next available register of the given class,
43     /// returning `None` if there are no more registers available.
44     pub fn reg_for_class(&mut self, class: RegClass) -> Option<Reg> {
45         self.available(class).then(|| {
46             let bitset = self[class];
47             let index = bitset.trailing_zeros();
48             self.allocate(class, index);
49             Reg::from(class, index as usize)
50         })
51     }
52 
53     /// Request a specific register.
54     pub fn reg(&mut self, reg: Reg) -> Option<Reg> {
55         let index = reg.hw_enc();
56         self.named_reg_available(reg).then(|| {
57             self.allocate(reg.class(), index.into());
58             reg
59         })
60     }
61 
62     /// Marks the specified register as available, utilizing the
63     /// register class to determine the bitset that requires updating.
64     pub fn free(&mut self, reg: Reg) {
65         let index = reg.hw_enc() as u32;
66         self[reg.class()] |= 1 << index;
67     }
68 
69     /// Returns true if the specified register is allocatable.
70     pub fn named_reg_available(&self, reg: Reg) -> bool {
71         let bitset = self[reg.class()];
72         let index = 1 << reg.hw_enc();
73         (!bitset & index) == 0
74     }
75 
76     fn available(&self, class: RegClass) -> bool {
77         let bitset = self[class];
78         bitset != 0
79     }
80 
81     fn allocate(&mut self, class: RegClass, index: u32) {
82         self[class] &= !(1 << index);
83     }
84 }
85 
86 #[cfg(test)]
87 mod tests {
88     use super::{Reg, RegSet};
89     use regalloc2::RegClass;
90 
91     const UNIVERSE: u32 = (1 << 16) - 1;
92 
93     #[test]
94     fn test_any_gpr() {
95         let mut set = RegSet::new(UNIVERSE, 0);
96         for _ in 0..16 {
97             let gpr = set.reg_for_class(RegClass::Int);
98             assert!(gpr.is_some())
99         }
100 
101         assert!(!set.available(RegClass::Int));
102         assert!(set.reg_for_class(RegClass::Int).is_none())
103     }
104 
105     #[test]
106     fn test_gpr() {
107         let all = UNIVERSE & !(1 << 5);
108         let target = Reg::int(5);
109         let mut set = RegSet::new(all, 0);
110         assert!(set.reg(target).is_none());
111     }
112 
113     #[test]
114     fn test_free_reg() {
115         let mut set = RegSet::new(UNIVERSE, 0);
116         let gpr = set.reg_for_class(RegClass::Int).unwrap();
117         set.free(gpr);
118         assert!(set.reg(gpr).is_some());
119     }
120 }
121