1 use crate::{isa::reg::Reg, regset::RegSet};
2 use regalloc2::RegClass;
3 
4 /// The register allocator.
5 ///
6 /// The register allocator uses a single-pass algorithm;
7 /// its implementation uses a bitset as a freelist
8 /// to track per-class register availability.
9 ///
10 /// If a particular register is not available upon request
11 /// the register allocation will perform a "spill", essentially
12 /// moving Local and Register values in the stack to memory.
13 /// This processs ensures that whenever a register is requested,
14 /// it is going to be available.
15 pub(crate) struct RegAlloc {
16     pub scratch: Reg,
17     regset: RegSet,
18 }
19 
20 impl RegAlloc {
21     /// Create a new register allocator
22     /// from a register set.
23     pub fn new(regset: RegSet, scratch: Reg) -> Self {
24         Self { regset, scratch }
25     }
26 
27     /// Allocate the next available register for the given class,
28     /// spilling if not available.
29     pub fn reg_for_class<F>(&mut self, class: RegClass, spill: &mut F) -> Reg
30     where
31         F: FnMut(&mut RegAlloc),
32     {
33         self.regset.reg_for_class(class).unwrap_or_else(|| {
34             spill(self);
35             self.regset.reg_for_class(class).unwrap_or_else(|| {
36                 panic!("expected register for class {:?}, to be avilable", class)
37             })
38         })
39     }
40 
41     /// Returns true if the specified register is allocatable.
42     pub fn reg_available(&self, reg: Reg) -> bool {
43         self.regset.named_reg_available(reg)
44     }
45 
46     /// Request a specific register, spilling if not available.
47     pub fn reg<F>(&mut self, named: Reg, mut spill: F) -> Reg
48     where
49         F: FnMut(&mut RegAlloc),
50     {
51         // If the scratch register is explicitly requested
52         // just return it, it's usage should never cause spills.
53         if named == self.scratch {
54             return named;
55         }
56 
57         self.regset.reg(named).unwrap_or_else(|| {
58             spill(self);
59             self.regset
60                 .reg(named)
61                 .expect(&format!("register {:?} to be available", named))
62         })
63     }
64 
65     /// Free the given register.
66     pub fn free(&mut self, reg: Reg) {
67         if reg != self.scratch {
68             self.regset.free(reg);
69         }
70     }
71 }
72