Skip to main content

lamina_codegen/aarch64/
frame.rs

1//! Stack frame layout for AArch64 code generation.
2//!
3//! This module provides frame mapping for virtual registers to stack slots.
4
5use lamina_mir::{Function, Register};
6use std::collections::{HashMap, HashSet};
7
8/// Maps virtual registers to stack slot offsets.
9pub struct FrameMap {
10    pub slots: HashMap<Register, i32>,
11    pub frame_size: i32,
12}
13
14impl FrameMap {
15    /// Creates a frame map from a function, assigning stack slots to all virtual registers.
16    pub fn from_function(f: &Function) -> Self {
17        let mut regs: HashSet<Register> = HashSet::new();
18        for p in &f.sig.params {
19            regs.insert(p.reg.clone());
20        }
21        for b in &f.blocks {
22            for ins in &b.instructions {
23                if let Some(d) = ins.def_reg() {
24                    regs.insert(d.clone());
25                }
26                for u in ins.use_regs() {
27                    regs.insert(u.clone());
28                }
29            }
30        }
31
32        let mut reg_vec: Vec<Register> = regs.into_iter().collect();
33        reg_vec.sort_by(|a, b| format!("{:?}", a).cmp(&format!("{:?}", b)));
34
35        let mut slots = HashMap::new();
36        let mut offset: i32 = -8;
37        for r in reg_vec {
38            if matches!(r, Register::Virtual(_)) {
39                slots.insert(r, offset);
40                offset -= 8;
41            }
42        }
43        let mut frame_size = -offset - 8;
44        if frame_size < 0 {
45            frame_size = 0;
46        }
47        frame_size = (frame_size + 15) & !15;
48        Self { slots, frame_size }
49    }
50
51    /// Returns the stack slot offset for a register, if it has one.
52    pub fn slot_of(&self, r: &Register) -> Option<i32> {
53        self.slots.get(r).copied()
54    }
55}