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::RegisterAllocator 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}
21
22const SYSV_MAP_GPRS: &[&str] = &["r12", "r13", "r14", "r15", "rbx"];
23const SYSV_SCRATCH_GPRS: &[&str] = &["r10", "r11"];
24
25const WIN64_MAP_GPRS: &[&str] = &["rbx", "rsi", "rdi", "r12", "r13", "r14", "r15"];
26const WIN64_SCRATCH_GPRS: &[&str] = &["r10", "r11"];
27
28impl Default for X64RegAlloc {
29    fn default() -> Self {
30        Self::new_default()
31    }
32}
33
34impl X64RegAlloc {
35    /// Creates a new register allocator for the specified target OS.
36    pub fn new(target_os: TargetOperatingSystem) -> Self {
37        Self::with_target_os(target_os)
38    }
39
40    /// Creates a new register allocator with default target OS (Linux).
41    pub fn new_default() -> Self {
42        Self::with_target_os(TargetOperatingSystem::Linux)
43    }
44
45    fn with_target_os(target_os: TargetOperatingSystem) -> Self {
46        let (map_gprs, scratch_gprs) = match target_os {
47            TargetOperatingSystem::Windows => (WIN64_MAP_GPRS, WIN64_SCRATCH_GPRS),
48            _ => (SYSV_MAP_GPRS, SYSV_SCRATCH_GPRS),
49        };
50
51        let mut free_gprs = VecDeque::new();
52        for &r in map_gprs {
53            free_gprs.push_back(r);
54        }
55
56        let mut scratch_free = VecDeque::new();
57        for &r in scratch_gprs {
58            scratch_free.push_back(r);
59        }
60
61        Self {
62            target_os,
63            free_gprs,
64            used_gprs: HashSet::new(),
65            vreg_to_preg: HashMap::new(),
66            scratch_free,
67            scratch_used: HashSet::new(),
68        }
69    }
70
71    /// Sets conservative mode for complex functions, using fewer registers to reduce pressure.
72    pub fn set_conservative_mode(&mut self) {
73        self.free_gprs.clear();
74
75        match self.target_os {
76            TargetOperatingSystem::Windows => {
77                self.free_gprs.push_back("rbx");
78                self.free_gprs.push_back("rsi");
79                self.free_gprs.push_back("r12");
80                self.free_gprs.push_back("r13");
81                self.free_gprs.push_back("r14");
82            }
83            _ => {
84                self.free_gprs.push_back("r12");
85                self.free_gprs.push_back("r13");
86                self.free_gprs.push_back("r14");
87                self.free_gprs.push_back("rbx");
88            }
89        }
90
91        self.used_gprs.retain(|r| self.free_gprs.contains(r));
92
93        let free_set: HashSet<&str> = self.free_gprs.iter().copied().collect();
94        self.vreg_to_preg.retain(|_, preg| free_set.contains(preg));
95    }
96
97    #[inline]
98    pub fn alloc_scratch(&mut self) -> Option<&'static str> {
99        MirRegisterAllocator::alloc_scratch(self)
100    }
101
102    #[inline]
103    pub fn free_scratch(&mut self, phys: &'static str) {
104        MirRegisterAllocator::free_scratch(self, phys);
105    }
106
107    #[inline]
108    pub fn is_occupied(&self, phys: &'static str) -> bool {
109        MirRegisterAllocator::is_occupied(self, phys)
110    }
111
112    #[inline]
113    pub fn occupy(&mut self, phys: &'static str) {
114        MirRegisterAllocator::occupy(self, phys);
115    }
116
117    #[inline]
118    pub fn release(&mut self, phys: &'static str) {
119        MirRegisterAllocator::release(self, phys);
120    }
121
122    #[inline]
123    pub fn get_mapping_for(&self, v: &VirtualReg) -> Option<&'static str> {
124        MirRegisterAllocator::get_mapping(self, v)
125    }
126
127    #[inline]
128    pub fn ensure_mapping(&mut self, v: VirtualReg) -> Option<&'static str> {
129        MirRegisterAllocator::ensure_mapping(self, v)
130    }
131
132    #[inline]
133    pub fn ensure_mapping_for_gpr(&mut self, v: VirtualReg) -> Option<&'static str> {
134        MirRegisterAllocator::ensure_mapping(self, v)
135    }
136
137    #[inline]
138    pub fn mapped_for_register(&self, r: &Register) -> Option<&'static str> {
139        MirRegisterAllocator::mapped_for_register(self, r)
140    }
141}
142
143impl MirRegisterAllocator for X64RegAlloc {
144    type PhysReg = &'static str;
145
146    fn alloc_scratch(&mut self) -> Option<Self::PhysReg> {
147        if let Some(phys) = self.scratch_free.pop_front() {
148            self.scratch_used.insert(phys);
149            Some(phys)
150        } else {
151            None
152        }
153    }
154
155    fn free_scratch(&mut self, phys: Self::PhysReg) {
156        if self.scratch_used.remove(&phys) {
157            self.scratch_free.push_back(phys);
158        }
159    }
160
161    fn get_mapping(&self, vreg: &VirtualReg) -> Option<Self::PhysReg> {
162        self.vreg_to_preg.get(vreg).copied()
163    }
164
165    fn ensure_mapping(&mut self, vreg: VirtualReg) -> Option<Self::PhysReg> {
166        if let Some(&phys) = self.vreg_to_preg.get(&vreg) {
167            return Some(phys);
168        }
169
170        if vreg.class != RegisterClass::Gpr {
171            return None;
172        }
173
174        if let Some(phys) = self.free_gprs.pop_front() {
175            self.used_gprs.insert(phys);
176            self.vreg_to_preg.insert(vreg, phys);
177            Some(phys)
178        } else if let Some((vreg_to_replace, &phys)) = self.vreg_to_preg.iter().next() {
179            let vreg_to_replace = *vreg_to_replace;
180            self.vreg_to_preg.remove(&vreg_to_replace);
181            self.vreg_to_preg.insert(vreg, phys);
182            Some(phys)
183        } else {
184            None
185        }
186    }
187
188    fn mapped_for_register(&self, reg: &Register) -> Option<Self::PhysReg> {
189        match reg {
190            Register::Virtual(v) => self.vreg_to_preg.get(v).copied(),
191            Register::Physical(p) => Some(p.name),
192        }
193    }
194
195    fn occupy(&mut self, phys: Self::PhysReg) {
196        if self.used_gprs.insert(phys)
197            && let Some(pos) = self.free_gprs.iter().position(|&p| p == phys)
198        {
199            self.free_gprs.remove(pos);
200        }
201    }
202
203    fn release(&mut self, phys: Self::PhysReg) {
204        if self.used_gprs.remove(&phys) {
205            self.free_gprs.push_back(phys);
206        }
207    }
208
209    fn is_occupied(&self, phys: Self::PhysReg) -> bool {
210        self.used_gprs.contains(&phys)
211    }
212}
lamina_codegen/x86_64/regalloc.rs

lamina_codegen/x86_64/
regalloc.rs
