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, Instruction, 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    /// Registers written by some instruction. A register that is only ever
13    /// read is a stack-allocation placeholder: its slot *is* the storage, so
14    /// its address is the value. Anything else holds a value that must be
15    /// loaded before use.
16    defined: HashSet<Register>,
17}
18
19impl FrameMap {
20    /// Creates a frame map from a function, assigning stack slots to all virtual registers.
21    pub fn from_function(f: &Function) -> Self {
22        let mut regs: HashSet<Register> = HashSet::new();
23        let mut defined: HashSet<Register> = HashSet::new();
24        for p in &f.sig.params {
25            regs.insert(p.reg.clone());
26        }
27        for b in &f.blocks {
28            for ins in &b.instructions {
29                if let Some(d) = ins.def_reg() {
30                    regs.insert(d.clone());
31                    defined.insert(d.clone());
32                }
33                if let Instruction::StackAddr { slot, .. } = ins {
34                    regs.insert(Register::Virtual(*slot));
35                }
36                for u in ins.use_regs() {
37                    regs.insert(u.clone());
38                }
39            }
40        }
41
42        let mut reg_vec: Vec<Register> = regs.into_iter().collect();
43        reg_vec.sort_by(|a, b| format!("{a:?}").cmp(&format!("{b:?}")));
44
45        let mut slots = HashMap::new();
46        let mut offset: i32 = -8;
47        for r in reg_vec {
48            if let Register::Virtual(v) = &r {
49                // A register with a reservation addresses the lowest byte of a
50                // multi-slot block, so it sits at the deepest slot of that block.
51                let extra = f.stack_reservations.get(v).copied().unwrap_or(0) as i32;
52                offset -= 8 * extra;
53                slots.insert(r.clone(), offset);
54                offset -= 8;
55            }
56        }
57        let mut frame_size = -offset - 8;
58        if frame_size < 0 {
59            frame_size = 0;
60        }
61        frame_size = (frame_size + 15) & !15;
62        Self {
63            slots,
64            frame_size,
65            defined,
66        }
67    }
68
69    pub fn recompute_frame_size_from_slots(&mut self) {
70        let mut min_off = 0i32;
71        for &o in self.slots.values() {
72            if o < min_off {
73                min_off = o;
74            }
75        }
76        // The deepest slot starts at `min_off`, so the frame must span all of
77        // it; `-min_off - 8` left that slot below sp.
78        let mut frame_size = (-min_off).max(0);
79        frame_size = (frame_size + 15) & !15;
80        self.frame_size = frame_size;
81    }
82
83    /// True when the register holds a value rather than being a stack
84    /// allocation placeholder whose slot address is the value.
85    pub fn holds_value(&self, r: &Register) -> bool {
86        self.defined.contains(r)
87    }
88
89    /// Returns the stack slot offset for a register, if it has one.
90    pub fn slot_of(&self, r: &Register) -> Option<i32> {
91        self.slots.get(r).copied()
92    }
93}
lamina_codegen/aarch64/frame.rs

lamina_codegen/aarch64/
frame.rs
