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