lamina_codegen/riscv/
frame.rs

1/// RISC-V stack frame management utilities
2use crate::riscv::target::RiscVTarget;
3use std::io::{Error, Write};
4
5pub struct RiscVFrame;
6
7/// Bytes the prologue reserves for the saved `ra` and `fp` pair.
8const SAVED_PAIR_BYTES: i32 = 16;
9
10/// Range of a 12-bit signed immediate, as used by `addi` and the load and store offsets.
11const IMM12_MIN: i64 = -2048;
12const IMM12_MAX: i64 = 2047;
13
14impl RiscVFrame {
15    /// Round the locals area up so `sp` keeps the 16-byte alignment the ABI requires.
16    ///
17    /// Callers size this as `slots * 8`, so an odd slot count would otherwise leave `sp`
18    /// 8-byte aligned for the whole call.
19    fn locals_bytes(
20        stack_size: usize,
21        callee_saved_registers: &[&str],
22        target: RiscVTarget,
23    ) -> usize {
24        let callee_saved_bytes = callee_saved_registers.len() * target.word_bytes() as usize;
25        (stack_size + callee_saved_bytes).div_ceil(16) * 16
26    }
27
28    fn callee_saved_offset(stack_size: usize, index: usize, target: RiscVTarget) -> i32 {
29        -(SAVED_PAIR_BYTES + stack_size as i32 + (index as i32 + 1) * target.word_bytes())
30    }
31
32    /// Generate function prologue
33    pub fn generate_prologue<W: Write>(
34        writer: &mut W,
35        stack_size: usize,
36        callee_saved_registers: &[&str],
37        target: RiscVTarget,
38    ) -> Result<(), Error> {
39        // 16 either way: both ABIs want sp 16-byte aligned.
40        let w = target.word_bytes();
41        let sw = target.store_word();
42        writeln!(writer, "    addi sp, sp, -{SAVED_PAIR_BYTES}")?;
43        writeln!(writer, "    {sw} ra, {}(sp)", SAVED_PAIR_BYTES - w)?;
44        writeln!(writer, "    {sw} fp, {}(sp)", SAVED_PAIR_BYTES - 2 * w)?;
45        writeln!(writer, "    addi fp, sp, {SAVED_PAIR_BYTES}")?;
46
47        // Allocate stack space for local variables if needed
48        let locals = Self::locals_bytes(stack_size, callee_saved_registers, target);
49        if locals > 0 {
50            Self::adjust_sp(writer, -(locals as i64))?;
51        }
52        for (index, register) in callee_saved_registers.iter().enumerate() {
53            let offset = Self::callee_saved_offset(stack_size, index, target);
54            Self::fp_store(writer, sw, register, offset, "t0")?;
55        }
56        Ok(())
57    }
58
59    /// Generate function epilogue
60    pub fn generate_epilogue<W: Write>(
61        writer: &mut W,
62        stack_size: usize,
63        callee_saved_registers: &[&str],
64        target: RiscVTarget,
65    ) -> Result<(), Error> {
66        let lw = target.load_word();
67        for (index, register) in callee_saved_registers.iter().enumerate().rev() {
68            let offset = Self::callee_saved_offset(stack_size, index, target);
69            Self::fp_load(writer, lw, register, offset, register)?;
70        }
71
72        // Deallocate stack space for local variables if needed
73        let locals = Self::locals_bytes(stack_size, callee_saved_registers, target);
74        if locals > 0 {
75            Self::adjust_sp(writer, locals as i64)?;
76        }
77
78        // Restore return address and frame pointer
79        let w = target.word_bytes();
80        writeln!(writer, "    {lw} ra, {}(fp)", -w)?;
81        writeln!(writer, "    {lw} fp, {}(fp)", -2 * w)?;
82        writeln!(writer, "    addi sp, sp, {SAVED_PAIR_BYTES}")?;
83        writeln!(writer, "    ret")?;
84        Ok(())
85    }
86
87    /// Like [`Self::generate_epilogue`], but jump to `target_sym` instead of `ret` (tail call).
88    pub fn generate_tail_epilogue<W: Write>(
89        writer: &mut W,
90        stack_size: usize,
91        callee_saved_registers: &[&str],
92        target_sym: &str,
93        target: RiscVTarget,
94    ) -> Result<(), Error> {
95        let lw = target.load_word();
96        for (index, register) in callee_saved_registers.iter().enumerate().rev() {
97            let offset = Self::callee_saved_offset(stack_size, index, target);
98            Self::fp_load(writer, lw, register, offset, register)?;
99        }
100        let locals = Self::locals_bytes(stack_size, callee_saved_registers, target);
101        if locals > 0 {
102            Self::adjust_sp(writer, locals as i64)?;
103        }
104        let w = target.word_bytes();
105        writeln!(writer, "    {lw} ra, {}(fp)", -w)?;
106        writeln!(writer, "    {lw} fp, {}(fp)", -2 * w)?;
107        writeln!(writer, "    addi sp, sp, {SAVED_PAIR_BYTES}")?;
108        writeln!(writer, "    j {target_sym}")?;
109        Ok(())
110    }
111
112    /// Move `sp` by `delta` bytes. Deltas outside the 12-bit range go through `t0`,
113    /// which the prologue, epilogue and call sequences do not hold a value in.
114    pub fn adjust_sp<W: Write>(writer: &mut W, delta: i64) -> Result<(), Error> {
115        if (IMM12_MIN..=IMM12_MAX).contains(&delta) {
116            writeln!(writer, "    addi sp, sp, {delta}")
117        } else {
118            Self::load_imm(writer, "t0", delta)?;
119            writeln!(writer, "    add sp, sp, t0")
120        }
121    }
122
123    /// Emit `{mnemonic} {reg}, {offset}(fp)`, building the address in `scratch` when the
124    /// offset does not fit the 12-bit immediate. `reg` is overwritten by a load, so it may
125    /// serve as the scratch for one.
126    pub fn fp_load<W: Write>(
127        writer: &mut W,
128        mnemonic: &str,
129        reg: &str,
130        offset: i32,
131        scratch: &str,
132    ) -> Result<(), Error> {
133        if (IMM12_MIN..=IMM12_MAX).contains(&i64::from(offset)) {
134            writeln!(writer, "    {mnemonic} {reg}, {offset}(fp)")
135        } else {
136            Self::load_imm(writer, scratch, i64::from(offset))?;
137            writeln!(writer, "    add {scratch}, fp, {scratch}")?;
138            writeln!(writer, "    {mnemonic} {reg}, 0({scratch})")
139        }
140    }
141
142    /// Store `reg` at `{offset}(fp)`. A store must not clobber `reg`, so when the offset is
143    /// out of range and `reg` is `t0`, `fp` itself is walked to the slot and back.
144    pub fn fp_store<W: Write>(
145        writer: &mut W,
146        mnemonic: &str,
147        reg: &str,
148        offset: i32,
149        scratch: &str,
150    ) -> Result<(), Error> {
151        if (IMM12_MIN..=IMM12_MAX).contains(&i64::from(offset)) {
152            writeln!(writer, "    {mnemonic} {reg}, {offset}(fp)")
153        } else if reg != scratch {
154            Self::fp_load(writer, mnemonic, reg, offset, scratch)
155        } else {
156            // fp + (offset + 2048) + (-2048) is fp + offset.
157            let walk = i64::from(offset) + 2048;
158            Self::walk_fp(writer, walk)?;
159            writeln!(writer, "    {mnemonic} {reg}, -2048(fp)")?;
160            Self::walk_fp(writer, -walk)
161        }
162    }
163
164    /// Materialise `fp + offset` into `dst`.
165    pub fn fp_address<W: Write>(writer: &mut W, dst: &str, offset: i32) -> Result<(), Error> {
166        if (IMM12_MIN..=IMM12_MAX).contains(&i64::from(offset)) {
167            writeln!(writer, "    addi {dst}, fp, {offset}")
168        } else {
169            Self::load_imm(writer, dst, i64::from(offset))?;
170            writeln!(writer, "    add {dst}, fp, {dst}")
171        }
172    }
173
174    /// Load `value` into `reg`. ras only encodes `li` for 12-bit immediates, so larger values
175    /// are built from `slli` and `addi` steps.
176    pub fn load_imm<W: Write>(writer: &mut W, reg: &str, value: i64) -> Result<(), Error> {
177        if (IMM12_MIN..=IMM12_MAX).contains(&value) {
178            return writeln!(writer, "    li {reg}, {value}");
179        }
180        let low = ((value & 0xfff) ^ 0x800) - 0x800;
181        Self::load_imm(writer, reg, value.wrapping_sub(low) >> 12)?;
182        writeln!(writer, "    slli {reg}, {reg}, 12")?;
183        if low != 0 {
184            writeln!(writer, "    addi {reg}, {reg}, {low}")?;
185        }
186        Ok(())
187    }
188
189    fn walk_fp<W: Write>(writer: &mut W, mut delta: i64) -> Result<(), Error> {
190        while delta != 0 {
191            let step = delta.clamp(IMM12_MIN, IMM12_MAX);
192            writeln!(writer, "    addi fp, fp, {step}")?;
193            delta -= step;
194        }
195        Ok(())
196    }
197
198    /// Offset of a local slot from `fp`.
199    ///
200    /// The prologue leaves `fp` at the caller's `sp`, with `ra` at `fp-8` and the caller's
201    /// `fp` at `fp-16`. Locals start below that pair. This used to return `fp-8` for slot
202    /// 0 and `fp-16` for slot 1, so the first local overwrote the return address and the
203    /// second overwrote the saved frame pointer.
204    pub fn calculate_stack_offset(slot_index: usize, target: RiscVTarget) -> i32 {
205        -(SAVED_PAIR_BYTES + (slot_index as i32 + 1) * target.word_bytes())
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use super::*;
212    use crate::riscv::target::Xlen;
213    use std::error::Error;
214
215    #[test]
216    fn locals_never_land_on_the_saved_pair() {
217        let rv64 = RiscVTarget::general(Xlen::Rv64);
218        for slot in 0..8 {
219            let off = RiscVFrame::calculate_stack_offset(slot, rv64);
220            assert!(off <= -24, "slot {slot} at fp{off} overlaps the saved pair");
221            assert_eq!(off % 8, 0, "slot {slot} at fp{off} is not 8-byte aligned");
222        }
223        assert_eq!(RiscVFrame::calculate_stack_offset(0, rv64), -24);
224        assert_eq!(RiscVFrame::calculate_stack_offset(1, rv64), -32);
225    }
226
227    #[test]
228    fn locals_area_keeps_sp_16_byte_aligned() {
229        let rv64 = RiscVTarget::general(Xlen::Rv64);
230        assert_eq!(RiscVFrame::locals_bytes(0, &[], rv64), 0);
231        assert_eq!(RiscVFrame::locals_bytes(8, &[], rv64), 16);
232        assert_eq!(RiscVFrame::locals_bytes(16, &[], rv64), 16);
233        assert_eq!(RiscVFrame::locals_bytes(24, &[], rv64), 32);
234        assert_eq!(RiscVFrame::locals_bytes(16, &["s1"], rv64), 32);
235    }
236
237    #[test]
238    fn prologue_and_epilogue_move_sp_by_the_same_amount() {
239        let mut pro = Vec::new();
240        let mut epi = Vec::new();
241        let rv64 = RiscVTarget::general(Xlen::Rv64);
242        RiscVFrame::generate_prologue(&mut pro, 24, &[], rv64).expect("prologue");
243        RiscVFrame::generate_epilogue(&mut epi, 24, &[], rv64).expect("epilogue");
244        let pro = String::from_utf8(pro).expect("utf8");
245        let epi = String::from_utf8(epi).expect("utf8");
246        assert!(pro.contains("addi sp, sp, -32"), "prologue was:\n{pro}");
247        assert!(epi.contains("addi sp, sp, 32"), "epilogue was:\n{epi}");
248    }
249
250    #[test]
251    fn frame_beyond_the_immediate_range_is_adjusted_through_t0() -> Result<(), Box<dyn Error>> {
252        let rv64 = RiscVTarget::general(Xlen::Rv64);
253        let mut prologue = Vec::new();
254        RiscVFrame::generate_prologue(&mut prologue, 40000, &[], rv64)?;
255        let prologue = String::from_utf8(prologue)?;
256        assert!(prologue.contains("    add sp, sp, t0\n"), "{prologue}");
257        assert_eq!(eval_li_sequence(&prologue, "t0"), -40000, "{prologue}");
258        assert!(!prologue.contains("addi sp, sp, -40000"), "{prologue}");
259        Ok(())
260    }
261
262    /// Runs the `li`, `slli` and `addi` lines that load `reg`, to check the value they build.
263    fn eval_li_sequence(asm: &str, reg: &str) -> i64 {
264        let mut value = 0i64;
265        for line in asm.lines().map(str::trim) {
266            let (op, args) = line.split_once(' ').unwrap_or((line, ""));
267            let mut parts = args.split(", ");
268            if parts.next() != Some(reg) {
269                continue;
270            }
271            match op {
272                "li" => value = parts.next().and_then(|v| v.parse().ok()).unwrap_or(0),
273                "slli" => {
274                    let shift: u32 = parts.nth(1).and_then(|v| v.parse().ok()).unwrap_or(0);
275                    value = value.wrapping_shl(shift);
276                }
277                "addi" => {
278                    let imm: i64 = parts.nth(1).and_then(|v| v.parse().ok()).unwrap_or(0);
279                    value = value.wrapping_add(imm);
280                }
281                _ => {}
282            }
283        }
284        value
285    }
286
287    #[test]
288    fn load_imm_builds_every_tested_value() -> Result<(), Box<dyn Error>> {
289        for value in [
290            0,
291            1,
292            -1,
293            2047,
294            -2048,
295            2048,
296            -2049,
297            40016,
298            -40016,
299            1 << 31,
300            i32::MIN as i64,
301            i64::MAX,
302            i64::MIN,
303        ] {
304            let mut out = Vec::new();
305            RiscVFrame::load_imm(&mut out, "t0", value)?;
306            let asm = String::from_utf8(out)?;
307            assert_eq!(eval_li_sequence(&asm, "t0"), value, "{asm}");
308        }
309        Ok(())
310    }
311
312    #[test]
313    fn out_of_range_fp_access_builds_the_address_in_scratch() -> Result<(), Box<dyn Error>> {
314        let mut out = Vec::new();
315        RiscVFrame::fp_store(&mut out, "sd", "a0", -40016, "t0")?;
316        RiscVFrame::fp_load(&mut out, "ld", "a1", -40016, "a1")?;
317        let out = String::from_utf8(out)?;
318        assert_eq!(
319            out,
320            concat!(
321                "    li t0, -10\n    slli t0, t0, 12\n    addi t0, t0, 944\n",
322                "    add t0, fp, t0\n    sd a0, 0(t0)\n",
323                "    li a1, -10\n    slli a1, a1, 12\n    addi a1, a1, 944\n",
324                "    add a1, fp, a1\n    ld a1, 0(a1)\n",
325            )
326        );
327        Ok(())
328    }
329
330    #[test]
331    fn out_of_range_store_of_the_scratch_register_walks_fp_and_back() -> Result<(), Box<dyn Error>>
332    {
333        let mut out = Vec::new();
334        RiscVFrame::fp_store(&mut out, "sd", "t0", -40016, "t0")?;
335        let out = String::from_utf8(out)?;
336        let walked: i64 = out
337            .lines()
338            .filter_map(|line| line.trim().strip_prefix("addi fp, fp, "))
339            .map(|v| v.parse::<i64>())
340            .collect::<Result<Vec<_>, _>>()?
341            .iter()
342            .sum();
343        assert_eq!(walked, 0, "fp must return to its value: {out}");
344        assert!(out.contains("sd t0, -2048(fp)"), "{out}");
345        assert!(
346            out.lines()
347                .filter_map(|line| line.trim().strip_prefix("addi fp, fp, "))
348                .all(|v| (-2048..=2047).contains(&v.parse::<i64>().unwrap_or(i64::MAX))),
349            "{out}"
350        );
351        Ok(())
352    }
353
354    #[test]
355    fn callee_saved_registers_have_frame_slots_and_are_restored() -> Result<(), Box<dyn Error>> {
356        let mut prologue = Vec::new();
357        let mut epilogue = Vec::new();
358        let rv64 = RiscVTarget::general(Xlen::Rv64);
359        RiscVFrame::generate_prologue(&mut prologue, 16, &["s1"], rv64)?;
360        RiscVFrame::generate_epilogue(&mut epilogue, 16, &["s1"], rv64)?;
361        let prologue = String::from_utf8(prologue)?;
362        let epilogue = String::from_utf8(epilogue)?;
363        assert!(prologue.contains("addi sp, sp, -32"));
364        assert!(prologue.contains("sd s1, -40(fp)"));
365        assert!(epilogue.contains("ld s1, -40(fp)"));
366        assert!(epilogue.contains("addi sp, sp, 32"));
367        Ok(())
368    }
369}
lamina_codegen/riscv/frame.rs

lamina_codegen/riscv/
frame.rs
