llvm-native-core 0.1.16

//! RISC-V Frame Lowering — stack frame management per psABI.
//!
//! Clean-room behavioral reconstruction from the RISC-V psABI and
//! the RISC-V ISA Manual. Zero LLVM source consultation.
//!
//! This module emits function prologues and epilogues, computes
//! frame layouts, handles callee-saved register spills, and
//! manages the frame pointer.

use super::riscv_instr_info::RiscVOpcode;
use super::riscv_register_info::*;
use crate::codegen::*;

// ---------------------------------------------------------------------------
// Stack layout constants
// ---------------------------------------------------------------------------

/// Stack alignment mandated by the RISC-V psABI (16 bytes for both
/// RV32 and RV64).
const STACK_ALIGNMENT: i64 = 16;

/// The offset (in bytes) where the return address is saved relative
/// to the stack pointer after the prologue.  On RV64, RA is saved
/// at `SP + frame_size - 8`.  On RV32, at `SP + frame_size - 4`.
pub const RA_SAVE_OFFSET_FROM_SP: i64 = -8; // relative to original SP

// ---------------------------------------------------------------------------
// RiscVFrameInfo
// ---------------------------------------------------------------------------

/// Holds all computed frame information for a function.
///
/// # Fields
/// * `frame_size` — total frame size in bytes (aligned to 16).
/// * `saved_regs` — list of callee-saved physical register IDs that
///   need to be saved/restored.
/// * `has_frame_pointer` — whether a frame pointer (s0) is needed.
/// * `local_area_offset` — offset from SP to the local variable area.
/// * `callee_saved_size` — total bytes consumed by callee-saved regs.
/// * `saved_ra_offset` — offset of RA save slot from SP (post-allocation).
/// * `saved_fp_offset` — offset of FP save slot from SP (post-allocation).
/// * `is_64bit` — whether this is an RV64 target.
#[derive(Debug, Clone)]
pub struct RiscVFrameInfo {
    /// Total frame size in bytes (aligned to 16).
    pub frame_size: i64,
    /// Callee-saved physical register IDs in save order.
    pub saved_regs: Vec<u16>,
    /// Whether a frame pointer (s0/x8) is needed.
    pub has_frame_pointer: bool,
    /// Whether the function makes calls (affects RA save).
    pub has_calls: bool,
    /// Offset from SP to the start of local variables.
    pub local_area_offset: i64,
    /// Total size of callee-saved register spill area.
    pub callee_saved_size: i64,
    /// Offset of RA save slot from SP after frame allocation.
    pub saved_ra_offset: i64,
    /// Offset of FP save slot from SP after frame allocation
    /// (only valid if has_frame_pointer).
    pub saved_fp_offset: i64,
    /// Fixed object area size (for byval args, etc.).
    pub fixed_frame_size: i64,
    /// Whether the target is RV64.
    pub is_64bit: bool,
}

impl RiscVFrameInfo {
    /// Create a new empty frame info with sensible defaults.
    pub fn new(is_64bit: bool) -> Self {
        RiscVFrameInfo {
            frame_size: 0,
            saved_regs: Vec::new(),
            has_frame_pointer: false,
            has_calls: false,
            local_area_offset: 0,
            callee_saved_size: 0,
            saved_ra_offset: 0,
            saved_fp_offset: 0,
            fixed_frame_size: 0,
            is_64bit,
        }
    }

    /// Returns the total frame size including fixed objects.
    pub fn total_frame_size(&self) -> i64 {
        self.frame_size + self.fixed_frame_size
    }

    /// Returns the offset from SP where local variables begin.
    pub fn local_start_offset(&self) -> i64 {
        self.local_area_offset
    }

    /// Returns the register width in bytes (8 for RV64, 4 for RV32).
    pub fn reg_width(&self) -> i64 {
        if self.is_64bit {
            8
        } else {
            4
        }
    }
}

impl Default for RiscVFrameInfo {
    fn default() -> Self {
        Self::new(true)
    }
}

// ---------------------------------------------------------------------------
// RiscVFrameLowering
// ---------------------------------------------------------------------------

/// Manages stack frame creation and teardown for RISC-V functions.
///
/// The frame lowering is responsible for:
/// 1. Computing the frame layout (size, spill slots).
/// 2. Emitting the prologue (stack allocation, RA/FP save).
/// 3. Emitting the epilogue (stack deallocation, RA/FP restore, ret).
///
/// # RV64 Prologue (canonical form):
/// ```asm
/// addi    sp, sp, -N           # allocate stack space
/// sd      ra, offset(sp)       # save return address
/// sd      s0, offset(sp)       # save frame pointer (if used)
/// addi    s0, sp, N            # set up frame pointer (if needed)
/// sd      s1, offset(sp)       # save callee-saved regs
/// ```
///
/// # RV64 Epilogue (canonical form):
/// ```asm
/// ld      s1, offset(sp)       # restore callee-saved
/// ld      s0, offset(sp)       # restore frame pointer
/// ld      ra, offset(sp)       # restore return address
/// addi    sp, sp, N            # deallocate stack
/// ret                          # jalr x0, ra, 0
/// ```
///
/// RV32 uses `sw`/`lw` instead of `sd`/`ld`.
pub struct RiscVFrameLowering {
    /// Whether this is a 64-bit RISC-V target.
    pub is_64bit: bool,
}

impl RiscVFrameLowering {
    /// Construct a new frame lowering instance.
    pub fn new(is_64bit: bool) -> Self {
        RiscVFrameLowering { is_64bit }
    }

    // ------------------------------------------------------------------
    // Prologue / Epilogue
    // ------------------------------------------------------------------

    /// Emit the function prologue as a sequence of `MachineInstr`.
    ///
    /// The prologue:
    /// 1. Allocates stack space: `addi sp, sp, -frame_size`
    /// 2. Saves RA: `sd/sw ra, offset(sp)`
    /// 3. Saves FP (if frame pointer needed): `sd/sw s0, offset(sp)`
    /// 4. Sets up FP: `addi s0, sp, frame_size`
    /// 5. Saves callee-saved registers.
    pub fn emit_prologue(&self, info: &RiscVFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let frame_size = info.frame_size;

        if frame_size == 0 && info.saved_regs.is_empty() && !info.has_calls {
            // Leaf function with no stack usage — no prologue needed.
            return instrs;
        }

        // 1. Allocate stack space
        if frame_size > 0 {
            // addi sp, sp, -frame_size
            let mut addi = MachineInstr::new(RiscVOpcode::ADDI as u32);
            addi.operands.push(MachineOperand::PhysReg(SP as u32));
            addi.operands.push(MachineOperand::PhysReg(SP as u32));
            addi.push_imm(-frame_size);
            instrs.push(addi);
        }

        // 2. Save RA
        let ra_offset = info.saved_ra_offset;
        if info.has_calls {
            if self.is_64bit {
                // sd ra, offset(sp)
                let mut sd = MachineInstr::new(RiscVOpcode::SD as u32);
                sd.operands.push(MachineOperand::PhysReg(RA as u32));
                sd.operands.push(MachineOperand::PhysReg(SP as u32));
                sd.push_imm(ra_offset);
                instrs.push(sd);
            } else {
                // sw ra, offset(sp)
                let mut sw = MachineInstr::new(RiscVOpcode::SW as u32);
                sw.operands.push(MachineOperand::PhysReg(RA as u32));
                sw.operands.push(MachineOperand::PhysReg(SP as u32));
                sw.push_imm(ra_offset);
                instrs.push(sw);
            }
        }

        // 3. Save FP if frame pointer is needed
        if info.has_frame_pointer {
            let fp_offset = info.saved_fp_offset;
            if self.is_64bit {
                // sd s0, offset(sp)
                let mut sd = MachineInstr::new(RiscVOpcode::SD as u32);
                sd.operands.push(MachineOperand::PhysReg(S0 as u32));
                sd.operands.push(MachineOperand::PhysReg(SP as u32));
                sd.push_imm(fp_offset);
                instrs.push(sd);
            } else {
                // sw s0, offset(sp)
                let mut sw = MachineInstr::new(RiscVOpcode::SW as u32);
                sw.operands.push(MachineOperand::PhysReg(S0 as u32));
                sw.operands.push(MachineOperand::PhysReg(SP as u32));
                sw.push_imm(fp_offset);
                instrs.push(sw);
            }

            // 4. Set up frame pointer: addi s0, sp, frame_size
            let mut addi = MachineInstr::new(RiscVOpcode::ADDI as u32);
            addi.operands.push(MachineOperand::PhysReg(S0 as u32));
            addi.operands.push(MachineOperand::PhysReg(SP as u32));
            addi.push_imm(frame_size);
            instrs.push(addi);
        }

        // 5. Save callee-saved registers
        for &(reg, offset) in &self.assign_callee_saved_spill_slots(info) {
            if self.is_64bit {
                // sd reg, offset(sp)
                let mut sd = MachineInstr::new(RiscVOpcode::SD as u32);
                sd.operands.push(MachineOperand::PhysReg(reg as u32));
                sd.operands.push(MachineOperand::PhysReg(SP as u32));
                sd.push_imm(offset);
                instrs.push(sd);
            } else {
                // sw reg, offset(sp)
                let mut sw = MachineInstr::new(RiscVOpcode::SW as u32);
                sw.operands.push(MachineOperand::PhysReg(reg as u32));
                sw.operands.push(MachineOperand::PhysReg(SP as u32));
                sw.push_imm(offset);
                instrs.push(sw);
            }
        }

        instrs
    }

