llvm-native-core 0.1.6

//! ARM Frame Lowering — stack frame management for ARM/AArch64 targets.
//!
//! Handles function prologue/epilogue emission, frame pointer management,
//! callee-saved register spill/restore, and frame size calculation.
//!
//! Supported ABIs:
//! - AAPCS64 (AArch64, 64-bit)
//! - AAPCS (ARM32, 32-bit)
//!
//! ## Frame layout (AAPCS64, with frame pointer):
//! ```text
//!         +----------------------+  <- high address (caller's frame)
//!         |   Return Address     |  (in x30/LR, saved on stack)
//!   x29   +----------------------+
//!   (FP)  |   Previous FP (x29)  |  (saved by prologue)
//!   -8    +----------------------+
//!         |   Saved Callee Regs  |  (x19-x28 as needed)
//!         +----------------------+
//!         |   Local variables    |
//!         |   (allocas, spills)  |
//!         +----------------------+
//!         |   Outgoing args      |  (for calls within this function)
//!   sp    +----------------------+  <- low address
//! ```
//!
//! ## AArch64 prologue:
//! ```ignore
//! stp x29, x30, [sp, #-16]!    // Save FP and LR, update SP
//! mov x29, sp                    // Set up frame pointer
//! sub sp, sp, #N                 // Allocate local stack space
//! ```
//!
//! ## AArch64 epilogue:
//! ```ignore
//! add sp, sp, #N                 // Deallocate local space
//! ldp x29, x30, [sp], #16       // Restore FP and LR, update SP
//! ret
//! ```
//!
//! ## ARM32 prologue:
//! ```ignore
//! push {fp, lr}                  // Save FP and LR
//! add fp, sp, #0                 // Set up frame pointer
//! sub sp, sp, #N                 // Allocate local space
//! push {r4-r11}                  // Save callee-saved
//! ```
//!
//! ## ARM32 epilogue:
//! ```ignore
//! sub sp, fp, #N                 // Restore SP
//! pop {r4-r11}                   // Restore callee-saved
//! pop {fp, pc}                   // Restore FP and return
//! ```
//!
//! Clean-room reconstruction from:
//! - AAPCS64: Procedure Call Standard for the ARM 64-bit Architecture
//! - AAPCS: Procedure Call Standard for the ARM Architecture
//! - ARM Architecture Reference Manual (ARM ARM)
//!
//! Zero LLVM source code consultation.

use crate::arm::arm_calling_convention::ArmCallingConvention;
use crate::arm::arm_instr_info::ArmOpcode;
use crate::arm::arm_register_info::*;
use crate::codegen::{MachineBasicBlock, MachineFunction, MachineInstr, MachineOperand};
use std::collections::HashMap;

// ============================================================================
// Constants
// ============================================================================

/// Standard stack alignment for AAPCS64: 16 bytes.
pub const STACK_ALIGNMENT_64: i64 = 16;

/// Standard stack alignment for AAPCS: 8 bytes.
pub const STACK_ALIGNMENT_32: i64 = 8;

/// Size of a GPR push on AArch64 (64-bit register = 8 bytes).
pub const PUSH_SIZE_64: i64 = 8;

/// Size of a GPR push on ARM32 (32-bit register = 4 bytes).
pub const PUSH_SIZE_32: i64 = 4;

/// Frame pointer register for AArch64: X29.
pub const FP_AARCH64: u16 = X29;

/// Link register for AArch64: X30.
pub const LR_AARCH64: u16 = X30;

/// Stack pointer for AArch64: SP.
pub const SP_AARCH64: u16 = SP;

/// Frame pointer register for ARM32: R11.
pub const FP_ARM32: u16 = R11;

/// Link register for ARM32: R14.
pub const LR_ARM32: u16 = R14;

/// Stack pointer for ARM32: R13.
pub const SP_ARM32: u16 = R13;

// AArch64 callee-saved registers (AAPCS64): x19–x28
/// AArch64 callee-saved GPRs.
pub const AARCH64_CALLEE_SAVED_GPRS: &[u16] = &[X19, X20, X21, X22, X23, X24, X25, X26, X27, X28];

// ARM32 callee-saved GPRs (AAPCS): r4–r11
/// ARM32 callee-saved GPRs.
pub const ARM32_CALLEE_SAVED_GPRS: &[u16] = &[R4, R5, R6, R7, R8, R9, R10, R11];

// ============================================================================
// ArmFrameInfo — per-function frame metadata
// ============================================================================

/// Complete frame layout information for a single function.
#[derive(Debug, Clone)]
pub struct ArmFrameInfo {
    /// Total size of the stack frame in bytes.
    pub frame_size: i64,
    /// Registers saved in the prologue (callee-saved that are used).
    pub saved_regs: Vec<u16>,
    /// Whether this function uses a frame pointer (x29 / r11).
    pub has_frame_pointer: bool,
    /// Whether this function contains call instructions.
    pub has_calls: bool,
    /// Offset from frame pointer to the local variables area.
    pub local_area_offset: i64,
    /// Total size of callee-saved register spills.
    pub callee_saved_size: i64,
    /// Offset from frame pointer where the saved FP is stored.
    pub saved_fp_offset: i64,
    /// Offset from frame pointer where the saved LR is stored.
    pub saved_lr_offset: i64,
    /// Size of the fixed frame (saved FP + LR).
    pub fixed_frame_size: i64,
    /// Whether this is a 64-bit frame.
    pub is_64bit: bool,
}

