xref: /wasmtime-44.0.1/winch/codegen/src/regset.rs (revision af4d94c8)
1 use crate::isa::reg::Reg;
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 impl RegSet {
12     /// Create a new register set.
13     pub fn new(gpr: u32, fpr: u32) -> Self {
14         Self { gpr, _fpr: fpr }
15     }
16 
17     /// Request a general purpose register.
18     pub fn any_gpr(&mut self) -> Option<Reg> {
19         self.gpr_available().then(|| {
20             let index = self.gpr.trailing_zeros();
21             self.allocate(index);
22             Reg::int(index as usize)
23         })
24     }
25 
26     /// Request a specific general purpose register.
27     pub fn gpr(&mut self, reg: Reg) -> Option<Reg> {
28         let index = reg.hw_enc();
29         self.named_gpr_available(index as u32).then(|| {
30             self.allocate(index as u32);
31             Reg::int(index as usize)
32         })
33     }
34 
35     /// Free the given general purpose register.
36     pub fn free_gpr(&mut self, reg: Reg) {
37         let index = reg.hw_enc() as u32;
38         self.gpr |= 1 << index;
39     }
40 
41     /// Returns true if the given general purpose register
42     /// is available.
43     pub fn named_gpr_available(&self, index: u32) -> bool {
44         let index = 1 << index;
45         (!self.gpr & index) == 0
46     }
47 
48     fn gpr_available(&self) -> bool {
49         self.gpr != 0
50     }
51 
52     fn allocate(&mut self, index: u32) {
53         self.gpr &= !(1 << index);
54     }
55 }
56 
57 #[cfg(test)]
58 mod tests {
59     use super::{Reg, RegSet};
60 
61     const UNIVERSE: u32 = (1 << 16) - 1;
62 
63     #[test]
64     fn test_any_gpr() {
65         let mut set = RegSet::new(UNIVERSE, 0);
66         for _ in 0..16 {
67             let gpr = set.any_gpr();
68             assert!(gpr.is_some())
69         }
70 
71         assert!(!set.gpr_available());
72         assert!(set.any_gpr().is_none())
73     }
74 
75     #[test]
76     fn test_gpr() {
77         let all = UNIVERSE & !(1 << 5);
78         let target = Reg::int(5);
79         let mut set = RegSet::new(all, 0);
80         assert!(set.gpr(target).is_none());
81     }
82 
83     #[test]
84     fn test_free_gpr() {
85         let mut set = RegSet::new(UNIVERSE, 0);
86         let gpr = set.any_gpr().unwrap();
87         set.free_gpr(gpr);
88         assert!(set.gpr(gpr).is_some());
89     }
90 }
91