    /// Emit the function epilogue as a sequence of `MachineInstr`.
    ///
    /// The epilogue:
    /// 1. Restores callee-saved registers.
    /// 2. Restores FP (if frame pointer used).
    /// 3. Restores RA.
    /// 4. Deallocates stack: `addi sp, sp, frame_size`
    /// 5. Returns: `ret`
    pub fn emit_epilogue(&self, info: &RiscVFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let frame_size = info.frame_size;

        if frame_size == 0 && info.saved_regs.is_empty() && !info.has_calls {
            // Leaf function, nothing to unwind.
            instrs.push(MachineInstr::new(RiscVOpcode::RET as u32));
            return instrs;
        }

        // 1. Restore callee-saved registers (in reverse save order)
        let spill_slots = self.assign_callee_saved_spill_slots(info);
        for &(reg, offset) in spill_slots.iter().rev() {
            if self.is_64bit {
                // ld reg, offset(sp)
                let mut ld = MachineInstr::new(RiscVOpcode::LD as u32);
                ld.operands.push(MachineOperand::PhysReg(reg as u32));
                ld.operands.push(MachineOperand::PhysReg(SP as u32));
                ld.push_imm(offset);
                instrs.push(ld);
            } else {
                // lw reg, offset(sp)
                let mut lw = MachineInstr::new(RiscVOpcode::LW as u32);
                lw.operands.push(MachineOperand::PhysReg(reg as u32));
                lw.operands.push(MachineOperand::PhysReg(SP as u32));
                lw.push_imm(offset);
                instrs.push(lw);
            }
        }

        // 2. Restore FP if frame pointer used
        if info.has_frame_pointer {
            let fp_offset = info.saved_fp_offset;
            if self.is_64bit {
                let mut ld = MachineInstr::new(RiscVOpcode::LD as u32);
                ld.operands.push(MachineOperand::PhysReg(S0 as u32));
                ld.operands.push(MachineOperand::PhysReg(SP as u32));
                ld.push_imm(fp_offset);
                instrs.push(ld);
            } else {
                let mut lw = MachineInstr::new(RiscVOpcode::LW as u32);
                lw.operands.push(MachineOperand::PhysReg(S0 as u32));
                lw.operands.push(MachineOperand::PhysReg(SP as u32));
                lw.push_imm(fp_offset);
                instrs.push(lw);
            }
        }

        // 3. Restore RA
        if info.has_calls {
            let ra_offset = info.saved_ra_offset;
            if self.is_64bit {
                // ld ra, offset(sp)
                let mut ld = MachineInstr::new(RiscVOpcode::LD as u32);
                ld.operands.push(MachineOperand::PhysReg(RA as u32));
                ld.operands.push(MachineOperand::PhysReg(SP as u32));
                ld.push_imm(ra_offset);
                instrs.push(ld);
            } else {
                // lw ra, offset(sp)
                let mut lw = MachineInstr::new(RiscVOpcode::LW as u32);
                lw.operands.push(MachineOperand::PhysReg(RA as u32));
                lw.operands.push(MachineOperand::PhysReg(SP as u32));
                lw.push_imm(ra_offset);
                instrs.push(lw);
            }
        }

        // 4. Deallocate stack
        if frame_size > 0 {
            // addi sp, sp, frame_size
            let mut addi = MachineInstr::new(RiscVOpcode::ADDI as u32);
            addi.operands.push(MachineOperand::PhysReg(SP as u32));
            addi.operands.push(MachineOperand::PhysReg(SP as u32));
            addi.push_imm(frame_size);
            instrs.push(addi);
        }

        // 5. Return
        instrs.push(MachineInstr::new(RiscVOpcode::RET as u32));

        instrs
    }

    // ------------------------------------------------------------------
    // Frame layout computation
    // ------------------------------------------------------------------

    /// Build a `RiscVFrameInfo` from the given MachineFunction.
    ///
    /// This computes:
    /// - Whether FP is needed (variable-size allocas, many spills, etc.)
    /// - Which callee-saved registers are used and need saving
    /// - Total frame size (locally allocated + spill area + RA/FP slots)
    pub fn build_frame_info(&self, mf: &MachineFunction) -> RiscVFrameInfo {
        let mut info = RiscVFrameInfo::new(self.is_64bit);

        // Collect callee-saved GPRs that are used in the function
        let used_regs = self.collect_used_regs(mf);
        let callee_saved_candidates = if self.is_64bit {
            vec![S2, S3, S4, S5, S6, S7, S8, S9, S10, S11]
        } else {
            vec![S2, S3, S4, S5, S6, S7, S8, S9, S10, S11]
        };

        for &reg in &callee_saved_candidates {
            if used_regs.contains(&(reg as u32)) {
                info.saved_regs.push(reg);
            }
        }

        // Determine if RA needs saving (any function that calls others)
        // Simplified heuristic: assume non-leaf functions need RA save
        info.has_calls = mf.blocks.iter().any(|bb| {
            bb.instructions.iter().any(|mi| {
                mi.opcode == RiscVOpcode::JAL as u32
                    || mi.opcode == RiscVOpcode::JALR as u32
                    || mi.opcode == RiscVOpcode::CALL as u32
            })
        });

        // Determine if FP is needed
        info.has_frame_pointer = self.needs_frame_pointer(&info);

        // Compute slot offsets
        let reg_width: i64 = if self.is_64bit { 8 } else { 4 };

        // Layout (from low address to high, stack grows down):
        // [locals ...] [callee-saved regs ...] [FP save] [RA save]
        //
        // Offsets from SP (after allocation, SP at low end):
        //   SP + 0              = bottom of stack
        //   SP + local_area     = start of spill area
        //   ...
        //   SP + frame_size - 2*reg_width = FP save
        //   SP + frame_size - 1*reg_width = RA save

        let mut offset = 0i64;

        // Local area offset (always at SP + 0)
        info.local_area_offset = offset;

        // Reserve space for locals (at least 0, increased by caller if needed)
        let local_size = info.frame_size - info.callee_saved_size;

        // Compute callee-saved register spill area size
        let num_saved = info.saved_regs.len() as i64;
        info.callee_saved_size = num_saved * reg_width;

        // Compute frame size
        let mut total = local_size + info.callee_saved_size;

        // Space for FP save (if needed)
        if info.has_frame_pointer {
            total += reg_width;
            info.saved_fp_offset = total - reg_width;
        } else {
            info.saved_fp_offset = 0;
        }

        // Space for RA save
        total += reg_width;
        info.saved_ra_offset = total - reg_width;

        info.frame_size = total;

        // Align frame size
        info.frame_size = Self::align_frame_size(info.frame_size);

        // Recompute FP offset with aligned size
        if info.has_frame_pointer {
            info.saved_fp_offset = info.frame_size - 2 * reg_width;
        }
        info.saved_ra_offset = info.frame_size - reg_width;

        info
    }