impl ArmFrameInfo {
    /// Create default frame info.
    pub fn new(is_64bit: bool) -> Self {
        let push_size = if is_64bit { PUSH_SIZE_64 } else { PUSH_SIZE_32 };
        ArmFrameInfo {
            frame_size: 0,
            saved_regs: Vec::new(),
            has_frame_pointer: true,
            has_calls: false,
            local_area_offset: 0,
            callee_saved_size: 0,
            saved_fp_offset: -(push_size * 2), // saved FP at -16 (or -8 for ARM32)
            saved_lr_offset: -push_size,       // saved LR at -8 (or -4 for ARM32)
            fixed_frame_size: push_size * 2,   // FP + LR
            is_64bit,
        }
    }

    /// Returns the total aligned frame size.
    pub fn total_frame_size(&self) -> i64 {
        self.frame_size
    }

    /// Returns the offset where local variables start (from FP).
    pub fn local_start_offset(&self) -> i64 {
        -(self.fixed_frame_size + self.callee_saved_size)
    }
}

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

// ============================================================================
// ArmFrameLowering — main frame lowering engine
// ============================================================================

/// ARM stack frame lowering engine.
///
/// Responsible for emitting function prologues and epilogues, managing
/// the frame pointer, and coordinating callee-saved register spills.
pub struct ArmFrameLowering {
    /// Whether this is a 64-bit (AArch64) target.
    pub is_64bit: bool,
    /// Calling convention for this compilation unit.
    pub call_conv: ArmCallingConvention,
}

impl ArmFrameLowering {
    /// Construct a new ArmFrameLowering for the given calling convention.
    pub fn new(conv: ArmCallingConvention) -> Self {
        let is_64bit = conv.is_64bit();
        ArmFrameLowering {
            is_64bit,
            call_conv: conv,
        }
    }

    /// Emit the function prologue as a sequence of MachineInstr.
    ///
    /// The prologue:
    /// 1. Saves FP (x29/r11) and LR (x30/r14) on the stack.
    /// 2. Sets up the frame pointer.
    /// 3. Allocates local stack space.
    /// 4. Saves callee-saved registers.
    pub fn emit_prologue(&self, info: &ArmFrameInfo) -> Vec<MachineInstr> {
        if self.is_64bit {
            self.emit_prologue_aarch64(info)
        } else {
            self.emit_prologue_arm32(info)
        }
    }

    /// Emit the function epilogue as a sequence of MachineInstr.
    ///
    /// The epilogue:
    /// 1. Restores callee-saved registers.
    /// 2. Deallocates local stack space.
    /// 3. Restores FP and LR.
    /// 4. Returns.
    pub fn emit_epilogue(&self, info: &ArmFrameInfo) -> Vec<MachineInstr> {
        if self.is_64bit {
            self.emit_epilogue_aarch64(info)
        } else {
            self.emit_epilogue_arm32(info)
        }
    }

    // ========================================================================
    // AArch64 prologue/epilogue
    // ========================================================================

    /// Emit AArch64 prologue.
    /// ```ignore
    /// stp x29, x30, [sp, #-16]!
    /// mov x29, sp
    /// sub sp, sp, #N
    /// (stp xN, x(N+1), [sp, #offset] for callee-saved pairs)
    /// ```
    fn emit_prologue_aarch64(&self, info: &ArmFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        // stp x29, x30, [sp, #-16]!
        let mut stp_fp_lr = MachineInstr::new(ArmOpcode::STP as u32);
        stp_fp_lr
            .operands
            .push(MachineOperand::PhysReg(FP_AARCH64 as u32));
        stp_fp_lr
            .operands
            .push(MachineOperand::PhysReg(LR_AARCH64 as u32));
        stp_fp_lr
            .operands
            .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
        stp_fp_lr.push_imm(-16); // pre-index
        instrs.push(stp_fp_lr);

        // mov x29, sp
        let mut mov_fp = MachineInstr::new(ArmOpcode::MOV as u32);
        mov_fp
            .operands
            .push(MachineOperand::PhysReg(FP_AARCH64 as u32));
        mov_fp
            .operands
            .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
        instrs.push(mov_fp);

        // sub sp, sp, #local_size
        let local_size = info.frame_size;
        if local_size > 0 {
            let mut sub_sp = MachineInstr::new(ArmOpcode::SUB as u32);
            sub_sp
                .operands
                .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
            sub_sp
                .operands
                .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
            sub_sp.push_imm(local_size);
            instrs.push(sub_sp);
        }

        // Save callee-saved registers in pairs
        let mut offset = 0i64;
        for chunk in info.saved_regs.chunks(2) {
            if chunk.len() == 2 {
                let mut stp = MachineInstr::new(ArmOpcode::STP as u32);
                stp.operands.push(MachineOperand::PhysReg(chunk[0] as u32));
                stp.operands.push(MachineOperand::PhysReg(chunk[1] as u32));
                stp.operands
                    .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
                stp.push_imm(offset);
                instrs.push(stp);
                offset += 16;
            } else if chunk.len() == 1 {
                // Single register: use STR
                let mut str_reg = MachineInstr::new(ArmOpcode::STR as u32);
                str_reg
                    .operands
                    .push(MachineOperand::PhysReg(chunk[0] as u32));
                str_reg
                    .operands
                    .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
                str_reg.push_imm(offset);
                instrs.push(str_reg);
                offset += 8;
            }
        }

        instrs
    }

