lamina_codegen/x86_64/
regalloc.rs

1//! x86_64 register allocator with platform-aware register selection.
2
3use std::collections::{HashMap, HashSet, VecDeque};
4
5use crate::regalloc::{Allocation, LocalRegisterAllocator as MirRegisterAllocator};
6use lamina_mir::{Register, RegisterClass, VirtualReg};
7use lamina_platform::TargetOperatingSystem;
8
9/// x86_64 register allocator supporting System V AMD64 and Microsoft x64 ABIs.
10///
11/// Uses platform-appropriate GPR pools for stable virtual-to-physical register mappings
12/// and separate scratch pools for short-lived temporaries.
13pub struct X64RegAlloc {
14    target_os: TargetOperatingSystem,
15    free_gprs: VecDeque<&'static str>,
16    used_gprs: HashSet<&'static str>,
17    vreg_to_preg: HashMap<VirtualReg, &'static str>,
18    scratch_free: VecDeque<&'static str>,
19    scratch_used: HashSet<&'static str>,
20    global_plan: bool,
21}
22
23const SYSV_MAP_GPRS: &[&str] = &["r12", "r13", "r14", "r15", "rbx"];
24const SYSV_LEAF_MAP_GPRS: &[&str] = &["r12", "r13", "r14", "r15", "rbx", "r10", "r11"];
25const SYSV_SCRATCH_GPRS: &[&str] = &["r10", "r11"]; // keep as-is for non-leaf
26const SYSV_LEAF_SCRATCH_GPRS: &[&str] = &["rcx", "rdx"]; // for leaf functions
27
28const WIN64_MAP_GPRS: &[&str] = &["rbx", "rsi", "rdi", "r12", "r13", "r14", "r15"];
29const WIN64_LEAF_MAP_GPRS: &[&str] = &[
30    "rbx", "rsi", "rdi", "r12", "r13", "r14", "r15", "r10", "r11",
31];
32const WIN64_SCRATCH_GPRS: &[&str] = &["r10", "r11"];
33const WIN64_LEAF_SCRATCH_GPRS: &[&str] = &["rcx", "rdx"];
34
35impl Default for X64RegAlloc {
36    fn default() -> Self {
37        Self::new_default()
38    }
39}
40
41impl X64RegAlloc {
42    /// Creates a new register allocator for the specified target OS.
43    pub fn new(target_os: TargetOperatingSystem) -> Self {
44        Self::with_target_os(target_os)
45    }
46
47    /// Creates a new register allocator with default target OS (Linux).
48    pub fn new_default() -> Self {
49        Self::with_target_os(TargetOperatingSystem::Linux)
50    }
51
52    fn with_target_os(target_os: TargetOperatingSystem) -> Self {
53        let (map_gprs, scratch_gprs) = match target_os {
54            TargetOperatingSystem::Windows => (WIN64_MAP_GPRS, WIN64_SCRATCH_GPRS),
55            _ => (SYSV_MAP_GPRS, SYSV_SCRATCH_GPRS),
56        };
57
58        let mut free_gprs = VecDeque::new();
59        for &r in map_gprs {
60            free_gprs.push_back(r);
61        }
62
63        let mut scratch_free = VecDeque::new();
64        for &r in scratch_gprs {
65            scratch_free.push_back(r);
66        }
67
68        Self {
69            target_os,
70            free_gprs,
71            used_gprs: HashSet::new(),
72            vreg_to_preg: HashMap::new(),
73            scratch_free,
74            scratch_used: HashSet::new(),
75            global_plan: false,
76        }
77    }
78
79    /// Creates a new register allocator for leaf functions (no calls), expanding the available
80    /// register pool by including r10/r11 as mappable GPRs since they won't be clobbered.
81    pub fn new_leaf(target_os: TargetOperatingSystem) -> Self {
82        let (map_gprs, scratch_gprs) = match target_os {
83            TargetOperatingSystem::Windows => (WIN64_LEAF_MAP_GPRS, WIN64_LEAF_SCRATCH_GPRS),
84            _ => (SYSV_LEAF_MAP_GPRS, SYSV_LEAF_SCRATCH_GPRS),
85        };
86        let mut free_gprs = VecDeque::new();
87        for &r in map_gprs {
88            free_gprs.push_back(r);
89        }
90        let mut scratch_free = VecDeque::new();
91        for &r in scratch_gprs {
92            scratch_free.push_back(r);
93        }
94        Self {
95            target_os,
96            free_gprs,
97            used_gprs: HashSet::new(),
98            vreg_to_preg: HashMap::new(),
99            scratch_free,
100            scratch_used: HashSet::new(),
101            global_plan: false,
102        }
103    }
104
105    /// GPR candidates for global allocation; emitter constraints apply separately.
106    pub fn gpr_pool_for_global_allocation(
107        _target_os: TargetOperatingSystem,
108        _leaf: bool,
109    ) -> Vec<&'static str> {
110        vec![
111            "r12", "r13", "r14", "r15", "rbx", "rsi", "rdi", "r8", "r9", "r10", "r11", "rdx", "rcx",
112        ]
113    }
114
115    /// Rebuild allocator state from a precomputed GPR plan (global linear scan or graph coloring).
116    pub fn from_global_plan(
117        target_os: TargetOperatingSystem,
118        leaf: bool,
119        plan: &HashMap<VirtualReg, Allocation<&'static str>>,
120    ) -> Self {
121        let mut s = if leaf {
122            Self::new_leaf(target_os)
123        } else {
124            Self::new(target_os)
125        };
126        s.free_gprs.clear();
127        s.global_plan = true;
128        for (&vreg, alloc) in plan {
129            if vreg.class != RegisterClass::Gpr {
130                continue;
131            }
132            if let Allocation::Register(phys) = alloc
133                && s.vreg_to_preg.insert(vreg, *phys).is_none()
134            {
135                s.used_gprs.insert(*phys);
136                if let Some(pos) = s.free_gprs.iter().position(|&p| p == *phys) {
137                    let _ = s.free_gprs.remove(pos);
138                }
139            }
140        }
141        s
142    }
143
144    /// Sets conservative mode for complex functions, using fewer registers to reduce pressure.
145    pub fn set_conservative_mode(&mut self) {
146        self.free_gprs.clear();
147
148        match self.target_os {
149            TargetOperatingSystem::Windows => {
150                self.free_gprs.push_back("rbx");
151                self.free_gprs.push_back("rsi");
152                self.free_gprs.push_back("r12");
153                self.free_gprs.push_back("r13");
154                self.free_gprs.push_back("r14");
155            }
156            _ => {
157                self.free_gprs.push_back("r12");
158                self.free_gprs.push_back("r13");
159                self.free_gprs.push_back("r14");
160                self.free_gprs.push_back("rbx");
161            }
162        }
163
164        self.used_gprs.retain(|r| self.free_gprs.contains(r));
165
166        let free_set: HashSet<&str> = self.free_gprs.iter().copied().collect();
167        self.vreg_to_preg.retain(|_, preg| free_set.contains(preg));
168    }
169
170    #[inline]
171    pub fn alloc_scratch(&mut self) -> Option<&'static str> {
172        MirRegisterAllocator::alloc_scratch(self)
173    }
174
175    #[inline]
176    pub fn free_scratch(&mut self, phys: &'static str) {
177        MirRegisterAllocator::free_scratch(self, phys);
178    }
179
180    #[inline]
181    pub fn is_occupied(&self, phys: &'static str) -> bool {
182        MirRegisterAllocator::is_occupied(self, phys)
183    }
184
185    #[inline]
186    pub fn occupy(&mut self, phys: &'static str) {
187        MirRegisterAllocator::occupy(self, phys);
188    }
189
190    #[inline]
191    pub fn release(&mut self, phys: &'static str) {
192        MirRegisterAllocator::release(self, phys);
193    }
194
195    #[inline]
196    pub fn get_mapping_for(&self, v: &VirtualReg) -> Option<&'static str> {
197        MirRegisterAllocator::get_mapping(self, v)
198    }
199
200    #[inline]
201    pub fn ensure_mapping(&mut self, v: VirtualReg) -> Option<&'static str> {
202        MirRegisterAllocator::ensure_mapping(self, v)
203    }
204
205    #[inline]
206    pub fn ensure_mapping_for_gpr(&mut self, v: VirtualReg) -> Option<&'static str> {
207        MirRegisterAllocator::ensure_mapping(self, v)
208    }
209
210    #[inline]
211    pub fn mapped_for_register(&self, r: &Register) -> Option<&'static str> {
212        MirRegisterAllocator::mapped_for_register(self, r)
213    }
214}
215
216impl MirRegisterAllocator for X64RegAlloc {
217    type PhysReg = &'static str;
218
219    fn alloc_scratch(&mut self) -> Option<Self::PhysReg> {
220        if let Some(phys) = self.scratch_free.pop_front() {
221            self.scratch_used.insert(phys);
222            Some(phys)
223        } else {
224            None
225        }
226    }
227
228    fn free_scratch(&mut self, phys: Self::PhysReg) {
229        if self.scratch_used.remove(&phys) {
230            self.scratch_free.push_back(phys);
231        }
232    }
233
234    fn get_mapping(&self, vreg: &VirtualReg) -> Option<Self::PhysReg> {
235        self.vreg_to_preg.get(vreg).copied()
236    }
237
238    fn ensure_mapping(&mut self, vreg: VirtualReg) -> Option<Self::PhysReg> {
239        if let Some(&phys) = self.vreg_to_preg.get(&vreg) {
240            return Some(phys);
241        }
242
243        if self.global_plan || vreg.class != RegisterClass::Gpr {
244            return None;
245        }
246
247        if let Some(phys) = self.free_gprs.pop_front() {
248            self.used_gprs.insert(phys);
249            self.vreg_to_preg.insert(vreg, phys);
250            Some(phys)
251        } else if let Some((vreg_to_replace, &phys)) = self.vreg_to_preg.iter().next() {
252            let vreg_to_replace = *vreg_to_replace;
253            self.vreg_to_preg.remove(&vreg_to_replace);
254            self.vreg_to_preg.insert(vreg, phys);
255            Some(phys)
256        } else {
257            None
258        }
259    }
260
261    fn mapped_for_register(&self, reg: &Register) -> Option<Self::PhysReg> {
262        match reg {
263            Register::Virtual(v) => self.vreg_to_preg.get(v).copied(),
264            Register::Physical(p) => Some(p.name),
265        }
266    }
267
268    fn occupy(&mut self, phys: Self::PhysReg) {
269        if self.used_gprs.insert(phys)
270            && let Some(pos) = self.free_gprs.iter().position(|&p| p == phys)
271        {
272            self.free_gprs.remove(pos);
273        }
274    }
275
276    fn release(&mut self, phys: Self::PhysReg) {
277        if self.used_gprs.remove(&phys) {
278            self.free_gprs.push_back(phys);
279        }
280    }
281
282    fn is_occupied(&self, phys: Self::PhysReg) -> bool {
283        self.used_gprs.contains(&phys)
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290
291    #[test]
292    fn new_linux_has_sysv_pool() {
293        let ra = X64RegAlloc::new(TargetOperatingSystem::Linux);
294        for r in SYSV_MAP_GPRS {
295            assert!(ra.free_gprs.contains(r));
296        }
297    }
298
299    #[test]
300    fn new_windows_has_win64_pool() {
301        let ra = X64RegAlloc::new(TargetOperatingSystem::Windows);
302        assert!(WIN64_MAP_GPRS.iter().all(|reg| ra.free_gprs.contains(reg)));
303    }
304
305    #[test]
306    fn ensure_mapping_returns_same_reg_for_same_vreg() {
307        let mut ra = X64RegAlloc::new_default();
308        let vreg = VirtualReg::gpr(0);
309        let first = ra.ensure_mapping(vreg);
310        let second = ra.ensure_mapping(vreg);
311        assert!(first.is_some());
312        assert_eq!(first, second);
313    }
314
315    #[test]
316    fn alloc_and_free_scratch_cycle() {
317        let mut ra = X64RegAlloc::new_default();
318        let scratch = ra.alloc_scratch();
319        assert!(
320            scratch.is_some(),
321            "scratch should be available on fresh allocator"
322        );
323        let phys = scratch.unwrap();
324        assert!(ra.scratch_used.contains(phys));
325        ra.free_scratch(phys);
326        assert!(!ra.scratch_used.contains(phys));
327        assert!(ra.scratch_free.contains(&phys));
328    }
329
330    #[test]
331    fn conservative_mode_shrinks_pool() {
332        let mut ra = X64RegAlloc::new_default();
333        let full_size = ra.free_gprs.len();
334        ra.set_conservative_mode();
335        assert!(ra.free_gprs.len() < full_size);
336    }
337
338    #[test]
339    fn global_pool_includes_caller_saved_candidates() {
340        let non_leaf =
341            X64RegAlloc::gpr_pool_for_global_allocation(TargetOperatingSystem::Linux, false);
342        let leaf = X64RegAlloc::gpr_pool_for_global_allocation(TargetOperatingSystem::Linux, true);
343        assert_eq!(leaf, non_leaf);
344        assert!(non_leaf.contains(&"r10"));
345        assert!(non_leaf.contains(&"r11"));
346        for reg in ["rax", "rbp", "rsp"] {
347            assert!(!leaf.contains(&reg));
348        }
349        for reg in ["rbx", "rcx", "rdx", "rsi", "rdi", "r8", "r9"] {
350            assert!(leaf.contains(&reg));
351        }
352    }
353
354    #[test]
355    fn global_plan_keeps_shared_mappings_and_spills_fixed() {
356        let first = VirtualReg::gpr(0);
357        let second = VirtualReg::gpr(1);
358        let spilled = VirtualReg::gpr(2);
359        let plan = HashMap::from([
360            (first, Allocation::Register("r12")),
361            (second, Allocation::Register("r12")),
362            (spilled, Allocation::Spill(-8)),
363        ]);
364        let mut ra = X64RegAlloc::from_global_plan(TargetOperatingSystem::Linux, true, &plan);
365        assert_eq!(ra.ensure_mapping(spilled), None);
366        assert_eq!(ra.ensure_mapping(VirtualReg::gpr(3)), None);
367        assert_eq!(ra.get_mapping_for(&first), Some("r12"));
368        assert_eq!(ra.get_mapping_for(&second), Some("r12"));
369        assert_eq!(ra.alloc_scratch(), Some("rcx"));
370        assert_eq!(ra.alloc_scratch(), Some("rdx"));
371    }
372
373    #[test]
374    fn is_occupied_tracks_occupy_and_release() {
375        let mut ra = X64RegAlloc::new_default();
376        let phys = "rbx";
377        assert!(!ra.is_occupied(phys));
378        ra.occupy(phys);
379        assert!(ra.is_occupied(phys));
380        ra.release(phys);
381        assert!(!ra.is_occupied(phys));
382    }
383
384    #[test]
385    fn non_gpr_vreg_returns_none() {
386        let mut ra = X64RegAlloc::new_default();
387        let fp_vreg = VirtualReg {
388            id: 99,
389            class: RegisterClass::Fpr,
390        };
391        assert!(ra.ensure_mapping(fp_vreg).is_none());
392    }
393}
lamina_codegen/x86_64/regalloc.rs

lamina_codegen/x86_64/
regalloc.rs
