lumen_compiler/compiler/
regalloc.rs

1//! Simple linear-scan register allocator.
2//! Assigns registers to cell params, let bindings, and temporaries.
3
4use std::collections::HashMap;
5
6/// Register allocation state for a single cell
7#[derive(Debug)]
8pub struct RegAlloc {
9    next_reg: u8,
10    bindings: HashMap<String, u8>,
11}
12
13impl Default for RegAlloc {
14    fn default() -> Self {
15        Self::new()
16    }
17}
18
19impl RegAlloc {
20    pub fn new() -> Self {
21        Self {
22            next_reg: 0,
23            bindings: HashMap::new(),
24        }
25    }
26
27    /// Allocate a named register for a parameter or let binding
28    pub fn alloc_named(&mut self, name: &str) -> u8 {
29        let reg = self.next_reg;
30        self.bindings.insert(name.to_string(), reg);
31        self.next_reg += 1;
32        reg
33    }
34
35    /// Allocate a temporary register
36    pub fn alloc_temp(&mut self) -> u8 {
37        let reg = self.next_reg;
38        self.next_reg += 1;
39        reg
40    }
41
42    /// Look up a named binding
43    pub fn lookup(&self, name: &str) -> Option<u8> {
44        self.bindings.get(name).copied()
45    }
46
47    /// Get the maximum register count used
48    pub fn max_regs(&self) -> u8 {
49        self.next_reg
50    }
51
52    /// Manually bind a name to an existing register (for temporary shadowing)
53    pub fn bind(&mut self, name: &str, reg: u8) {
54        self.bindings.insert(name.to_string(), reg);
55    }
56
57    /// Unbind a name (for temporary shadowing)
58    pub fn unbind(&mut self, name: &str) {
59        self.bindings.remove(name);
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[test]
68    fn test_regalloc_basic() {
69        let mut ra = RegAlloc::new();
70        let r0 = ra.alloc_named("x");
71        let r1 = ra.alloc_named("y");
72        let r2 = ra.alloc_temp();
73        assert_eq!(r0, 0);
74        assert_eq!(r1, 1);
75        assert_eq!(r2, 2);
76        assert_eq!(ra.lookup("x"), Some(0));
77        assert_eq!(ra.lookup("y"), Some(1));
78        assert_eq!(ra.max_regs(), 3);
79    }
80}
lumen_compiler/compiler/regalloc.rs

lumen_compiler/compiler/
regalloc.rs