    /// Collect all physical register IDs used in the machine function.
    fn collect_used_regs(&self, mf: &MachineFunction) -> Vec<u32> {
        let mut used = Vec::new();
        for bb in &mf.blocks {
            for mi in &bb.instructions {
                for op in &mi.operands {
                    if let MachineOperand::PhysReg(r) = op {
                        if !used.contains(r) {
                            used.push(*r);
                        }
                    }
                }
            }
        }
        used
    }

    /// Determine whether a frame pointer (s0/x8) is required.
    ///
    /// A frame pointer is needed when:
    /// - The function has variable-sized allocas.
    /// - The frame size exceeds the 12-bit immediate range for
    ///   offset-based addressing (≥ 2048 bytes).
    /// - The function has many callee-saved registers to spill.
    /// - Debug info requires it.
    pub fn needs_frame_pointer(&self, info: &RiscVFrameInfo) -> bool {
        // Frame size ≥ 2048 requires FP to address stack slots
        // (immediate offsets are ±2^11 bytes signed)
        if info.frame_size >= 2048 {
            return true;
        }

        // Many saved regs ≥ 8 also warrants FP for simpler addressing
        if info.saved_regs.len() >= 8 {
            return true;
        }

        // If there are calls, FP may be useful for unwind info
        if info.has_calls && info.saved_regs.len() >= 4 {
            return true;
        }

        false
    }

    /// Calculate the total frame size needed for the given frame info.
    pub fn calculate_frame_size(&self, info: &RiscVFrameInfo) -> i64 {
        let mut size = info.fixed_frame_size;
        let reg_width: i64 = if self.is_64bit { 8 } else { 4 };

        // Callee-saved spill area
        size += info.saved_regs.len() as i64 * reg_width;

        // RA slot
        size += reg_width;

        // FP slot
        if info.has_frame_pointer {
            size += reg_width;
        }

        size
    }

    /// Align a frame size up to the stack alignment (16 bytes).
    fn align_frame_size(size: i64) -> i64 {
        if size <= 0 {
            return 0;
        }
        ((size + (STACK_ALIGNMENT - 1)) / STACK_ALIGNMENT) * STACK_ALIGNMENT
    }

    /// Assign spill slots (stack offsets) for callee-saved registers.
    ///
    /// Returns a vector of `(reg_id, offset_from_sp)` pairs.
    /// Offsets are relative to SP after frame allocation (SP at low
    /// end of frame).
    pub fn assign_callee_saved_spill_slots(&self, info: &RiscVFrameInfo) -> Vec<(u16, i64)> {
        let mut slots = Vec::new();
        let reg_width: i64 = if self.is_64bit { 8 } else { 4 };

        // Spill slots start after the local area
        let base_offset = info.local_area_offset;
        let mut current_offset = base_offset;

        for &reg in &info.saved_regs {
            slots.push((reg, current_offset));
            current_offset += reg_width;
        }

        slots
    }

    /// Compute stack offset for a frame index reference.
    ///
    /// Given an object at frame index, return its offset from the
    /// stack pointer (or frame pointer if FP is used).
    pub fn get_frame_index_reference(&self, info: &RiscVFrameInfo, object_offset: i64) -> i64 {
        if info.has_frame_pointer {
            // Relative to FP (s0 points to original SP/FP)
            object_offset - info.frame_size
        } else {
            // Relative to SP
            object_offset
        }
    }

    /// Returns true if the given register is used anywhere in the
    /// machine function.
    pub fn is_reg_used_in_function(&self, mf: &MachineFunction, reg: u32) -> bool {
        for bb in &mf.blocks {
            for mi in &bb.instructions {
                for op in &mi.operands {
                    if let MachineOperand::PhysReg(r) = op {
                        if *r == reg {
                            return true;
                        }
                    }
                }
            }
        }
        false
    }

    // ==================================================================
    // RV32E frame layout (16 registers, no s2-s11)
    // ==================================================================

    /// Build frame info for RV32E (embedded) target.
    /// RV32E has only x0-x15, no s2-s11 registers.
    /// Callee-saved: s0 (x8), s1 (x9).
    pub fn build_frame_info_rv32e(&self, mf: &MachineFunction) -> RiscVFrameInfo {
        let mut info = RiscVFrameInfo::new(false);
        info.is_64bit = false;

        // RV32E callee-saved candidates: only s0, s1
        let callee_saved_candidates: Vec<u16> = vec![S0, S1];
        let used_regs = self.collect_used_regs(mf);

        for &reg in &callee_saved_candidates {
            if used_regs.contains(&(reg as u32)) {
                info.saved_regs.push(reg);
            }
        }

        info.has_calls = mf.blocks.iter().any(|bb| {
            bb.instructions.iter().any(|mi| {
                mi.opcode == RiscVOpcode::JAL as u32
                    || mi.opcode == RiscVOpcode::JALR as u32
                    || mi.opcode == RiscVOpcode::CALL as u32
            })
        });

        info.has_frame_pointer = self.needs_frame_pointer(&info);

        let reg_width: i64 = 4; // RV32E is 32-bit
        let num_saved = info.saved_regs.len() as i64;
        info.callee_saved_size = num_saved * reg_width;

        let mut total = info.fixed_frame_size + info.callee_saved_size;

        if info.has_frame_pointer {
            total += reg_width;
            info.saved_fp_offset = total - reg_width;
        }

        total += reg_width;
        info.saved_ra_offset = total - reg_width;
        info.frame_size = Self::align_frame_size(total);

        if info.has_frame_pointer {
            info.saved_fp_offset = info.frame_size - 2 * reg_width;
        }
        info.saved_ra_offset = info.frame_size - reg_width;

        info
    }