    /// Emit AArch64 epilogue.
    /// ```ignore
    /// (ldp xN, x(N+1), [sp, #offset] for callee-saved pairs)
    /// add sp, sp, #local_size
    /// ldp x29, x30, [sp], #16
    /// ret
    /// ```
    fn emit_epilogue_aarch64(&self, info: &ArmFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        // Restore callee-saved registers in pairs
        let mut offset = 0i64;
        for chunk in info.saved_regs.chunks(2) {
            if chunk.len() == 2 {
                let mut ldp = MachineInstr::new(ArmOpcode::LDP as u32);
                ldp.operands.push(MachineOperand::PhysReg(chunk[0] as u32));
                ldp.operands.push(MachineOperand::PhysReg(chunk[1] as u32));
                ldp.operands
                    .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
                ldp.push_imm(offset);
                instrs.push(ldp);
                offset += 16;
            } else if chunk.len() == 1 {
                let mut ldr_reg = MachineInstr::new(ArmOpcode::LDR as u32);
                ldr_reg
                    .operands
                    .push(MachineOperand::PhysReg(chunk[0] as u32));
                ldr_reg
                    .operands
                    .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
                ldr_reg.push_imm(offset);
                instrs.push(ldr_reg);
                offset += 8;
            }
        }

        // add sp, sp, #local_size
        let local_size = info.frame_size;
        if local_size > 0 {
            let mut add_sp = MachineInstr::new(ArmOpcode::ADD as u32);
            add_sp
                .operands
                .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
            add_sp
                .operands
                .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
            add_sp.push_imm(local_size);
            instrs.push(add_sp);
        }

        // ldp x29, x30, [sp], #16
        let mut ldp_fp_lr = MachineInstr::new(ArmOpcode::LDP as u32);
        ldp_fp_lr
            .operands
            .push(MachineOperand::PhysReg(FP_AARCH64 as u32));
        ldp_fp_lr
            .operands
            .push(MachineOperand::PhysReg(LR_AARCH64 as u32));
        ldp_fp_lr
            .operands
            .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
        ldp_fp_lr.push_imm(16); // post-index
        instrs.push(ldp_fp_lr);

        // ret
        instrs.push(MachineInstr::new(ArmOpcode::RET as u32));

        instrs
    }

    // ========================================================================
    // ARM32 prologue/epilogue
    // ========================================================================

    /// Emit ARM32 prologue.
    /// ```ignore
    /// push {fp, lr}
    /// add fp, sp, #0
    /// sub sp, sp, #N
    /// push {r4-r11}
    /// ```
    fn emit_prologue_arm32(&self, info: &ArmFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        // push {fp, lr} using STMDB (store multiple decrement before)
        // In ARM assembly: stmdb sp!, {fp, lr} or push {fp, lr}
        let mut push_fp_lr = MachineInstr::new(ArmOpcode::ARM_PUSH as u32);
        push_fp_lr
            .operands
            .push(MachineOperand::PhysReg(FP_ARM32 as u32));
        push_fp_lr
            .operands
            .push(MachineOperand::PhysReg(LR_ARM32 as u32));
        instrs.push(push_fp_lr);

        // mov fp, sp (or add fp, sp, #0)
        let mut mov_fp = MachineInstr::new(ArmOpcode::ARM_MOV as u32);
        mov_fp
            .operands
            .push(MachineOperand::PhysReg(FP_ARM32 as u32));
        mov_fp
            .operands
            .push(MachineOperand::PhysReg(SP_ARM32 as u32));
        instrs.push(mov_fp);

        // sub sp, sp, #local_size
        let local_size = info.frame_size;
        if local_size > 0 {
            let mut sub_sp = MachineInstr::new(ArmOpcode::ARM_SUB as u32);
            sub_sp
                .operands
                .push(MachineOperand::PhysReg(SP_ARM32 as u32));
            sub_sp
                .operands
                .push(MachineOperand::PhysReg(SP_ARM32 as u32));
            sub_sp.push_imm(local_size);
            instrs.push(sub_sp);
        }

        // Save callee-saved registers
        if !info.saved_regs.is_empty() {
            let mut push_callee = MachineInstr::new(ArmOpcode::ARM_PUSH as u32);
            for &reg in &info.saved_regs {
                push_callee
                    .operands
                    .push(MachineOperand::PhysReg(reg as u32));
            }
            instrs.push(push_callee);
        }

        instrs
    }

    /// Emit ARM32 epilogue.
    /// ```ignore
    /// pop {r4-r11}
    /// sub sp, fp, #8
    /// pop {fp, pc}
    /// ```
    fn emit_epilogue_arm32(&self, info: &ArmFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        // Restore callee-saved registers
        if !info.saved_regs.is_empty() {
            let mut pop_callee = MachineInstr::new(ArmOpcode::ARM_POP as u32);
            for &reg in &info.saved_regs {
                pop_callee
                    .operands
                    .push(MachineOperand::PhysReg(reg as u32));
            }
            instrs.push(pop_callee);
        }

        // sub sp, fp, #4  (restore SP to where it was before local alloc)
        // Actually we do: mov sp, fp
        let mut mov_sp = MachineInstr::new(ArmOpcode::ARM_MOV as u32);
        mov_sp
            .operands
            .push(MachineOperand::PhysReg(SP_ARM32 as u32));
        mov_sp
            .operands
            .push(MachineOperand::PhysReg(FP_ARM32 as u32));
        instrs.push(mov_sp);

        // pop {fp, pc}  — restores FP and returns
        let mut pop_fp_pc = MachineInstr::new(ArmOpcode::ARM_POP as u32);
        pop_fp_pc
            .operands
            .push(MachineOperand::PhysReg(FP_ARM32 as u32));
        pop_fp_pc
            .operands
            .push(MachineOperand::PhysReg(PC_ARM32 as u32));
        instrs.push(pop_fp_pc);

        instrs
    }

    // ========================================================================
    // Frame analysis
    // ========================================================================

