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