lamina_codegen/wasm/
regalloc.rs

1use crate::regalloc::{
2    PhysRegConvertible, PhysRegHandle, RegisterAllocator as MirRegisterAllocator,
3};
4use lamina_mir::register::{Register, RegisterClass, VirtualReg};
5
6/// WASM "register allocator" for stack-based virtual machine.
7///
8/// WASM doesn't use physical registers like x86_64, but we still need to track
9/// virtual registers for stack operations. This is a simplified allocator
10/// that maps virtual registers to stack positions.
11pub struct WasmRegAlloc {
12    vreg_to_stack: std::collections::HashMap<VirtualReg, u32>,
13    next_stack_slot: u32,
14}
15
16impl Default for WasmRegAlloc {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl WasmRegAlloc {
23    pub fn new() -> Self {
24        Self {
25            vreg_to_stack: std::collections::HashMap::new(),
26            next_stack_slot: 0,
27        }
28    }
29
30    /// Get the stack position for a virtual register
31    pub fn get_stack_position(&self, vreg: &VirtualReg) -> Option<u32> {
32        self.vreg_to_stack.get(vreg).copied()
33    }
34
35    /// Allocate a stack position for a virtual register
36    pub fn allocate_stack(&mut self, vreg: VirtualReg) -> u32 {
37        if let Some(pos) = self.vreg_to_stack.get(&vreg) {
38            *pos
39        } else {
40            let pos = self.next_stack_slot;
41            self.vreg_to_stack.insert(vreg, pos);
42            self.next_stack_slot += 1;
43            pos
44        }
45    }
46}
47
48impl MirRegisterAllocator for WasmRegAlloc {
49    type PhysReg = u32; // Stack position as "physical register"
50
51    fn alloc_scratch(&mut self) -> Option<Self::PhysReg> {
52        let pos = self.next_stack_slot;
53        self.next_stack_slot += 1;
54        Some(pos)
55    }
56
57    fn free_scratch(&mut self, _phys: Self::PhysReg) {}
58
59    fn get_mapping(&self, vreg: &VirtualReg) -> Option<Self::PhysReg> {
60        self.vreg_to_stack.get(vreg).copied()
61    }
62
63    fn ensure_mapping(&mut self, vreg: VirtualReg) -> Option<Self::PhysReg> {
64        if vreg.class != RegisterClass::Gpr {
65            return None;
66        }
67        Some(self.allocate_stack(vreg))
68    }
69
70    fn mapped_for_register(&self, reg: &Register) -> Option<Self::PhysReg> {
71        match reg {
72            Register::Virtual(v) => self.vreg_to_stack.get(v).copied(),
73            Register::Physical(_) => None,
74        }
75    }
76
77    fn occupy(&mut self, _phys: Self::PhysReg) {}
78
79    fn release(&mut self, _phys: Self::PhysReg) {}
80
81    fn is_occupied(&self, _phys: Self::PhysReg) -> bool {
82        false
83    }
84}
85
86impl PhysRegConvertible for u32 {
87    fn into_handle(self) -> PhysRegHandle {
88        PhysRegHandle::Named(format!("{}", self).leak())
89    }
90
91    fn from_handle(handle: PhysRegHandle) -> Option<Self> {
92        match handle {
93            PhysRegHandle::Named(name) => name.parse().ok(),
94        }
95    }
96}
lamina_codegen/wasm/regalloc.rs

lamina_codegen/wasm/
regalloc.rs