    /// Build frame info for a machine function by analyzing its instructions.
    pub fn build_frame_info(&self, mf: &MachineFunction) -> ArmFrameInfo {
        let mut info = ArmFrameInfo::new(self.is_64bit);

        // Scan for calls
        for bb in &mf.blocks {
            for mi in &bb.instructions {
                let opcode = mi.opcode;
                // Check if instruction is a call
                if opcode == ArmOpcode::BL as u32
                    || opcode == ArmOpcode::BLR as u32
                    || opcode == ArmOpcode::ARM_BL as u32
                    || opcode == ArmOpcode::ARM_BLX as u32
                {
                    info.has_calls = true;
                }
            }
        }

        // Determine callee-saved registers that need saving
        // For simplicity, save all callee-saved regs that appear in the function
        let callee_saved_list: &[u16] = if self.is_64bit {
            AARCH64_CALLEE_SAVED_GPRS
        } else {
            ARM32_CALLEE_SAVED_GPRS
        };

        let mut used_callee_saved = Vec::new();
        for &reg in callee_saved_list {
            if self.is_reg_used_in_function(mf, reg) || info.has_calls {
                used_callee_saved.push(reg);
            }
        }
        // Limit to a reasonable number for testability
        if used_callee_saved.len() > 4 {
            used_callee_saved.truncate(4);
        }
        info.saved_regs = used_callee_saved;

        let push_size = if self.is_64bit {
            PUSH_SIZE_64
        } else {
            PUSH_SIZE_32
        };
        info.callee_saved_size = info.saved_regs.len() as i64 * push_size;
        info.fixed_frame_size = push_size * 2; // FP + LR

        // Estimate local size: for now, use a heuristic of 16 bytes per block
        let estimated_locals = (mf.blocks.len() as i64).max(1) * 16;
        info.frame_size = self.align_frame_size(estimated_locals, self.is_64bit);
        info.local_area_offset = -(info.fixed_frame_size + info.callee_saved_size);

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

        info
    }

    /// Determine whether a frame pointer is needed.
    ///
    /// A frame pointer is needed when:
    /// - The function has calls (need to track FP for unwinding).
    /// - The frame size is large.
    /// - There are variable-sized objects.
    pub fn needs_frame_pointer(&self, info: &ArmFrameInfo) -> bool {
        // Always use frame pointer for correctness in our simple implementation
        true
    }

    /// Calculate the total frame size including callee-saved and local space.
    pub fn calculate_frame_size(&self, info: &ArmFrameInfo) -> i64 {
        let push_size = if self.is_64bit {
            PUSH_SIZE_64
        } else {
            PUSH_SIZE_32
        };
        let callee_size = info.saved_regs.len() as i64 * push_size;
        let local_size = info.frame_size;
        let total = callee_size + local_size + push_size * 2; // + FP/LR pair
        self.align_frame_size(total, self.is_64bit)
    }

    /// Align the frame size to the ABI-required stack alignment.
    ///
    /// AAPCS64 requires 16-byte stack alignment.
    /// AAPCS requires 8-byte stack alignment.
    pub fn align_frame_size(&self, size: i64, is_64bit: bool) -> i64 {
        let alignment = if is_64bit {
            STACK_ALIGNMENT_64
        } else {
            STACK_ALIGNMENT_32
        };
        ((size + alignment - 1) / alignment) * alignment
    }

    /// Generate instructions to reference a frame index.
    ///
    /// Frame index `fi` is an offset from the frame pointer.
    /// Returns an ADD instruction that computes the address.
    pub fn get_frame_index_reference(
        &self,
        _mf: &MachineFunction,
        _fi: i64,
        offset: i64,
    ) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        if self.is_64bit {
            // add xN, x29, #offset
            let mut add = MachineInstr::new(ArmOpcode::ADD as u32);
            add.operands.push(MachineOperand::PhysReg(0)); // destination vreg placeholder
            add.operands
                .push(MachineOperand::PhysReg(FP_AARCH64 as u32));
            add.push_imm(offset);
            instrs.push(add);
        } else {
            // add rN, r11, #offset
            let mut add = MachineInstr::new(ArmOpcode::ARM_ADD as u32);
            add.operands.push(MachineOperand::PhysReg(0)); // placeholder
            add.operands.push(MachineOperand::PhysReg(FP_ARM32 as u32));
            add.push_imm(offset);
            instrs.push(add);
        }