    /// Emit prologue for RV32E with reduced register set.
    pub fn emit_prologue_rv32e(&self, info: &RiscVFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let frame_size = info.frame_size;

        if frame_size == 0 && info.saved_regs.is_empty() && !info.has_calls {
            return instrs;
        }

        // Allocate stack: addi sp, sp, -frame_size
        if frame_size > 0 {
            let mut addi = MachineInstr::new(RiscVOpcode::ADDI as u32);
            addi.operands.push(MachineOperand::PhysReg(SP as u32));
            addi.operands.push(MachineOperand::PhysReg(SP as u32));
            addi.push_imm(-frame_size);
            instrs.push(addi);
        }

        // Save RA
        if info.has_calls {
            let mut sw = MachineInstr::new(RiscVOpcode::SW as u32);
            sw.operands.push(MachineOperand::PhysReg(RA as u32));
            sw.operands.push(MachineOperand::PhysReg(SP as u32));
            sw.push_imm(info.saved_ra_offset);
            instrs.push(sw);
        }

        // Save FP (s0)
        if info.has_frame_pointer {
            let mut sw = MachineInstr::new(RiscVOpcode::SW as u32);
            sw.operands.push(MachineOperand::PhysReg(S0 as u32));
            sw.operands.push(MachineOperand::PhysReg(SP as u32));
            sw.push_imm(info.saved_fp_offset);
            instrs.push(sw);

            // Set up FP: addi s0, sp, frame_size
            let mut addi = MachineInstr::new(RiscVOpcode::ADDI as u32);
            addi.operands.push(MachineOperand::PhysReg(S0 as u32));
            addi.operands.push(MachineOperand::PhysReg(SP as u32));
            addi.push_imm(frame_size);
            instrs.push(addi);
        }

        // Save remaining callee-saved (s1 only in RV32E)
        for &(reg, offset) in &self.assign_callee_saved_spill_slots(info) {
            let mut sw = MachineInstr::new(RiscVOpcode::SW as u32);
            sw.operands.push(MachineOperand::PhysReg(reg as u32));
            sw.operands.push(MachineOperand::PhysReg(SP as u32));
            sw.push_imm(offset);
            instrs.push(sw);
        }

        instrs
    }

    /// Emit epilogue for RV32E.
    pub fn emit_epilogue_rv32e(&self, info: &RiscVFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let frame_size = info.frame_size;

        if frame_size == 0 && info.saved_regs.is_empty() && !info.has_calls {
            instrs.push(MachineInstr::new(RiscVOpcode::RET as u32));
            return instrs;
        }

        // Restore callee-saved (reverse order)
        let spill_slots = self.assign_callee_saved_spill_slots(info);
        for &(reg, offset) in spill_slots.iter().rev() {
            let mut lw = MachineInstr::new(RiscVOpcode::LW as u32);
            lw.operands.push(MachineOperand::PhysReg(reg as u32));
            lw.operands.push(MachineOperand::PhysReg(SP as u32));
            lw.push_imm(offset);
            instrs.push(lw);
        }

        // Restore FP
        if info.has_frame_pointer {
            let mut lw = MachineInstr::new(RiscVOpcode::LW as u32);
            lw.operands.push(MachineOperand::PhysReg(S0 as u32));
            lw.operands.push(MachineOperand::PhysReg(SP as u32));
            lw.push_imm(info.saved_fp_offset);
            instrs.push(lw);
        }

        // Restore RA
        if info.has_calls {
            let mut lw = MachineInstr::new(RiscVOpcode::LW as u32);
            lw.operands.push(MachineOperand::PhysReg(RA as u32));
            lw.operands.push(MachineOperand::PhysReg(SP as u32));
            lw.push_imm(info.saved_ra_offset);
            instrs.push(lw);
        }

        // Deallocate stack
        if frame_size > 0 {
            let mut addi = MachineInstr::new(RiscVOpcode::ADDI as u32);
            addi.operands.push(MachineOperand::PhysReg(SP as u32));
            addi.operands.push(MachineOperand::PhysReg(SP as u32));
            addi.push_imm(frame_size);
            instrs.push(addi);
        }

        instrs.push(MachineInstr::new(RiscVOpcode::RET as u32));
        instrs
    }

    // ==================================================================
    // Large frame / multi-instruction prologue / epilogue
    // ==================================================================

    /// Emit prologue for large frames (> 12-bit immediate).
    /// Uses LUI + ADDI sequence to allocate large frames.
    pub fn emit_prologue_large_frame(&self, info: &RiscVFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let frame_size = info.frame_size;

        if frame_size <= 2047 {
            // Use the standard prologue for small frames
            return self.emit_prologue(info);
        }

        // Large frame: use LUI + ADDI + SUB to allocate
        // LUI tmp, upper; ADDI tmp, tmp, lower; SUB sp, sp, tmp
        let tmp = T1 as u32;
        let upper = ((frame_size + 0x800) as u64 >> 12) as i64;
        let lower = frame_size - (upper << 12);

        if upper > 0 {
            let mut lui = MachineInstr::new(RiscVOpcode::LUI as u32);
            lui.operands.push(MachineOperand::PhysReg(tmp));
            lui.push_imm(upper);
            instrs.push(lui);
        }

        if lower != 0 {
            let mut addi = MachineInstr::new(RiscVOpcode::ADDI as u32);
            addi.operands.push(MachineOperand::PhysReg(tmp));
            addi.push_reg(tmp);
            addi.push_imm(lower);
            instrs.push(addi);
        }

        // SUB sp, sp, tmp
        let mut sub = MachineInstr::new(RiscVOpcode::SUB as u32);
        sub.operands.push(MachineOperand::PhysReg(SP as u32));
        sub.operands.push(MachineOperand::PhysReg(SP as u32));
        sub.push_reg(tmp);
        instrs.push(sub);

        // Save RA
        if info.has_calls {
            let op = if self.is_64bit {
                RiscVOpcode::SD as u32
            } else {
                RiscVOpcode::SW as u32
            };
            let mut store = MachineInstr::new(op);
            store.operands.push(MachineOperand::PhysReg(RA as u32));
            store.operands.push(MachineOperand::PhysReg(SP as u32));
            store.push_imm(info.saved_ra_offset);
            instrs.push(store);
        }

        // Save FP
        if info.has_frame_pointer {
            let op = if self.is_64bit {
                RiscVOpcode::SD as u32
            } else {
                RiscVOpcode::SW as u32
            };
            let mut store = MachineInstr::new(op);
            store.operands.push(MachineOperand::PhysReg(S0 as u32));
            store.operands.push(MachineOperand::PhysReg(SP as u32));
            store.push_imm(info.saved_fp_offset);
            instrs.push(store);

            // FP = SP + frame_size (large constant needs LUI+ADDI)
            let fp_upper = ((frame_size + 0x800) as u64 >> 12) as i64;
            let fp_lower = frame_size - (fp_upper << 12);
            let tmp2 = T2 as u32;
            if fp_upper > 0 {
                let mut lui = MachineInstr::new(RiscVOpcode::LUI as u32);
                lui.operands.push(MachineOperand::PhysReg(tmp2));
                lui.push_imm(fp_upper);
                instrs.push(lui);
            }
            if fp_lower != 0 || fp_upper == 0 {
                let mut addi = MachineInstr::new(RiscVOpcode::ADDI as u32);
                addi.operands.push(MachineOperand::PhysReg(tmp2));
                if fp_upper > 0 {
                    addi.push_reg(tmp2);
                } else {
                    addi.operands.push(MachineOperand::PhysReg(ZERO as u32));
                }
                addi.push_imm(fp_lower);
                instrs.push(addi);
            }
            let mut add = MachineInstr::new(RiscVOpcode::ADD as u32);
            add.operands.push(MachineOperand::PhysReg(S0 as u32));
            add.operands.push(MachineOperand::PhysReg(SP as u32));
            add.push_reg(tmp2);
            instrs.push(add);
        }

        // Save callee-saved
        for &(reg, offset) in &self.assign_callee_saved_spill_slots(info) {
            let op = if self.is_64bit {
                RiscVOpcode::SD as u32
            } else {
                RiscVOpcode::SW as u32
            };
            let mut store = MachineInstr::new(op);
            store.operands.push(MachineOperand::PhysReg(reg as u32));
            store.operands.push(MachineOperand::PhysReg(SP as u32));
            store.push_imm(offset);
            instrs.push(store);
        }

        instrs
    }

