lamina_codegen/wasm/
regalloc.rs

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

lamina_codegen/wasm/
regalloc.rs
