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