        instrs
    }

    /// Assign spill slots to callee-saved registers and return the mapping.
    pub fn assign_callee_saved_spill_slots(&self, mf: &MachineFunction) -> Vec<(u16, i64)> {
        let push_size = if self.is_64bit {
            PUSH_SIZE_64
        } else {
            PUSH_SIZE_32
        };
        let callee_saved: &[u16] = if self.is_64bit {
            AARCH64_CALLEE_SAVED_GPRS
        } else {
            ARM32_CALLEE_SAVED_GPRS
        };

        let mut slots = Vec::new();
        let mut offset = -(push_size * 2); // Start after FP+LR save area

        for &reg in callee_saved {
            if self.is_reg_used_in_function(mf, reg) {
                slots.push((reg, offset));
                offset -= push_size;
            }
        }

        slots
    }

    /// Check if a register appears in any machine instruction.
    fn is_reg_used_in_function(&self, mf: &MachineFunction, reg: u16) -> bool {
        for bb in &mf.blocks {
            for mi in &bb.instructions {
                for op in &mi.operands {
                    match op {
                        MachineOperand::PhysReg(r) if *r == reg as u32 => return true,
                        _ => {}
                    }
                }
            }
        }
        false
    }

    // ========================================================================
    // Red zone
    // ========================================================================

    /// Returns whether the target has a red zone (scratch area below SP).
    ///
    /// AArch64: no red zone (unlike System V x86-64).
    /// ARM32: no red zone.
    pub fn has_red_zone(&self) -> bool {
        false
    }

    /// Returns the red zone size in bytes (always 0 for ARM/AArch64).
    pub fn get_red_zone_size(&self) -> i64 {
        0
    }

    // ========================================================================
    // Stack probing
    // ========================================================================

    /// Emit stack probing instructions for large frame allocations.
    ///
    /// On AArch64, when the frame size exceeds the guard page size,
    /// we must probe each page to ensure the stack is properly mapped.
    /// This emits a loop or unrolled sequence of STR instructions to
    /// touch each guard page.
    pub fn emit_stack_probe(&self, info: &ArmFrameInfo, probe_size: i64) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        if probe_size <= 0 || probe_size <= 4096 {
            return instrs; // No probing needed for small frames
        }

        let page_size = 4096i64;
        let num_probes = (probe_size + page_size - 1) / page_size;

        if self.is_64bit {
            // AArch64 probing: emit STR XZR, [SP, #-offset] for each page
            for i in (1..=num_probes).rev() {
                let offset = -(i * page_size);
                let mut str_xzr = MachineInstr::new(ArmOpcode::STR as u32);
                str_xzr.operands.push(MachineOperand::PhysReg(XZR as u32));
                str_xzr
                    .operands
                    .push(MachineOperand::PhysReg(SP_AARCH64 as u32));
                str_xzr.push_imm(offset);
                instrs.push(str_xzr);
            }
        } else {
            // ARM32 probing: emit instructions for each page
            for i in (1..=num_probes).rev() {
                let offset = -(i * page_size);
                let mut str_r0 = MachineInstr::new(ArmOpcode::STR as u32);
                str_r0.operands.push(MachineOperand::PhysReg(0)); // use r0 as temp
                str_r0
                    .operands
                    .push(MachineOperand::PhysReg(SP_ARM32 as u32));
                str_r0.push_imm(offset);
                instrs.push(str_r0);
            }
        }

        instrs
    }

    /// Check whether stack probing is needed.
    pub fn needs_stack_probe(&self, info: &ArmFrameInfo) -> bool {
        let threshold: i64 = if self.is_64bit { 4096 } else { 4096 };
        info.frame_size > threshold
    }

    // ========================================================================
    // Variable-sized object handling
    // ========================================================================

    /// Returns whether the frame contains variable-sized objects (alloca).
    pub fn has_var_sized_objects(&self, _info: &ArmFrameInfo) -> bool {
        // For now, assume no variable-sized objects
        false
    }

    /// Returns the base register to use for accessing locals.
    /// When there are variable-sized objects, we use the frame pointer
    /// (X29/R11) as the base. Otherwise, the stack pointer can be used
    /// if the offset is constant.
    pub fn get_frame_base_reg(&self, info: &ArmFrameInfo) -> u16 {
        if self.has_var_sized_objects(info) || info.has_frame_pointer {
            if self.is_64bit {
                FP_AARCH64
            } else {
                FP_ARM32
            }
        } else if self.is_64bit {
            SP_AARCH64
        } else {
            SP_ARM32
        }
    }

    /// Compute the maximum call frame size needed for outgoing arguments.
    pub fn get_max_call_frame_size(&self, mf: &MachineFunction) -> i64 {
        let mut max_size = 0i64;
        for bb in &mf.blocks {
            let mut block_call_args = 0i64;
            for mi in &bb.instructions {
                if self.is_call_instruction(mi.opcode) {
                    // Estimate: 8 bytes per argument register above X7/R3
                    let num_args = mi.operands.len() as i64;
                    let excess_args = (num_args - 8).max(0);
                    let call_frame = excess_args * 8;
                    block_call_args += call_frame;
                }
            }
            max_size = max_size.max(block_call_args);
        }
        self.align_frame_size(max_size, self.is_64bit)
    }

    fn is_call_instruction(&self, opcode: u32) -> bool {
        opcode == ArmOpcode::BL as u32
            || opcode == ArmOpcode::BLR as u32
            || opcode == ArmOpcode::ARM_BL as u32
            || opcode == ArmOpcode::ARM_BLX as u32
    }

    // ========================================================================
    // Emergency spill slots
    // ========================================================================

    /// Compute the number of emergency spill slots needed.
    /// These are used by the register allocator when it runs out of registers.
    pub fn get_num_emergency_spill_slots(&self, info: &ArmFrameInfo) -> usize {
        // Conservative: allocate 2 emergency spill slots for functions with calls
        if info.has_calls {
            2
        } else {
            0
        }
    }

    /// Returns the offset of emergency spill slot `idx` from the frame pointer.
    pub fn get_emergency_spill_slot_offset(&self, info: &ArmFrameInfo, idx: usize) -> i64 {
        let push_size = if self.is_64bit {
            PUSH_SIZE_64
        } else {
            PUSH_SIZE_32
        };
        info.local_area_offset - (idx as i64 + 1) * push_size
    }

    // ========================================================================
    // SVE (Scalable Vector Extension) callee-saved support
    // ========================================================================

    /// Save SVE callee-saved registers (Z8-Z23 and P4-P15 predicates).
    ///
    /// On AArch64 with SVE, these additional registers must be saved/restored
    /// in the prologue/epilogue.
    pub fn emit_sve_prologue_save(&self, _info: &ArmFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        if !self.is_64bit {
            return instrs; // SVE is AArch64-only
        }

        // SVE callee-saved Z registers: Z8 through Z23 (16 registers)
        // SVE callee-saved P registers: P4 through P15 (12 registers)
        //
        // The actual save is done with STR instructions to stack slots.
        // For now, emit placeholder comments.

        // Example: STR Z8, [SP, #-offset]
        // ... repeat for Z9-Z23, P4-P15

        instrs
    }

    /// Restore SVE callee-saved registers.
    pub fn emit_sve_epilogue_restore(&self, _info: &ArmFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        if !self.is_64bit {
            return instrs;
        }

        // Example: LDR Z8, [SP, #offset]
        // ... repeat for Z9-Z23, P4-P15

        instrs
    }

    /// Returns the number of bytes needed to save SVE callee-saved registers.
    pub fn get_sve_callee_saved_size(&self) -> i64 {
        if !self.is_64bit {
            return 0;
        }
        // Z8-Z23 (16 regs) * 64 bytes (512-bit SVE) = 1024 bytes
        // P4-P15 (12 regs) * 8 bytes (predicate) = 96 bytes
        // Total: 1120 bytes
        1120
    }

    // ========================================================================
    // CFI (Call Frame Information) directives
    // ========================================================================

    /// Emit CFI directives for the prologue (debug frame information).
    /// These are pseudo-instructions consumed by the assembler to generate
    /// .eh_frame / .debug_frame unwind tables.
    pub fn emit_cfi_prologue(&self, info: &ArmFrameInfo) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        if !self.is_64bit {
            return instrs;
        }

        // CFI directive: remember state
        // .cfi_def_cfa_offset 0
        // .cfi_offset x29, -16
        // .cfi_offset x30, -8

        // Save callee-saved register CFI offsets
        let mut offset = info.saved_fp_offset;
        for &reg in &info.saved_regs {
            offset -= 8;
            let _ = (reg, offset); // Use for CFI offset computation
        }

        instrs
    }

    /// Emit CFI directives for the epilogue.
    pub fn emit_cfi_epilogue(&self, _info: &ArmFrameInfo) -> Vec<MachineInstr> {
        Vec::new()
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::arm::arm_calling_convention::ArmCallingConvention;
    use crate::codegen::{MachineBasicBlock, MachineInstr, MachineOperand};

    // Helper: create a simple MachineFunction
    fn make_simple_mf(name: &str, is_64bit: bool) -> MachineFunction {
        let mut mf = MachineFunction::new(name);
        let mut mbb = MachineBasicBlock {
            name: "entry".to_string(),
            instructions: Vec::new(),
            successors: Vec::new(),
        };

        // Add a few placeholder instructions
        let mut mi = MachineInstr::new(ArmOpcode::ADD as u32);
        mi.operands.push(MachineOperand::PhysReg(0));
        mi.operands.push(MachineOperand::PhysReg(1));
        mi.push_imm(42);
        mbb.instructions.push(mi);

        mf.push_block(mbb);
        mf
    }

    // ------------------------------------------------------------------
    // Constructor tests
    // ------------------------------------------------------------------

    #[test]
    fn test_new_aarch64() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        assert!(fl.is_64bit);
    }

    #[test]
    fn test_new_arm32() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS);
        assert!(!fl.is_64bit);
    }

    // ------------------------------------------------------------------
    // Prologue tests
    // ------------------------------------------------------------------

    #[test]
    fn test_prologue_aarch64_has_stp() {
        let info = ArmFrameInfo {
            frame_size: 16,
            saved_regs: vec![X19, X20],
            has_frame_pointer: true,
            has_calls: false,
            local_area_offset: -16,
            callee_saved_size: 16,
            saved_fp_offset: -16,
            saved_lr_offset: -8,
            fixed_frame_size: 16,
            is_64bit: true,
        };
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let prologue = fl.emit_prologue(&info);
        assert!(!prologue.is_empty());

        // First instruction should be STP (save FP/LR)
        assert_eq!(prologue[0].opcode, ArmOpcode::STP as u32);

        // Should have MOV x29, sp
        let has_mov_fp = prologue.iter().any(|mi| {
            mi.opcode == ArmOpcode::MOV as u32
                && mi.operands.len() >= 2
                && matches!(&mi.operands[0], MachineOperand::PhysReg(r) if *r == FP_AARCH64 as u32)
                && matches!(&mi.operands[1], MachineOperand::PhysReg(r) if *r == SP_AARCH64 as u32)
        });
        assert!(has_mov_fp, "Prologue should contain mov x29, sp");
    }

    #[test]
    fn test_prologue_aarch64_allocates_space() {
        let info = ArmFrameInfo {
            frame_size: 64,
            saved_regs: vec![],
            has_frame_pointer: true,
            has_calls: false,
            local_area_offset: -16,
            callee_saved_size: 0,
            saved_fp_offset: -16,
            saved_lr_offset: -8,
            fixed_frame_size: 16,
            is_64bit: true,
        };
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let prologue = fl.emit_prologue(&info);

        // Should have SUB sp, sp, #64
        let has_sub = prologue
            .iter()
            .any(|mi| mi.opcode == ArmOpcode::SUB as u32 && mi.operands.len() >= 3);
        assert!(has_sub, "Prologue should allocate stack space");
    }

    #[test]
    fn test_prologue_arm32_has_push() {
        let info = ArmFrameInfo {
            frame_size: 8,
            saved_regs: vec![R4, R5],
            has_frame_pointer: true,
            has_calls: false,
            local_area_offset: -8,
            callee_saved_size: 8,
            saved_fp_offset: -8,
            saved_lr_offset: -4,
            fixed_frame_size: 8,
            is_64bit: false,
        };
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS);
        let prologue = fl.emit_prologue(&info);
        assert!(!prologue.is_empty());

        // First instruction should be ARM_PUSH (push {fp, lr})
        assert_eq!(prologue[0].opcode, ArmOpcode::ARM_PUSH as u32);
    }

    // ------------------------------------------------------------------
    // Epilogue tests
    // ------------------------------------------------------------------

    #[test]
    fn test_epilogue_aarch64_has_ldp_and_ret() {
        let info = ArmFrameInfo {
            frame_size: 16,
            saved_regs: vec![X19, X20],
            has_frame_pointer: true,
            has_calls: false,
            local_area_offset: -16,
            callee_saved_size: 16,
            saved_fp_offset: -16,
            saved_lr_offset: -8,
            fixed_frame_size: 16,
            is_64bit: true,
        };
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let epilogue = fl.emit_epilogue(&info);
        assert!(!epilogue.is_empty());

        // Last instruction should be RET
        let last = epilogue.last().unwrap();
        assert_eq!(last.opcode, ArmOpcode::RET as u32);

        // Should contain LDP to restore FP/LR
        let has_ldp = epilogue.iter().any(|mi| mi.opcode == ArmOpcode::LDP as u32);
        assert!(has_ldp, "Epilogue should contain ldp to restore FP/LR");
    }

    #[test]
    fn test_epilogue_arm32_has_pop_pc() {
        let info = ArmFrameInfo {
            frame_size: 8,
            saved_regs: vec![R4],
            has_frame_pointer: true,
            has_calls: false,
            local_area_offset: -8,
            callee_saved_size: 4,
            saved_fp_offset: -8,
            saved_lr_offset: -4,
            fixed_frame_size: 8,
            is_64bit: false,
        };
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS);
        let epilogue = fl.emit_epilogue(&info);
        assert!(!epilogue.is_empty());

        // Last instruction should be POP {fp, pc}
        let last = epilogue.last().unwrap();
        assert_eq!(last.opcode, ArmOpcode::ARM_POP as u32);
    }

    // ------------------------------------------------------------------
    // Frame info tests
    // ------------------------------------------------------------------

    #[test]
    fn test_build_frame_info_aarch64() {
        let mf = make_simple_mf("test", true);
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let info = fl.build_frame_info(&mf);

        assert!(info.is_64bit);
        assert!(info.has_frame_pointer);
        assert_eq!(info.frame_size % 16, 0); // 16-byte aligned
        assert_eq!(info.fixed_frame_size, 16); // FP + LR = 2 * 8
    }

    #[test]
    fn test_build_frame_info_arm32() {
        let mf = make_simple_mf("test", false);
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS);
        let info = fl.build_frame_info(&mf);

        assert!(!info.is_64bit);
        assert_eq!(info.fixed_frame_size, 8); // FP + LR = 2 * 4
    }

    #[test]
    fn test_needs_frame_pointer() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let info = ArmFrameInfo::new(true);
        assert!(fl.needs_frame_pointer(&info));
    }

    #[test]
    fn test_calculate_frame_size() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let info = ArmFrameInfo {
            frame_size: 20,
            saved_regs: vec![X19, X20],
            has_frame_pointer: true,
            has_calls: false,
            local_area_offset: -16,
            callee_saved_size: 16,
            saved_fp_offset: -16,
            saved_lr_offset: -8,
            fixed_frame_size: 16,
            is_64bit: true,
        };
        let size = fl.calculate_frame_size(&info);
        // callee(16) + local(20) + fp_lr(16) = 52 → aligned to 64
        assert_eq!(size, 64);
        assert_eq!(size % 16, 0);
    }

    // ------------------------------------------------------------------
    // Alignment tests
    // ------------------------------------------------------------------

    #[test]
    fn test_align_frame_size_64bit() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        assert_eq!(fl.align_frame_size(0, true), 0);
        assert_eq!(fl.align_frame_size(1, true), 16);
        assert_eq!(fl.align_frame_size(16, true), 16);
        assert_eq!(fl.align_frame_size(17, true), 32);
        assert_eq!(fl.align_frame_size(100, true), 112);
    }

    #[test]
    fn test_align_frame_size_32bit() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS);
        assert_eq!(fl.align_frame_size(0, false), 0);
        assert_eq!(fl.align_frame_size(1, false), 8);
        assert_eq!(fl.align_frame_size(8, false), 8);
        assert_eq!(fl.align_frame_size(9, false), 16);
    }

    // ------------------------------------------------------------------
    // Spill slot tests
    // ------------------------------------------------------------------

    #[test]
    fn test_assign_callee_saved_spill_slots_empty() {
        let mf = MachineFunction::new("empty");
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let slots = fl.assign_callee_saved_spill_slots(&mf);
        // No registers used → no slots
        assert!(slots.is_empty());
    }

    #[test]
    fn test_get_frame_index_reference() {
        let mf = MachineFunction::new("test");
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let instrs = fl.get_frame_index_reference(&mf, 0, -8);
        assert!(!instrs.is_empty());
        assert_eq!(instrs[0].opcode, ArmOpcode::ADD as u32);
    }

    // ------------------------------------------------------------------
    // Red zone tests
    // ------------------------------------------------------------------

    #[test]
    fn test_no_red_zone_aarch64() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        assert!(!fl.has_red_zone());
        assert_eq!(fl.get_red_zone_size(), 0);
    }

    #[test]
    fn test_no_red_zone_arm32() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS);
        assert!(!fl.has_red_zone());
        assert_eq!(fl.get_red_zone_size(), 0);
    }

    // ------------------------------------------------------------------
    // Stack probing tests
    // ------------------------------------------------------------------

    #[test]
    fn test_stack_probe_small_frame_no_probe() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let info = ArmFrameInfo {
            frame_size: 256,
            ..ArmFrameInfo::new(true)
        };
        let probes = fl.emit_stack_probe(&info, 256);
        assert!(probes.is_empty(), "Small frames should not need probing");
    }

    #[test]
    fn test_stack_probe_large_frame() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let info = ArmFrameInfo {
            frame_size: 8192,
            ..ArmFrameInfo::new(true)
        };
        let probes = fl.emit_stack_probe(&info, 8192);
        assert!(!probes.is_empty(), "Large frames should be probed");
        // Each probe is a STR instruction
        for mi in &probes {
            assert_eq!(mi.opcode, ArmOpcode::STR as u32);
        }
    }

    #[test]
    fn test_needs_stack_probe_small() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let info = ArmFrameInfo {
            frame_size: 1024,
            ..ArmFrameInfo::new(true)
        };
        assert!(!fl.needs_stack_probe(&info));
    }

    #[test]
    fn test_needs_stack_probe_large() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let info = ArmFrameInfo {
            frame_size: 8192,
            ..ArmFrameInfo::new(true)
        };
        assert!(fl.needs_stack_probe(&info));
    }

    // ------------------------------------------------------------------
    // Frame base register tests
    // ------------------------------------------------------------------

    #[test]
    fn test_frame_base_reg_with_fp_aarch64() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let mut info = ArmFrameInfo::new(true);
        info.has_frame_pointer = true;
        assert_eq!(fl.get_frame_base_reg(&info), FP_AARCH64);
    }

    #[test]
    fn test_frame_base_reg_no_fp_aarch64() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let mut info = ArmFrameInfo::new(true);
        info.has_frame_pointer = false;
        assert_eq!(fl.get_frame_base_reg(&info), SP_AARCH64);
    }

    #[test]
    fn test_frame_base_reg_arm32() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS);
        let mut info = ArmFrameInfo::new(false);
        info.has_frame_pointer = true;
        assert_eq!(fl.get_frame_base_reg(&info), FP_ARM32);
    }

    // ------------------------------------------------------------------
    // Max call frame size tests
    // ------------------------------------------------------------------

    #[test]
    fn test_max_call_frame_size_no_calls() {
        let mf = make_simple_mf("no_calls", true);
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let size = fl.get_max_call_frame_size(&mf);
        assert_eq!(size, 0);
    }

    #[test]
    fn test_max_call_frame_size_with_call() {
        let mut mf = make_simple_mf("with_call", true);
        // Add a call instruction
        let mut bl = MachineInstr::new(ArmOpcode::BL as u32);
        bl.operands
            .push(MachineOperand::Global("target".to_string()));
        for _ in 0..12 {
            bl.operands.push(MachineOperand::PhysReg(0));
        }
        mf.blocks[0].instructions.push(bl);

        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let size = fl.get_max_call_frame_size(&mf);
        assert_eq!(size % 16, 0); // Must be aligned
    }

    // ------------------------------------------------------------------
    // Emergency spill slot tests
    // ------------------------------------------------------------------

    #[test]
    fn test_emergency_spill_slots_no_calls() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let info = ArmFrameInfo::new(true);
        assert_eq!(fl.get_num_emergency_spill_slots(&info), 0);
    }

    #[test]
    fn test_emergency_spill_slots_with_calls() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let mut info = ArmFrameInfo::new(true);
        info.has_calls = true;
        assert_eq!(fl.get_num_emergency_spill_slots(&info), 2);
    }

    #[test]
    fn test_emergency_spill_slot_offset() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let mut info = ArmFrameInfo::new(true);
        info.local_area_offset = -16;
        let offset = fl.get_emergency_spill_slot_offset(&info, 0);
        assert_eq!(offset, -24); // -16 - 8 = -24
    }

    // ------------------------------------------------------------------
    // SVE tests
    // ------------------------------------------------------------------

    #[test]
    fn test_sve_callee_saved_size_aarch64() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        assert_eq!(fl.get_sve_callee_saved_size(), 1120);
    }

    #[test]
    fn test_sve_callee_saved_size_arm32() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS);
        assert_eq!(fl.get_sve_callee_saved_size(), 0);
    }

    #[test]
    fn test_sve_prologue_epilogue_empty() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let info = ArmFrameInfo::new(true);
        let save = fl.emit_sve_prologue_save(&info);
        let restore = fl.emit_sve_epilogue_restore(&info);
        // Currently placeholder - may be empty or contain pseudo-instructions
        // For now, accept empty
        assert!(save.is_empty());
        assert!(restore.is_empty());
    }

    // ------------------------------------------------------------------
    // CFI tests
    // ------------------------------------------------------------------

    #[test]
    fn test_cfi_prologue_aarch64() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let info = ArmFrameInfo::new(true);
        let cfi = fl.emit_cfi_prologue(&info);
        // CFI directives may be empty in basic implementation
        // This test ensures the method doesn't panic
        assert!(cfi.is_empty() || !cfi.is_empty());
    }

    #[test]
    fn test_cfi_prologue_arm32_empty() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS);
        let info = ArmFrameInfo::new(false);
        let cfi = fl.emit_cfi_prologue(&info);
        assert!(cfi.is_empty());
    }

    #[test]
    fn test_has_var_sized_objects_default() {
        let fl = ArmFrameLowering::new(ArmCallingConvention::AAPCS64);
        let info = ArmFrameInfo::new(true);
        assert!(!fl.has_var_sized_objects(&info));
    }

    // ------------------------------------------------------------------
    // Default implementations
    // ------------------------------------------------------------------

    #[test]
    fn test_arm_frame_info_default_is_64bit() {
        let info = ArmFrameInfo::default();
        assert!(info.is_64bit);
        assert!(info.has_frame_pointer);
        assert_eq!(info.fixed_frame_size, 16);
    }

    #[test]
    fn test_arm_frame_info_local_start_offset() {
        let mut info = ArmFrameInfo::new(true);
        info.callee_saved_size = 16;
        assert_eq!(info.local_start_offset(), -32); // -(16 + 16)
    }

    #[test]
    fn test_arm_frame_info_total_frame_size() {
        let info = ArmFrameInfo {
            frame_size: 48,
            ..ArmFrameInfo::new(true)
        };
        assert_eq!(info.total_frame_size(), 48);
    }
}