lamina_codegen/aarch64/
regalloc.rs

1use std::collections::{HashMap, HashSet, VecDeque};
2
3use crate::regalloc::RegisterAllocator as MirRegisterAllocator;
4use lamina_mir::{Register, RegisterClass, VirtualReg};
5
6// AArch64 register allocator for MIR virtual registers.
7
8/// AArch64 register allocator for MIR virtual registers.
9///
10/// Tracks a fixed pool of available caller-saved GPRs and provides
11/// stable mappings from virtual registers to physical registers.
12/// Returns None when no register is available (caller should spill to stack).
13pub struct A64RegAlloc {
14    free_gprs: VecDeque<&'static str>,
15    used_gprs: HashSet<&'static str>,
16    vreg_to_preg: HashMap<VirtualReg, &'static str>,
17    scratch_free: VecDeque<&'static str>,
18    scratch_used: HashSet<&'static str>,
19}
20
21// Only use caller-saved registers (x9-x15) until we implement callee-saved save/restore
22// x19-x28 are callee-saved and must be saved/restored in prologue/epilogue if used
23const MAP_GPRS: &[&str] = &["x9", "x10", "x11", "x12", "x13", "x14", "x15"];
24
25const SCRATCH_GPRS: &[&str] = &["x9", "x10", "x11", "x12", "x13", "x14", "x15"];
26
27impl Default for A64RegAlloc {
28    fn default() -> Self {
29        Self::new()
30    }
31}
32
33impl A64RegAlloc {
34    /// Creates a new register allocator with default register pools.
35    pub fn new() -> Self {
36        let mut free_gprs = VecDeque::new();
37        for &r in MAP_GPRS {
38            free_gprs.push_back(r);
39        }
40
41        let mut scratch_free = VecDeque::new();
42        for &r in SCRATCH_GPRS {
43            scratch_free.push_back(r);
44        }
45
46        Self {
47            free_gprs,
48            used_gprs: HashSet::new(),
49            vreg_to_preg: HashMap::new(),
50            scratch_free,
51            scratch_used: HashSet::new(),
52        }
53    }
54
55    /// Sets conservative mode for complex functions, using fewer registers to reduce pressure.
56    pub fn set_conservative_mode(&mut self) {
57        self.free_gprs.clear();
58        self.free_gprs.push_back("x13");
59        self.free_gprs.push_back("x14");
60        self.free_gprs.push_back("x15");
61        self.free_gprs.push_back("x19");
62        self.free_gprs.push_back("x20");
63        self.free_gprs.push_back("x21");
64        self.free_gprs.push_back("x22");
65        self.free_gprs.push_back("x23");
66        self.used_gprs.retain(|r| self.free_gprs.contains(r));
67        let free_set: HashSet<&str> = self.free_gprs.iter().copied().collect();
68        self.vreg_to_preg.retain(|_, preg| free_set.contains(preg));
69    }
70
71    #[inline]
72    pub fn alloc_scratch(&mut self) -> Option<&'static str> {
73        MirRegisterAllocator::alloc_scratch(self)
74    }
75
76    #[inline]
77    pub fn free_scratch(&mut self, phys: &'static str) {
78        MirRegisterAllocator::free_scratch(self, phys);
79    }
80
81    #[inline]
82    pub fn is_occupied(&self, phys: &'static str) -> bool {
83        MirRegisterAllocator::is_occupied(self, phys)
84    }
85
86    #[inline]
87    pub fn occupy(&mut self, phys: &'static str) {
88        MirRegisterAllocator::occupy(self, phys);
89    }
90
91    #[inline]
92    pub fn release(&mut self, phys: &'static str) {
93        MirRegisterAllocator::release(self, phys);
94    }
95
96    #[inline]
97    pub fn get_mapping_for(&self, v: &VirtualReg) -> Option<&'static str> {
98        MirRegisterAllocator::get_mapping(self, v)
99    }
100
101    #[inline]
102    pub fn ensure_mapping(&mut self, v: VirtualReg) -> Option<&'static str> {
103        MirRegisterAllocator::ensure_mapping(self, v)
104    }
105
106    #[inline]
107    pub fn ensure_mapping_for_gpr(&mut self, v: VirtualReg) -> Option<&'static str> {
108        MirRegisterAllocator::ensure_mapping(self, v)
109    }
110
111    #[inline]
112    pub fn mapped_for_register(&self, r: &Register) -> Option<&'static str> {
113        MirRegisterAllocator::mapped_for_register(self, r)
114    }
115}
116
117impl MirRegisterAllocator for A64RegAlloc {
118    type PhysReg = &'static str;
119
120    fn alloc_scratch(&mut self) -> Option<Self::PhysReg> {
121        if let Some(phys) = self.scratch_free.pop_front() {
122            self.scratch_used.insert(phys);
123            Some(phys)
124        } else {
125            None
126        }
127    }
128
129    fn free_scratch(&mut self, phys: Self::PhysReg) {
130        if self.scratch_used.remove(&phys) {
131            self.scratch_free.push_back(phys);
132        }
133    }
134
135    fn get_mapping(&self, vreg: &VirtualReg) -> Option<Self::PhysReg> {
136        self.vreg_to_preg.get(vreg).copied()
137    }
138
139    fn ensure_mapping(&mut self, vreg: VirtualReg) -> Option<Self::PhysReg> {
140        if let Some(&phys) = self.vreg_to_preg.get(&vreg) {
141            return Some(phys);
142        }
143
144        if vreg.class != RegisterClass::Gpr {
145            return None;
146        }
147
148        if let Some(phys) = self.free_gprs.pop_front() {
149            self.used_gprs.insert(phys);
150            self.vreg_to_preg.insert(vreg, phys);
151            Some(phys)
152        } else if let Some((vreg_to_replace, &phys)) = self.vreg_to_preg.iter().next() {
153            let vreg_to_replace = *vreg_to_replace;
154            self.vreg_to_preg.remove(&vreg_to_replace);
155            self.vreg_to_preg.insert(vreg, phys);
156            Some(phys)
157        } else {
158            None
159        }
160    }
161
162    fn mapped_for_register(&self, reg: &Register) -> Option<Self::PhysReg> {
163        match reg {
164            Register::Virtual(v) => self.vreg_to_preg.get(v).copied(),
165            Register::Physical(p) => Some(p.name),
166        }
167    }
168
169    fn occupy(&mut self, phys: Self::PhysReg) {
170        if self.used_gprs.insert(phys)
171            && let Some(pos) = self.free_gprs.iter().position(|&p| p == phys)
172        {
173            self.free_gprs.remove(pos);
174        }
175    }
176
177    fn release(&mut self, phys: Self::PhysReg) {
178        if self.used_gprs.remove(&phys) {
179            self.free_gprs.push_back(phys);
180        }
181    }
182
183    fn is_occupied(&self, phys: Self::PhysReg) -> bool {
184        self.used_gprs.contains(&phys)
185    }
186}
lamina_codegen/aarch64/regalloc.rs

lamina_codegen/aarch64/
regalloc.rs
