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, 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, 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.into()), RegImm::reg(dst), size),
80             Val::I64(imm) => masm.mov(RegImm::imm(imm), RegImm::reg(dst), size),
81             Val::Local(index) => {
82                 let slot = frame
83                     .get_local(index)
84                     .expect(&format!("valid locat at index = {}", index));
85                 let addr = masm.local_address(&slot);
86                 masm.load(addr, dst, slot.ty.into());
87             }
88             v => panic!("Unsupported value {:?}", v),
89         };
90     }
91 
92     /// Allocate the next available general purpose register,
93     /// spilling if none available.
94     pub fn any_gpr<M: MacroAssembler>(&mut self, context: &mut CodeGenContext<M>) -> Reg {
95         self.regset.any_gpr().unwrap_or_else(|| {
96             self.spill(context);
97             self.regset.any_gpr().expect("any gpr to be available")
98         })
99     }
100 
101     /// Request a specific general purpose register,
102     /// spilling if not available.
103     pub fn gpr<M: MacroAssembler>(&mut self, context: &mut CodeGenContext<M>, named: Reg) -> Reg {
104         self.regset.gpr(named).unwrap_or_else(|| {
105             self.spill(context);
106             self.regset
107                 .gpr(named)
108                 .expect(&format!("gpr {:?} to be available", named))
109         })
110     }
111 
112     /// Mark a particular general purpose register as available.
113     pub fn free_gpr(&mut self, reg: Reg) {
114         self.regset.free_gpr(reg);
115     }
116 
117     /// Spill locals and registers to memory.
118     // TODO optimize the spill range;
119     //
120     // At any point in the program, the stack
121     // might already contain Memory entries;
122     // we could effectively ignore that range;
123     // only focusing on the range that contains
124     // spillable values.
125     fn spill<M: MacroAssembler>(&mut self, context: &mut CodeGenContext<M>) {
126         context.stack.inner_mut().iter_mut().for_each(|v| match v {
127             Val::Reg(r) => {
128                 let offset = context.masm.push(*r);
129                 self.free_gpr(*r);
130                 *v = Val::Memory(offset);
131             }
132             Val::Local(index) => {
133                 let slot = context
134                     .frame
135                     .get_local(*index)
136                     .expect("valid local at slot");
137                 let addr = context.masm.local_address(&slot);
138                 context.masm.load(addr, self.scratch, slot.ty.into());
139                 let offset = context.masm.push(self.scratch);
140                 *v = Val::Memory(offset);
141             }
142             _ => {}
143         });
144     }
145 }
146