lamina_codegen/powerpc/
regalloc.rs

1//! Simple PowerPC64 register allocator.
2//!
3//! Maps virtual registers to caller-saved GPRs in order.
4//! When exhausted, the caller should spill to a stack slot instead.
5
6use crate::regalloc::{Allocation, LocalRegisterAllocator as MirRegisterAllocator};
7use lamina_mir::{Register, RegisterClass, VirtualReg};
8use lamina_platform::TargetOperatingSystem;
9use std::collections::HashMap;
10
11/// Caller-saved general-purpose registers available for allocation.
12///
13/// We exclude r1 (sp), r2 (TOC), r13 (thread pointer on Linux).
14const AVAILABLE_GPRS: &[&str] = &["6", "7", "8", "9", "10", "11", "12"];
15
16pub struct Ppc64RegAlloc {
17    #[allow(dead_code)]
18    target_os: TargetOperatingSystem,
19    global_mappings: HashMap<VirtualReg, &'static str>,
20    allocated_gprs: HashMap<&'static str, VirtualReg>,
21    stack_slots: HashMap<VirtualReg, i32>,
22    next_stack_slot: i32,
23}
24
25impl Default for Ppc64RegAlloc {
26    fn default() -> Self {
27        Self::new(TargetOperatingSystem::Linux)
28    }
29}
30
31impl Ppc64RegAlloc {
32    pub fn new(target_os: TargetOperatingSystem) -> Self {
33        Self {
34            target_os,
35            global_mappings: HashMap::new(),
36            allocated_gprs: HashMap::new(),
37            stack_slots: HashMap::new(),
38            next_stack_slot: -8,
39        }
40    }
41
42    pub fn get_mapping_for(&self, vreg: &VirtualReg) -> Option<&'static str> {
43        if let Some(reg) = self.global_mappings.get(vreg) {
44            return Some(*reg);
45        }
46        for (reg, allocated_vreg) in &self.allocated_gprs {
47            if allocated_vreg == vreg {
48                return Some(*reg);
49            }
50        }
51        None
52    }
53
54    pub fn gpr_pool_for_global_allocation() -> Vec<&'static str> {
55        AVAILABLE_GPRS.to_vec()
56    }
57
58    pub fn from_global_plan(
59        target_os: TargetOperatingSystem,
60        plan: &HashMap<VirtualReg, Allocation<&'static str>>,
61    ) -> Self {
62        let mut s = Self::new(target_os);
63        for (&vreg, alloc) in plan {
64            if vreg.class != RegisterClass::Gpr {
65                continue;
66            }
67            if let Allocation::Register(phys) = alloc
68                && AVAILABLE_GPRS.contains(phys)
69            {
70                s.global_mappings.insert(vreg, *phys);
71                s.allocated_gprs.insert(*phys, vreg);
72            }
73        }
74        s
75    }
76}
77
78impl MirRegisterAllocator for Ppc64RegAlloc {
79    type PhysReg = &'static str;
80
81    fn alloc_scratch(&mut self) -> Option<Self::PhysReg> {
82        for &reg in AVAILABLE_GPRS {
83            if !self.allocated_gprs.contains_key(reg) {
84                return Some(reg);
85            }
86        }
87        AVAILABLE_GPRS.first().copied()
88    }
89
90    fn free_scratch(&mut self, phys: Self::PhysReg) {
91        self.allocated_gprs.remove(phys);
92    }
93
94    fn get_mapping(&self, vreg: &VirtualReg) -> Option<Self::PhysReg> {
95        self.get_mapping_for(vreg)
96    }
97
98    fn ensure_mapping(&mut self, vreg: VirtualReg) -> Option<Self::PhysReg> {
99        if vreg.class != RegisterClass::Gpr {
100            return None;
101        }
102        if let Some(phys) = self.get_mapping(&vreg) {
103            return Some(phys);
104        }
105        if let Some(phys) = self.alloc_scratch() {
106            self.allocated_gprs.insert(phys, vreg);
107            return Some(phys);
108        }
109        let slot = self.next_stack_slot;
110        self.stack_slots.insert(vreg, slot);
111        self.next_stack_slot -= 8;
112        None
113    }
114
115    fn mapped_for_register(&self, reg: &Register) -> Option<Self::PhysReg> {
116        match reg {
117            Register::Virtual(v) => self.get_mapping(v),
118            Register::Physical(p) => Some(p.name),
119        }
120    }
121
122    fn occupy(&mut self, _phys: Self::PhysReg) {}
123
124    fn release(&mut self, phys: Self::PhysReg) {
125        self.allocated_gprs.remove(phys);
126    }
127
128    fn is_occupied(&self, phys: Self::PhysReg) -> bool {
129        self.allocated_gprs.contains_key(phys)
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use super::*;
136
137    #[test]
138    fn global_plan_preserves_reused_register_mappings() {
139        let first = VirtualReg::gpr(0);
140        let second = VirtualReg::gpr(1);
141        let plan = HashMap::from([
142            (first, Allocation::Register("6")),
143            (second, Allocation::Register("6")),
144        ]);
145        let allocator = Ppc64RegAlloc::from_global_plan(TargetOperatingSystem::Linux, &plan);
146        assert_eq!(allocator.get_mapping(&first), Some("6"));
147        assert_eq!(allocator.get_mapping(&second), Some("6"));
148    }
149
150    #[test]
151    fn test_allocate_gpr_sequential() {
152        let mut ra = Ppc64RegAlloc::new(TargetOperatingSystem::Linux);
153        let v0 = VirtualReg::gpr(0);
154        let v1 = VirtualReg::gpr(1);
155        let r0 = ra.ensure_mapping(v0).unwrap();
156        let r1 = ra.ensure_mapping(v1).unwrap();
157        assert_ne!(r0, r1);
158    }
159
160    #[test]
161    fn test_allocate_same_vreg_stable() {
162        let mut ra = Ppc64RegAlloc::new(TargetOperatingSystem::Linux);
163        let v = VirtualReg::gpr(0);
164        let r1 = ra.ensure_mapping(v);
165        let r2 = ra.ensure_mapping(v);
166        assert_eq!(r1, r2);
167    }
168
169    #[test]
170    fn test_alloc_scratch_not_none() {
171        let mut ra = Ppc64RegAlloc::new(TargetOperatingSystem::Linux);
172        let s = ra.alloc_scratch();
173        assert!(s.is_some());
174    }
175
176    #[test]
177    fn test_release_allows_reuse() {
178        let mut ra = Ppc64RegAlloc::new(TargetOperatingSystem::Linux);
179        let s = ra.alloc_scratch().unwrap();
180        ra.free_scratch(s);
181        let s2 = ra.alloc_scratch();
182        assert!(s2.is_some());
183    }
184}
lamina_codegen/powerpc/regalloc.rs

lamina_codegen/powerpc/
regalloc.rs