    /// Emit epilogue for large frames.
    pub fn emit_epilogue_large_frame(&self, info: &RiscVFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let frame_size = info.frame_size;

        if frame_size <= 2047 {
            return self.emit_epilogue(info);
        }

        // Restore callee-saved (reverse)
        let spill_slots = self.assign_callee_saved_spill_slots(info);
        for &(reg, offset) in spill_slots.iter().rev() {
            let op = if self.is_64bit {
                RiscVOpcode::LD as u32
            } else {
                RiscVOpcode::LW as u32
            };
            let mut load = MachineInstr::new(op);
            load.operands.push(MachineOperand::PhysReg(reg as u32));
            load.operands.push(MachineOperand::PhysReg(SP as u32));
            load.push_imm(offset);
            instrs.push(load);
        }

        // Restore FP
        if info.has_frame_pointer {
            let op = if self.is_64bit {
                RiscVOpcode::LD as u32
            } else {
                RiscVOpcode::LW as u32
            };
            let mut load = MachineInstr::new(op);
            load.operands.push(MachineOperand::PhysReg(S0 as u32));
            load.operands.push(MachineOperand::PhysReg(SP as u32));
            load.push_imm(info.saved_fp_offset);
            instrs.push(load);
        }

        // Restore RA
        if info.has_calls {
            let op = if self.is_64bit {
                RiscVOpcode::LD as u32
            } else {
                RiscVOpcode::LW as u32
            };
            let mut load = MachineInstr::new(op);
            load.operands.push(MachineOperand::PhysReg(RA as u32));
            load.operands.push(MachineOperand::PhysReg(SP as u32));
            load.push_imm(info.saved_ra_offset);
            instrs.push(load);
        }

        // Deallocate large frame: LUI+ADDI+ADD instead of ADDI
        let tmp = T1 as u32;
        let upper = ((frame_size + 0x800) as u64 >> 12) as i64;
        let lower = frame_size - (upper << 12);

        if upper > 0 {
            let mut lui = MachineInstr::new(RiscVOpcode::LUI as u32);
            lui.operands.push(MachineOperand::PhysReg(tmp));
            lui.push_imm(upper);
            instrs.push(lui);
        }
        if lower != 0 {
            let mut addi = MachineInstr::new(RiscVOpcode::ADDI as u32);
            addi.operands.push(MachineOperand::PhysReg(tmp));
            addi.push_reg(tmp);
            addi.push_imm(lower);
            instrs.push(addi);
        }
        let mut add = MachineInstr::new(RiscVOpcode::ADD as u32);
        add.operands.push(MachineOperand::PhysReg(SP as u32));
        add.operands.push(MachineOperand::PhysReg(SP as u32));
        add.push_reg(tmp);
        instrs.push(add);

        instrs.push(MachineInstr::new(RiscVOpcode::RET as u32));
        instrs
    }

    // ==================================================================
    // Variable-sized objects (alloca / VLA) frame support
    // ==================================================================

