lamina_codegen/riscv/
regalloc.rs

1use crate::regalloc::RegisterAllocator as MirRegisterAllocator;
2use lamina_mir::register::{Register, RegisterClass, VirtualReg};
3use lamina_platform::TargetOperatingSystem;
4
5/// RISC-V register allocator with platform-aware register selection.
6///
7/// RISC-V has 32 general-purpose registers (x0-x31):
8/// - x0: zero (hardwired zero)
9/// - x1: ra (return address)
10/// - x2: sp (stack pointer)
11/// - x3: gp (global pointer)
12/// - x4: tp (thread pointer)
13/// - x5-x7, x28-x31: temporaries
14/// - x8: fp/s0 (frame pointer/saved register)
15/// - x9-x15: s1-s7 (saved registers)
16/// - x16-x27: a0-a7 (argument registers), t0-t6 (temporaries)
17///
18/// The allocator uses a conservative subset until prologue/epilogue support
19/// covers every saved register.
20pub struct RiscVRegAlloc {
21    #[allow(dead_code)]
22    target_os: TargetOperatingSystem,
23    available_gprs: Vec<&'static str>,
24    allocated_gprs: std::collections::HashMap<&'static str, VirtualReg>,
25    stack_slots: std::collections::HashMap<VirtualReg, i32>,
26    next_stack_slot: i32,
27}
28
29impl Default for RiscVRegAlloc {
30    fn default() -> Self {
31        Self::new(TargetOperatingSystem::Linux)
32    }
33}
34
35impl RiscVRegAlloc {
36    const AVAILABLE_REGISTERS: &'static [&'static str] = &[
37        "x5", "x6", "x7", // t0-t2
38        "x9", "x10", "x11", "x12", "x13", "x14", "x15", // s1-s7
39        "x16", "x17", "x18", "x19", "x20", "x21", "x22", "x23", // a0-a7
40        "x24", "x25", "x26", "x27", // t3-t6
41        "x28", "x29", "x30", "x31", // t3-t6 (duplicate for simplicity)
42    ];
43
44    pub fn new(target_os: TargetOperatingSystem) -> Self {
45        Self {
46            target_os,
47            available_gprs: Self::AVAILABLE_REGISTERS.to_vec(),
48            allocated_gprs: std::collections::HashMap::new(),
49            stack_slots: std::collections::HashMap::new(),
50            next_stack_slot: -8,
51        }
52    }
53
54    /// Set conservative mode (limit to fewer registers)
55    pub fn set_conservative_mode(&mut self) {
56        self.available_gprs = vec![
57            "x9", "x10", "x11", "x12", "x13", "x14", "x15", // s1-s7
58            "x16", "x17", "x18", "x19", // a0-a3
59        ];
60    }
61
62    /// Get stack slot for a virtual register
63    pub fn get_stack_slot(&self, vreg: &VirtualReg) -> Option<i32> {
64        self.stack_slots.get(vreg).copied()
65    }
66}
67
68impl MirRegisterAllocator for RiscVRegAlloc {
69    type PhysReg = &'static str;
70
71    fn alloc_scratch(&mut self) -> Option<Self::PhysReg> {
72        for &reg in &self.available_gprs {
73            if !self.allocated_gprs.contains_key(reg) {
74                return Some(reg);
75            }
76        }
77        self.available_gprs.first().copied()
78    }
79
80    fn free_scratch(&mut self, phys: Self::PhysReg) {
81        self.allocated_gprs.remove(phys);
82    }
83
84    fn get_mapping(&self, vreg: &VirtualReg) -> Option<Self::PhysReg> {
85        for (reg, allocated_vreg) in &self.allocated_gprs {
86            if allocated_vreg == vreg {
87                return Some(*reg);
88            }
89        }
90        None
91    }
92
93    fn ensure_mapping(&mut self, vreg: VirtualReg) -> Option<Self::PhysReg> {
94        if vreg.class != RegisterClass::Gpr {
95            return None;
96        }
97
98        if let Some(phys) = self.get_mapping(&vreg) {
99            return Some(phys);
100        }
101
102        if let Some(phys) = self.alloc_scratch() {
103            self.allocated_gprs.insert(phys, vreg);
104            return Some(phys);
105        }
106
107        let stack_slot = self.next_stack_slot;
108        self.stack_slots.insert(vreg, stack_slot);
109        self.next_stack_slot -= 8;
110        None // Return None to indicate stack allocation
111    }
112
113    fn mapped_for_register(&self, reg: &Register) -> Option<Self::PhysReg> {
114        match reg {
115            Register::Virtual(v) => self.get_mapping(v),
116            Register::Physical(p) => Some(p.name),
117        }
118    }
119
120    fn occupy(&mut self, _phys: Self::PhysReg) {}
121
122    fn release(&mut self, phys: Self::PhysReg) {
123        self.allocated_gprs.remove(phys);
124    }
125
126    fn is_occupied(&self, phys: Self::PhysReg) -> bool {
127        self.allocated_gprs.contains_key(phys)
128    }
129}
lamina_codegen/riscv/regalloc.rs

lamina_codegen/riscv/
regalloc.rs
