1 use crate::{
2     codegen::CodeGenContext,
3     frame::Frame,
4     isa::reg::Reg,
5     masm::{MacroAssembler, OperandSize, RegImm},
6     regset::RegSet,
7     stack::Val,
8 };
9 
10 /// The register allocator.
11 ///
12 /// The register allocator uses a single-pass algorithm;
13 /// its implementation uses a bitset as a freelist
14 /// to track per-class register availability.
15 ///
16 /// If a particular register is not available upon request
17 /// the register allocation will perform a "spill", essentially
18 /// moving Local and Register values in the stack to memory.
19 /// This processs ensures that whenever a register is requested,
20 /// it is going to be available.
21 pub(crate) struct RegAlloc {
22     pub scratch: Reg,
23     regset: RegSet,
24 }
25 
26 impl RegAlloc {
27     /// Create a new register allocator
28     /// from a register set.
29     pub fn new(regset: RegSet, scratch: Reg) -> Self {
30         Self { regset, scratch }
31     }
32 
33     /// Loads the stack top value into a register, if it isn't already one;
34     /// spilling if there are no registers available.
35     pub fn pop_to_reg<M: MacroAssembler>(
36         &mut self,
37         context: &mut CodeGenContext<M>,
38         size: OperandSize,
39     ) -> Reg {
40         if let Some(reg) = context.stack.pop_reg() {
41             return reg;
42         }
43 
44         let dst = self.any_gpr(context);
45         let val = context.stack.pop().expect("a value at stack top");
46         Self::move_val_to_reg(val, dst, &mut context.masm, context.frame, size);
47         dst
48     }
49 
50     /// Checks if the stack top contains the given register. The register
51     /// gets allocated otherwise, potentially causing a spill.
52     /// Once the requested register is allocated, the value at the top of the stack
53     /// gets loaded into the register.
54     pub fn pop_to_named_reg<M: MacroAssembler>(
55         &mut self,
56         context: &mut CodeGenContext<M>,
57         named: Reg,
58         size: OperandSize,
59     ) -> Reg {
60         if let Some(reg) = context.stack.pop_named_reg(named) {
61             return reg;
62         }
63 
64         let dst = self.gpr(context, named);
65         let val = context.stack.pop().expect("a value at stack top");
66         Self::move_val_to_reg(val, dst, &mut context.masm, context.frame, size);
67         dst
68     }
69 
70     fn move_val_to_reg<M: MacroAssembler>(
71         src: Val,
72         dst: Reg,
73         masm: &mut M,
74         frame: &Frame,
75         size: OperandSize,
76     ) {
77         match src {
78             Val::Reg(src) => masm.mov(RegImm::reg(src), RegImm::reg(dst), size),
79             Val::I32(imm) => masm.mov(RegImm::imm(imm), RegImm::reg(dst), size),
80             Val::Local(index) => {
81                 let slot = frame
82                     .get_local(index)
83                     .expect(&format!("valid locat at index = {}", index));
84                 let addr = masm.local_address(&slot);
85                 masm.load(addr, dst, slot.ty.into());
86             }
87             v => panic!("Unsupported value {:?}", v),
88         };
89     }
90 
91     /// Allocate the next available general purpose register,
92     /// spilling if none available.
93     pub fn any_gpr<M: MacroAssembler>(&mut self, context: &mut CodeGenContext<M>) -> Reg {
94         self.regset.any_gpr().unwrap_or_else(|| {
95             self.spill(context);
96             self.regset.any_gpr().expect("any gpr to be available")
97         })
98     }
99 
100     /// Request a specific general purpose register,
101     /// spilling if not available.
102     pub fn gpr<M: MacroAssembler>(&mut self, context: &mut CodeGenContext<M>, named: Reg) -> Reg {
103         self.regset.gpr(named).unwrap_or_else(|| {
104             self.spill(context);
105             self.regset
106                 .gpr(named)
107                 .expect(&format!("gpr {:?} to be available", named))
108         })
109     }
110 
111     /// Mark a particular general purpose register as available.
112     pub fn free_gpr(&mut self, reg: Reg) {
113         self.regset.free_gpr(reg);
114     }
115 
116     /// Spill locals and registers to memory.
117     // TODO optimize the spill range;
118     //
119     // At any point in the program, the stack
120     // might already contain Memory entries;
121     // we could effectively ignore that range;
122     // only focusing on the range that contains
123     // spillable values.
124     fn spill<M: MacroAssembler>(&mut self, context: &mut CodeGenContext<M>) {
125         context.stack.inner_mut().iter_mut().for_each(|v| match v {
126             Val::Reg(r) => {
127                 let offset = context.masm.push(*r);
128                 self.free_gpr(*r);
129                 *v = Val::Memory(offset);
130             }
131             Val::Local(index) => {
132                 let slot = context
133                     .frame
134                     .get_local(*index)
135                     .expect("valid local at slot");
136                 let addr = context.masm.local_address(&slot);
137                 context
138                     .masm
139                     .store(RegImm::reg(self.scratch), addr, slot.ty.into());
140                 let offset = context.masm.push(self.scratch);
141                 *v = Val::Memory(offset);
142             }
143             _ => {}
144         });
145     }
146 }
147