    /// Check if the function has variable-sized allocas.
    pub fn has_variable_sized_objects(&self, mf: &MachineFunction) -> bool {
        for bb in &mf.blocks {
            for mi in &bb.instructions {
                if mi.opcode == RiscVOpcode::ADDI as u32 {
                    // Check if this ADDI allocates stack space (simplified)
                    if let Some(MachineOperand::PhysReg(r)) = mi.operands.first() {
                        if *r == SP as u32 {
                            // SP modification from a non-constant might be a VLA
                            if mi.operands.len() >= 3 {
                                if let Some(MachineOperand::Reg(_)) = mi.operands.get(1) {
                                    if let Some(MachineOperand::Reg(_)) = mi.operands.get(2) {
                                        // Non-immediate SP adjustment suggests dynamic alloc
                                        return true;
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        false
    }

    /// Build frame info for functions with variable-sized objects.
    /// Forces frame pointer usage for dynamic stack adjustments.
    pub fn build_frame_info_with_vla(&self, mf: &MachineFunction, vla_size: i64) -> RiscVFrameInfo {
        let mut info = self.build_frame_info(mf);

        // VLA functions MUST use frame pointer
        info.has_frame_pointer = true;

        // Add VLA space to fixed frame
        info.fixed_frame_size += vla_size;

        // Recalculate frame size with the extra VLA space
        let reg_width: i64 = if self.is_64bit { 8 } else { 4 };
        let mut total = info.fixed_frame_size + info.callee_saved_size;

        if info.has_frame_pointer {
            total += reg_width;
            info.saved_fp_offset = total - reg_width;
        }

        total += reg_width;
        info.saved_ra_offset = total - reg_width;
        info.frame_size = Self::align_frame_size(total);

        if info.has_frame_pointer {
            info.saved_fp_offset = info.frame_size - 2 * reg_width;
        }
        info.saved_ra_offset = info.frame_size - reg_width;

        info
    }

    /// Emit the dynamic stack adjustment for VLA at runtime.
    /// SUB sp, fp, vla_size_reg — sets SP = FP - vla_size.
    pub fn emit_vla_adjustment(&self, vla_size_reg: u32) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        // Negate the VLA size for stack grow-down
        // Ideally: SUB sp, s0, vla_size_reg
        // But we can also: SUB sp, sp, vla_size_reg (relative to current SP)
        let mut sub = MachineInstr::new(RiscVOpcode::SUB as u32);
        sub.operands.push(MachineOperand::PhysReg(SP as u32));
        sub.operands.push(MachineOperand::PhysReg(S0 as u32));
        sub.push_reg(vla_size_reg);
        instrs.push(sub);

        instrs
    }

    // ==================================================================
    // Outgoing argument area for calls
    // ==================================================================

    /// Calculate the outgoing argument area size needed for calls.
    /// RISC-V psABI: first 8 integer args in a0-a7, first 8 FP args
    /// in fa0-fa7.  Additional args go on the stack.
    pub fn calculate_outgoing_arg_area(&self, max_call_args: usize, has_fp_args: bool) -> i64 {
        let int_reg_args = 8usize;
        let fp_reg_args = if has_fp_args { 8usize } else { 0usize };

        // Args beyond register capacity go on stack
        let overflow_int = if max_call_args > int_reg_args {
            max_call_args - int_reg_args
        } else {
            0
        };

        let reg_width: i64 = if self.is_64bit { 8 } else { 4 };

        // Outgoing arg area: aligned to 16 bytes
        let raw_size = overflow_int as i64 * reg_width;
        Self::align_frame_size(raw_size)
    }

    /// Build full frame info including outgoing argument area.
    pub fn build_frame_info_with_outgoing_args(
        &self,
        mf: &MachineFunction,
        max_call_args: usize,
        has_fp_args: bool,
    ) -> RiscVFrameInfo {
        let mut info = self.build_frame_info(mf);
        let outgoing_area = self.calculate_outgoing_arg_area(max_call_args, has_fp_args);

        info.fixed_frame_size += outgoing_area;

        let reg_width: i64 = if self.is_64bit { 8 } else { 4 };
        let mut total = info.fixed_frame_size + info.callee_saved_size;

        if info.has_frame_pointer {
            total += reg_width;
            info.saved_fp_offset = total - reg_width;
        }

        total += reg_width;
        info.saved_ra_offset = total - reg_width;
        info.frame_size = Self::align_frame_size(total);

        if info.has_frame_pointer {
            info.saved_fp_offset = info.frame_size - 2 * reg_width;
        }
        info.saved_ra_offset = info.frame_size - reg_width;

        info
    }

    /// Estimate the maximum number of call arguments from a machine function.
    pub fn estimate_max_call_args(&self, mf: &MachineFunction) -> usize {
        let mut max_args = 0usize;
        for bb in &mf.blocks {
            for mi in &bb.instructions {
                if mi.opcode == RiscVOpcode::JAL as u32 || mi.opcode == RiscVOpcode::CALL as u32 {
                    // Count how many a0-a7 are written before the call
                    let arg_count = mi
                        .operands
                        .iter()
                        .filter(|op| {
                            if let MachineOperand::PhysReg(r) = op {
                                let id = *r as u16;
                                id >= A0 && id <= A7
                            } else {
                                false
                            }
                        })
                        .count();
                    max_args = max_args.max(arg_count);
                }
            }
        }
        max_args
    }

    // ==================================================================
    // Stack realignment helpers
    // ==================================================================

    /// Check if stack realignment is needed (e.g., for over-aligned locals).
    pub fn needs_stack_realignment(&self, local_alignment: u32) -> bool {
        local_alignment > STACK_ALIGNMENT as u32
    }

    /// Calculate the stack adjustment needed for realignment.
    pub fn calculate_realignment_offset(&self, current_offset: i64, alignment: i64) -> i64 {
        if alignment == 0 {
            return 0;
        }
        let misalignment = current_offset % alignment;
        if misalignment == 0 {
            0
        } else {
            alignment - misalignment
        }
    }

    /// Emit stack realignment prologue sequence.
    /// AND sp, sp, -alignment  (clears lower bits for alignment)
    pub fn emit_stack_realignment(&self, alignment: i64) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let tmp = T1 as u32;

        // LI tmp, -alignment
        let neg_align = -alignment;
        if neg_align >= -2048 && neg_align <= 2047 {
            let mut addi = MachineInstr::new(RiscVOpcode::ADDI as u32);
            addi.operands.push(MachineOperand::PhysReg(tmp));
            addi.operands.push(MachineOperand::PhysReg(ZERO as u32));
            addi.push_imm(neg_align);
            instrs.push(addi);
        } else {
            // LUI + ADDI for large alignment masks
            let upper = ((neg_align as u64 + 0x800) >> 12) as i64;
            let mut lui = MachineInstr::new(RiscVOpcode::LUI as u32);
            lui.operands.push(MachineOperand::PhysReg(tmp));
            lui.push_imm(upper);
            instrs.push(lui);
            let lower = neg_align - (upper << 12);
            if lower != 0 {
                let mut addi = MachineInstr::new(RiscVOpcode::ADDI as u32);
                addi.operands.push(MachineOperand::PhysReg(tmp));
                addi.push_reg(tmp);
                addi.push_imm(lower);
                instrs.push(addi);
            }
        }

        // AND sp, sp, tmp
        let mut and_mi = MachineInstr::new(RiscVOpcode::AND as u32);
        and_mi.operands.push(MachineOperand::PhysReg(SP as u32));
        and_mi.operands.push(MachineOperand::PhysReg(SP as u32));
        and_mi.push_reg(tmp);
        instrs.push(and_mi);

        instrs
    }

    // ==================================================================
    // Save/restore via libcalls (for large callee-saved sets)
    // ==================================================================

    /// Determine if it's profitable to use save/restore libcalls.
    /// This is common on embedded RISC-V targets like SiFive E-series
    /// where code size matters more than speed.
    pub fn should_use_save_restore_libcall(&self, num_saved: usize) -> bool {
        // Use libcall when saving 5+ registers (libcall overhead ~4 insns)
        num_saved >= 5
    }

    /// Emit a call to __riscv_save_N prologue helper.
    pub fn emit_save_libcall(&self, num_regs: usize) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let symbol = format!("__riscv_save_{}", num_regs);

        // AUIPC ra, %pcrel_hi(symbol)
        let mut auipc = MachineInstr::new(RiscVOpcode::AUIPC as u32);
        auipc.operands.push(MachineOperand::PhysReg(RA as u32));
        auipc
            .operands
            .push(MachineOperand::Global(format!("%pcrel_hi({})", symbol)));
        instrs.push(auipc);

        // JALR ra, ra, %pcrel_lo(symbol)
        let mut jalr = MachineInstr::new(RiscVOpcode::JALR as u32);
        jalr.operands.push(MachineOperand::PhysReg(RA as u32));
        jalr.operands.push(MachineOperand::PhysReg(RA as u32));
        jalr.operands
            .push(MachineOperand::Global(format!("%pcrel_lo({})", symbol)));
        instrs.push(jalr);

        instrs
    }

    /// Emit a call to __riscv_restore_N epilogue helper.
    pub fn emit_restore_libcall(&self, num_regs: usize) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let symbol = format!("__riscv_restore_{}", num_regs);

        // TAIL __riscv_restore_N  (tail-call to restore + ret)
        let mut auipc = MachineInstr::new(RiscVOpcode::AUIPC as u32);
        auipc.operands.push(MachineOperand::PhysReg(T1 as u32));
        auipc
            .operands
            .push(MachineOperand::Global(format!("%pcrel_hi({})", symbol)));
        instrs.push(auipc);

        let mut jalr = MachineInstr::new(RiscVOpcode::JALR as u32);
        jalr.operands.push(MachineOperand::PhysReg(ZERO as u32));
        jalr.operands.push(MachineOperand::PhysReg(T1 as u32));
        jalr.operands
            .push(MachineOperand::Global(format!("%pcrel_lo({})", symbol)));
        instrs.push(jalr);

        instrs
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    /// Helper: create a simple MachineFunction with a given name.
    fn make_simple_mf(name: &str) -> MachineFunction {
        MachineFunction::new(name)
    }

    /// Helper: create a MachineFunction with a block containing a call.
    fn make_mf_with_call(name: &str) -> MachineFunction {
        let mut mf = MachineFunction::new(name);
        let mut bb = MachineBasicBlock {
            name: "entry".to_string(),
            instructions: Vec::new(),
            successors: Vec::new(),
        };
        // Add a CALL instruction
        let mut call = MachineInstr::new(RiscVOpcode::CALL as u32);
        call.operands
            .push(MachineOperand::Global("helper".to_string()));
        bb.instructions.push(call);
        mf.blocks.push(bb);
        mf
    }

    // ------------------------------------------------------------------
    // Construction
    // ------------------------------------------------------------------

    #[test]
    fn test_new_rv64() {
        let fl = RiscVFrameLowering::new(true);
        assert!(fl.is_64bit);
    }

    #[test]
    fn test_new_rv32() {
        let fl = RiscVFrameLowering::new(false);
        assert!(!fl.is_64bit);
    }

    #[test]
    fn test_frame_info_default() {
        let info = RiscVFrameInfo::default();
        assert!(info.is_64bit);
        assert_eq!(info.frame_size, 0);
        assert!(!info.has_frame_pointer);
        assert_eq!(info.reg_width(), 8);
    }

    #[test]
    fn test_frame_info_rv32_reg_width() {
        let info = RiscVFrameInfo::new(false);
        assert_eq!(info.reg_width(), 4);
    }

    // ------------------------------------------------------------------
    // Prologue / Epilogue
    // ------------------------------------------------------------------

    #[test]
    fn test_prologue_rv64_empty() {
        let fl = RiscVFrameLowering::new(true);
        let info = RiscVFrameInfo::new(true);
        let instrs = fl.emit_prologue(&info);
        // Leaf function with no stack usage: no prologue
        assert!(instrs.is_empty());
    }

    #[test]
    fn test_prologue_rv64_with_frame() {
        let fl = RiscVFrameLowering::new(true);
        let mut info = RiscVFrameInfo::new(true);
        info.frame_size = 32;
        info.has_calls = true;
        info.has_frame_pointer = true;
        info.saved_ra_offset = 24;
        info.saved_fp_offset = 16;
        info.saved_regs = vec![S1, S2];

        let instrs = fl.emit_prologue(&info);
        assert!(!instrs.is_empty(), "Prologue should contain instructions");

        // First instruction should be addi sp, sp, -32
        let first = &instrs[0];
        assert_eq!(first.opcode, RiscVOpcode::ADDI as u32);
        // Should contain SD for RA
        let has_sd_ra = instrs.iter().any(|mi| {
            mi.opcode == RiscVOpcode::SD as u32
                && mi.operands.len() >= 2
                && matches!(&mi.operands[0], MachineOperand::PhysReg(r) if *r == RA as u32)
        });
        assert!(has_sd_ra);
    }

    #[test]
    fn test_prologue_rv32_with_frame() {
        let fl = RiscVFrameLowering::new(false);
        let mut info = RiscVFrameInfo::new(false);
        info.frame_size = 16;
        info.has_calls = true;
        info.saved_ra_offset = 12;
        info.saved_regs = vec![S1];

        let instrs = fl.emit_prologue(&info);
        assert!(!instrs.is_empty());
        // RV32 uses SW instead of SD
        let has_sw = instrs.iter().any(|mi| mi.opcode == RiscVOpcode::SW as u32);
        assert!(has_sw);
    }

    #[test]
    fn test_epilogue_rv64_empty() {
        let fl = RiscVFrameLowering::new(true);
        let info = RiscVFrameInfo::new(true);
        let instrs = fl.emit_epilogue(&info);
        assert_eq!(instrs.len(), 1);
        assert_eq!(instrs[0].opcode, RiscVOpcode::RET as u32);
    }

    #[test]
    fn test_epilogue_rv64_with_frame() {
        let fl = RiscVFrameLowering::new(true);
        let mut info = RiscVFrameInfo::new(true);
        info.frame_size = 32;
        info.has_calls = true;
        info.has_frame_pointer = true;
        info.saved_ra_offset = 24;
        info.saved_fp_offset = 16;
        info.saved_regs = vec![S1];

        let instrs = fl.emit_epilogue(&info);
        // Should end with RET
        let last = instrs.last().unwrap();
        assert_eq!(last.opcode, RiscVOpcode::RET as u32);

        // Should contain LD for RA
        let has_ld_ra = instrs.iter().any(|mi| {
            mi.opcode == RiscVOpcode::LD as u32
                && mi.operands.len() >= 2
                && matches!(&mi.operands[0], MachineOperand::PhysReg(r) if *r == RA as u32)
        });
        assert!(has_ld_ra);
    }

    // ------------------------------------------------------------------
    // Frame layout
    // ------------------------------------------------------------------

    #[test]
    fn test_build_frame_info_leaf() {
        let fl = RiscVFrameLowering::new(true);
        let mf = make_simple_mf("leaf");
        let info = fl.build_frame_info(&mf);
        assert!(!info.has_calls);
        assert!(!info.has_frame_pointer);
        // RA slot is always reserved in build_frame_info even for leaf
        // (in practice DAG-based lowering may optimize it out)
    }

    #[test]
    fn test_build_frame_info_with_call() {
        let fl = RiscVFrameLowering::new(true);
        let mf = make_mf_with_call("caller");
        let info = fl.build_frame_info(&mf);
        assert!(info.has_calls);
    }

    #[test]
    fn test_needs_frame_pointer_small_frame() {
        let fl = RiscVFrameLowering::new(true);
        let info = RiscVFrameInfo::new(true);
        assert!(!fl.needs_frame_pointer(&info));
    }

    #[test]
    fn test_needs_frame_pointer_large_frame() {
        let fl = RiscVFrameLowering::new(true);
        let mut info = RiscVFrameInfo::new(true);
        info.frame_size = 4096; // > 2048
        assert!(fl.needs_frame_pointer(&info));
    }

    #[test]
    fn test_needs_frame_pointer_many_saved_regs() {
        let fl = RiscVFrameLowering::new(true);
        let mut info = RiscVFrameInfo::new(true);
        info.saved_regs = (0..8).map(|i| (S2 + i as u16)).collect();
        assert!(fl.needs_frame_pointer(&info));
    }

    #[test]
    fn test_align_frame_size() {
        assert_eq!(RiscVFrameLowering::align_frame_size(0), 0);
        assert_eq!(RiscVFrameLowering::align_frame_size(1), 16);
        assert_eq!(RiscVFrameLowering::align_frame_size(15), 16);
        assert_eq!(RiscVFrameLowering::align_frame_size(16), 16);
        assert_eq!(RiscVFrameLowering::align_frame_size(17), 32);
        assert_eq!(RiscVFrameLowering::align_frame_size(32), 32);
        assert_eq!(RiscVFrameLowering::align_frame_size(-5), 0);
    }

    #[test]
    fn test_calculate_frame_size() {
        let fl = RiscVFrameLowering::new(true);
        let mut info = RiscVFrameInfo::new(true);
        info.saved_regs = vec![S1, S2];
        // 2 * 8 = 16 for callee-saved, +8 for RA = 24
        assert_eq!(fl.calculate_frame_size(&info), 24);

        info.has_frame_pointer = true;
        // +8 for FP = 32
        assert_eq!(fl.calculate_frame_size(&info), 32);
    }

    #[test]
    fn test_assign_callee_saved_spill_slots() {
        let fl = RiscVFrameLowering::new(true);
        let mut info = RiscVFrameInfo::new(true);
        info.saved_regs = vec![S1, S2, S3];
        info.local_area_offset = 0;

        let slots = fl.assign_callee_saved_spill_slots(&info);
        assert_eq!(slots.len(), 3);
        assert_eq!(slots[0], (S1, 0));
        assert_eq!(slots[1], (S2, 8));
        assert_eq!(slots[2], (S3, 16));
    }

    #[test]
    fn test_assign_callee_saved_spill_slots_rv32() {
        let fl = RiscVFrameLowering::new(false);
        let mut info = RiscVFrameInfo::new(false);
        info.saved_regs = vec![S1, S2];
        info.local_area_offset = 0;

        let slots = fl.assign_callee_saved_spill_slots(&info);
        assert_eq!(slots.len(), 2);
        assert_eq!(slots[0], (S1, 0));
        assert_eq!(slots[1], (S2, 4)); // 4-byte slots on RV32
    }

    #[test]
    fn test_get_frame_index_reference_with_fp() {
        let fl = RiscVFrameLowering::new(true);
        let mut info = RiscVFrameInfo::new(true);
        info.has_frame_pointer = true;
        info.frame_size = 48;

        // Object at offset 0 from locals (SP + 0), relative to FP:
        // 0 - 48 = -48
        let ref_offset = fl.get_frame_index_reference(&info, 0);
        assert_eq!(ref_offset, -48);
    }

    #[test]
    fn test_get_frame_index_reference_without_fp() {
        let fl = RiscVFrameLowering::new(true);
        let info = RiscVFrameInfo::new(true);
        // No FP, so reference is relative to SP directly
        let ref_offset = fl.get_frame_index_reference(&info, 8);
        assert_eq!(ref_offset, 8);
    }

    #[test]
    fn test_is_reg_used_in_function() {
        let fl = RiscVFrameLowering::new(true);
        let mut mf = MachineFunction::new("test");
        let mut bb = MachineBasicBlock {
            name: "entry".to_string(),
            instructions: Vec::new(),
            successors: Vec::new(),
        };
        let mut mi = MachineInstr::new(RiscVOpcode::ADD as u32);
        mi.operands.push(MachineOperand::PhysReg(S1 as u32));
        mi.operands.push(MachineOperand::PhysReg(A0 as u32));
        bb.instructions.push(mi);
        mf.blocks.push(bb);

        assert!(fl.is_reg_used_in_function(&mf, S1 as u32));
        assert!(fl.is_reg_used_in_function(&mf, A0 as u32));
        assert!(!fl.is_reg_used_in_function(&mf, S3 as u32));
    }

    #[test]
    fn test_prologue_contains_frame_setup() {
        let fl = RiscVFrameLowering::new(true);
        let mut info = RiscVFrameInfo::new(true);
        info.frame_size = 64;
        info.has_calls = true;
        info.has_frame_pointer = true;
        info.saved_ra_offset = 56;
        info.saved_fp_offset = 48;
        info.saved_regs = vec![S1];

        let instrs = fl.emit_prologue(&info);
        // Should contain: addi sp, sp, -64
        // Should contain: sd ra, 56(sp)
        // Should contain: sd s0, 48(sp)
        // Should contain: addi s0, sp, 64
        let has_fp_setup = instrs.iter().any(|mi| {
            mi.opcode == RiscVOpcode::ADDI as u32
                && mi.operands.len() >= 3
                && matches!(&mi.operands[0], MachineOperand::PhysReg(r) if *r == S0 as u32)
                && matches!(&mi.operands[1], MachineOperand::PhysReg(r) if *r == SP as u32)
        });
        assert!(has_fp_setup);
    }

    #[test]
    fn test_epilogue_restores_in_reverse_order() {
        let fl = RiscVFrameLowering::new(true);
        let mut info = RiscVFrameInfo::new(true);
        info.frame_size = 32;
        info.has_calls = true;
        info.saved_ra_offset = 24;
        info.saved_regs = vec![S1, S2];

        let instrs = fl.emit_epilogue(&info);
        // The callee-saved restores should be in reverse order
        // Find the LD instructions
        let lds: Vec<_> = instrs
            .iter()
            .filter(|mi| mi.opcode == RiscVOpcode::LD as u32)
            .collect();
        // Should have at least S1, S2 loads
        assert!(lds.len() >= 2);
    }

    #[test]
    fn test_total_frame_size() {
        let mut info = RiscVFrameInfo::new(true);
        info.frame_size = 32;
        info.fixed_frame_size = 16;
        assert_eq!(info.total_frame_size(), 48);
    }

    // ==================================================================
    // RV32E tests
    // ==================================================================

    #[test]
    fn test_rv32e_frame_info() {
        let fl = RiscVFrameLowering::new(false);
        let mf = make_mf_with_call("test_rv32e");
        let info = fl.build_frame_info_rv32e(&mf);
        assert_eq!(info.reg_width(), 4);
        assert!(!info.has_frame_pointer || info.frame_size >= 0);
    }

    #[test]
    fn test_rv32e_prologue_epilogue() {
        let fl = RiscVFrameLowering::new(false);
        let info = RiscVFrameInfo {
            frame_size: 16,
            saved_regs: vec![S1],
            has_frame_pointer: false,
            has_calls: true,
            local_area_offset: 0,
            callee_saved_size: 4,
            saved_ra_offset: 8,
            saved_fp_offset: 0,
            fixed_frame_size: 0,
            is_64bit: false,
        };
        let prologue = fl.emit_prologue_rv32e(&info);
        assert!(!prologue.is_empty());
        assert!(prologue
            .iter()
            .any(|mi| mi.opcode == RiscVOpcode::SW as u32));

        let epilogue = fl.emit_epilogue_rv32e(&info);
        assert!(!epilogue.is_empty());
        assert_eq!(epilogue.last().unwrap().opcode, RiscVOpcode::RET as u32);
    }

    // ==================================================================
    // Large frame tests
    // ==================================================================

    #[test]
    fn test_large_frame_prologue() {
        let fl = RiscVFrameLowering::new(true);
        let info = RiscVFrameInfo {
            frame_size: 4096,
            saved_regs: vec![S2, S3],
            has_frame_pointer: true,
            has_calls: true,
            local_area_offset: 0,
            callee_saved_size: 16,
            saved_ra_offset: 4080,
            saved_fp_offset: 4072,
            fixed_frame_size: 0,
            is_64bit: true,
        };
        let prologue = fl.emit_prologue_large_frame(&info);
        assert!(!prologue.is_empty());
        // Should have LUI instruction for large frame
        assert!(prologue
            .iter()
            .any(|mi| mi.opcode == RiscVOpcode::LUI as u32));
    }

    #[test]
    fn test_large_frame_epilogue() {
        let fl = RiscVFrameLowering::new(true);
        let info = RiscVFrameInfo {
            frame_size: 4096,
            saved_regs: vec![],
            has_frame_pointer: false,
            has_calls: false,
            local_area_offset: 0,
            callee_saved_size: 0,
            saved_ra_offset: 4088,
            saved_fp_offset: 0,
            fixed_frame_size: 0,
            is_64bit: true,
        };
        let epilogue = fl.emit_epilogue_large_frame(&info);
        assert!(!epilogue.is_empty());
        assert_eq!(epilogue.last().unwrap().opcode, RiscVOpcode::RET as u32);
    }

    // ==================================================================
    // VLA / outgoing args tests
    // ==================================================================

    #[test]
    fn test_has_variable_sized_objects() {
        let fl = RiscVFrameLowering::new(true);
        let mf = make_simple_mf("test_vla");
        assert!(!fl.has_variable_sized_objects(&mf));
    }

    #[test]
    fn test_calculate_outgoing_arg_area() {
        let fl = RiscVFrameLowering::new(true);
        // 10 integer args: 8 in regs, 2 overflow → 2 * 8 = 16 bytes
        let area = fl.calculate_outgoing_arg_area(10, false);
        assert_eq!(area, 16);
        // 5 args: all fit in regs → 0 bytes
        let area2 = fl.calculate_outgoing_arg_area(5, false);
        assert_eq!(area2, 0);
        // 0 args
        let area3 = fl.calculate_outgoing_arg_area(0, false);
        assert_eq!(area3, 0);
    }

    #[test]
    fn test_needs_stack_realignment() {
        let fl = RiscVFrameLowering::new(true);
        assert!(!fl.needs_stack_realignment(8));
        assert!(!fl.needs_stack_realignment(16));
        assert!(fl.needs_stack_realignment(32));
        assert!(fl.needs_stack_realignment(64));
    }

    #[test]
    fn test_calculate_realignment_offset() {
        let fl = RiscVFrameLowering::new(true);
        assert_eq!(fl.calculate_realignment_offset(0, 32), 0);
        assert_eq!(fl.calculate_realignment_offset(32, 32), 0);
        assert_eq!(fl.calculate_realignment_offset(40, 32), 24);
    }

    #[test]
    fn test_save_restore_libcall() {
        let fl = RiscVFrameLowering::new(true);
        assert!(!fl.should_use_save_restore_libcall(2));
        assert!(!fl.should_use_save_restore_libcall(4));
        assert!(fl.should_use_save_restore_libcall(5));
        assert!(fl.should_use_save_restore_libcall(10));
    }

    #[test]
    fn test_emit_save_libcall() {
        let fl = RiscVFrameLowering::new(true);
        let instrs = fl.emit_save_libcall(5);
        assert!(!instrs.is_empty());
        assert!(instrs
            .iter()
            .any(|mi| mi.opcode == RiscVOpcode::AUIPC as u32));
        assert!(instrs
            .iter()
            .any(|mi| mi.opcode == RiscVOpcode::JALR as u32));
    }

    #[test]
    fn test_emit_restore_libcall() {
        let fl = RiscVFrameLowering::new(true);
        let instrs = fl.emit_restore_libcall(5);
        assert!(!instrs.is_empty());
        assert_eq!(instrs.last().unwrap().opcode, RiscVOpcode::JALR as u32);
    }

    #[test]
    fn test_estimate_max_call_args() {
        let fl = RiscVFrameLowering::new(true);
        let mf = make_mf_with_call("test_args");
        let args = fl.estimate_max_call_args(&mf);
        assert!(args <= 8); // Can't exceed register count without explicit operand setup
    }
}