llvm-native-core-ext 0.1.0

//! Prologue/Epilogue Inserter — inserts function prologue and epilogue
//! code into machine functions after register allocation.
//! Clean-room behavioral reconstruction.
//!
//! After register allocation, the prologue/epilogue inserter:
//! 1. Analyzes callee-saved register usage
//! 2. Finds scratch registers for frame setup
//! 3. Computes frame layout (spill slots, locals, callee-save area)
//! 4. Inserts prologue at function entry
//! 5. Inserts epilogue before each return instruction
//! 6. Replaces frame index operands with actual SP/FP-relative offsets
//! 7. Handles stack realignment for over-aligned locals
//! 8. Minimizes stack size via packed layout

use llvm_native_core::codegen::{MachineFunction, MachineInstr, MachineOperand, TargetRegInfo};
use std::collections::{HashMap, HashSet};

// ============================================================================
// Callee-Saved Register Analysis
// ============================================================================

/// Analysis result: which registers must be saved/restored.
#[derive(Debug, Clone)]
pub struct CalleeSavedAnalysis {
    /// Registers that must be saved (used in function body).
    pub must_save: HashSet<u32>,
    /// Total size of callee-saved area in bytes.
    pub total_save_size: u32,
    /// Offset of each saved register from the frame base.
    pub save_offsets: HashMap<u32, i64>,
    /// Registers available as scratch for frame setup.
    pub scratch_regs: Vec<u32>,
    /// Whether a frame pointer is needed.
    pub needs_frame_pointer: bool,
}

impl Default for CalleeSavedAnalysis {
    fn default() -> Self {
        CalleeSavedAnalysis {
            must_save: HashSet::new(),
            total_save_size: 0,
            save_offsets: HashMap::new(),
            scratch_regs: Vec::new(),
            needs_frame_pointer: false,
        }
    }
}

impl CalleeSavedAnalysis {
    /// Analyze a machine function to determine callee-saved register usage.
    pub fn analyze(mf: &MachineFunction, target: &TargetRegInfo, reg_size: u32) -> Self {
        let mut analysis = CalleeSavedAnalysis::default();
        let callee_saved = &target.callee_saved;

        // Scan all instructions to find which callee-saved regs are used
        let mut used_regs = HashSet::new();

        for bb in &mf.blocks {
            for mi in &bb.instructions {
                // Check operands for callee-saved register usage
                for op in &mi.operands {
                    if let MachineOperand::PhysReg(r) = *op {
                        if callee_saved.contains(&r) {
                            used_regs.insert(r);
                        }
                    }
                }
                // Also check def (output register)
                if let Some(def) = mi.def {
                    if callee_saved.contains(&def) {
                        used_regs.insert(def);
                    }
                }
            }
        }

        analysis.must_save = used_regs;
        let used_regs_ref = analysis.must_save.clone();

        // Compute save offsets (each saved reg takes one slot)
        let mut offset: i64 = 0;
        for &reg in callee_saved {
            if used_regs_ref.contains(&reg) {
                analysis.save_offsets.insert(reg, offset);
                offset -= reg_size as i64;
            }
        }
        analysis.total_save_size = (-offset) as u32;

        // Find scratch registers (allocatable but not saved)
        for &reg in &target.allocatable_regs() {
            if !used_regs_ref.contains(&reg) && !target.reserved.contains(&reg) {
                analysis.scratch_regs.push(reg);
            }
        }

        // Decide if frame pointer is needed
        analysis.needs_frame_pointer =
            !used_regs_ref.is_empty() || mf.blocks.iter().any(|bb| bb.instructions.len() > 50);

        analysis
    }
}

// ============================================================================
// Scratch Register Analysis
// ============================================================================

/// Analysis to find available scratch registers for frame setup.
#[derive(Debug, Clone)]
pub struct ScratchRegAnalysis {
    /// Registers available for use as temporaries in prologue/epilogue.
    pub available: Vec<u32>,
    /// Best scratch register for each purpose.
    pub best_for_frame_setup: Option<u32>,
    pub best_for_stack_probe: Option<u32>,
}

impl Default for ScratchRegAnalysis {
    fn default() -> Self {
        ScratchRegAnalysis {
            available: Vec::new(),
            best_for_frame_setup: None,
            best_for_stack_probe: None,
        }
    }
}

impl ScratchRegAnalysis {
    /// Find scratch registers for the function.
    pub fn find(mf: &MachineFunction, target: &TargetRegInfo, used_regs: &HashSet<u32>) -> Self {
        let mut analysis = ScratchRegAnalysis::default();

        // Find registers that are allocatable but not used
        for &reg in &target.allocatable_regs() {
            if !used_regs.contains(&reg) && !target.reserved.contains(&reg) {
                analysis.available.push(reg);
            }
        }

        // Best scratch for frame setup: prefer caller-saved regs
        // that don't need saving themselves
        if let Some(&reg) = analysis.available.first() {
            analysis.best_for_frame_setup = Some(reg);
        }

        // Best scratch for stack probe: any available reg
        analysis.best_for_stack_probe = analysis.best_for_frame_setup;

        let _ = mf; // for future analysis
        analysis
    }
}

// ============================================================================
// Frame Layout Computation
// ============================================================================

/// Complete frame layout for a function.
#[derive(Debug, Clone)]
pub struct FrameLayout {
    /// Total frame size in bytes.
    pub frame_size: i64,
    /// Offset from SP/FP to the callee-saved area.
    pub callee_saved_offset: i64,
    /// Size of the callee-saved area.
    pub callee_saved_size: i64,
    /// Offset to the local variables area.
    pub local_area_offset: i64,
    /// Size of the local variables area.
    pub local_area_size: i64,
    /// Offset to the spill slot area.
    pub spill_area_offset: i64,
    /// Size of the spill slot area.
    pub spill_area_size: i64,
    /// Offset to the outgoing argument area.
    pub outgoing_arg_offset: i64,
    /// Size of the outgoing argument area.
    pub outgoing_arg_size: i64,
    /// Whether the frame is packed (no wasted alignment padding).
    pub is_packed: bool,
    /// Required stack alignment.
    pub stack_alignment: i64,
}

impl Default for FrameLayout {
    fn default() -> Self {
        FrameLayout {
            frame_size: 0,
            callee_saved_offset: 0,
            callee_saved_size: 0,
            local_area_offset: 0,
            local_area_size: 0,
            spill_area_offset: 0,
            spill_area_size: 0,
            outgoing_arg_offset: 0,
            outgoing_arg_size: 0,
            is_packed: false,
            stack_alignment: 16,
        }
    }
}

impl FrameLayout {
    /// Compute the frame layout for a function.
    pub fn compute(
        callee_saved_size: i64,
        local_size: i64,
        spill_size: i64,
        outgoing_arg_size: i64,
        alignment: i64,
    ) -> Self {
        let mut layout = FrameLayout::default();
        layout.stack_alignment = alignment;
        layout.callee_saved_size = align_to(callee_saved_size, 8);
        layout.local_area_size = local_size;
        layout.spill_area_size = align_to(spill_size, 8);
        layout.outgoing_arg_size = align_to(outgoing_arg_size, alignment);

        // Compute offsets (stack grows downward, negative offsets)
        let mut offset: i64 = 0;

        // Callee-saved area (closest to FP)
        layout.callee_saved_offset = offset;
        offset -= layout.callee_saved_size;

        // Local variables
        layout.local_area_offset = offset;
        offset -= layout.local_area_size;

        // Spill slots
        layout.spill_area_offset = offset;
        offset -= layout.spill_area_size;

        // Outgoing arguments (furthest from FP)
        layout.outgoing_arg_offset = offset;
        offset -= layout.outgoing_arg_size;

        // Total frame size (positive, aligned)
        layout.frame_size = align_to(-offset, alignment);

        // Check if packed (no wasted space between areas)
        layout.is_packed = layout.callee_saved_size
            + layout.local_area_size
            + layout.spill_area_size
            + layout.outgoing_arg_size
            == layout.frame_size;

        layout
    }

    /// Compute offset for a specific frame object from the frame base.
    pub fn get_offset(&self, object_type: FrameObjectType) -> i64 {
        match object_type {
            FrameObjectType::CalleeSaved => self.callee_saved_offset,
            FrameObjectType::Local => self.local_area_offset,
            FrameObjectType::SpillSlot => self.spill_area_offset,
            FrameObjectType::OutgoingArg => self.outgoing_arg_offset,
        }
    }
}

/// Types of frame objects.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FrameObjectType {
    /// A callee-saved register spill slot.
    CalleeSaved,
    /// A local variable.
    Local,
    /// A spilled virtual register.
    SpillSlot,
    /// An outgoing argument slot.
    OutgoingArg,
}

/// Align a value to the given alignment.
fn align_to(value: i64, alignment: i64) -> i64 {
    if value % alignment == 0 {
        value
    } else {
        value + (alignment - (value % alignment))
    }
}

// ============================================================================
// Prologue Inserter
// ============================================================================

/// Inserts prologue code at function entry.
pub struct PrologueInserter {
    /// Target register info.
    pub target: TargetRegInfo,
    /// Frame pointer register ID.
    pub fp_reg: u32,
    /// Stack pointer register ID.
    pub sp_reg: u32,
    /// Register size on this architecture (4 or 8).
    pub reg_size: u32,
}

impl PrologueInserter {
    /// Create a new prologue inserter for the given target.
    pub fn new(target: TargetRegInfo, fp_reg: u32, sp_reg: u32) -> Self {
        PrologueInserter {
            target,
            fp_reg,
            sp_reg,
            reg_size: 8, // default to 64-bit
        }
    }

    /// Create a prologue inserter for x86-64.
    pub fn for_x86_64() -> Self {
        PrologueInserter {
            target: TargetRegInfo::for_x86_64(),
            fp_reg: 5, // RBP
            sp_reg: 4, // RSP
            reg_size: 8,
        }
    }

    /// Insert prologue code at the beginning of the function.
    ///
    /// The prologue is inserted before the first non-PHI instruction
    /// in the entry block.
    pub fn insert_prologue(
        &self,
        mf: &mut MachineFunction,
        callee_saved: &CalleeSavedAnalysis,
        frame_layout: &FrameLayout,
    ) {
        let mut prologue = Vec::new();

        // 1. Push frame pointer
        let mut push_fp = MachineInstr::new(11); // PUSH
        push_fp.operands.push(MachineOperand::PhysReg(self.fp_reg));
        prologue.push(push_fp);

        // 2. Set frame pointer
        let mut mov_fp = MachineInstr::new(1); // MOV
        mov_fp.operands.push(MachineOperand::PhysReg(self.fp_reg));
        mov_fp.operands.push(MachineOperand::PhysReg(self.sp_reg));
        prologue.push(mov_fp);

        // 3. Allocate stack space
        if frame_layout.frame_size > 0 {
            let alloc_size = frame_layout.frame_size;
            if alloc_size > 0 {
                let mut sub_sp = MachineInstr::new(3); // SUB
                sub_sp.operands.push(MachineOperand::PhysReg(self.sp_reg));
                sub_sp.operands.push(MachineOperand::Imm(alloc_size));
                prologue.push(sub_sp);
            }
        }

        // 4. Save callee-saved registers
        for &reg in &self.target.callee_saved {
            if callee_saved.must_save.contains(&reg) {
                let mut push_reg = MachineInstr::new(11); // PUSH
                push_reg.operands.push(MachineOperand::PhysReg(reg));
                prologue.push(push_reg);
            }
        }

        // Insert at start of first block
        if let Some(entry_block) = mf.blocks.first_mut() {
            for instr in prologue.into_iter().rev() {
                entry_block.instructions.insert(0, instr);
            }
        }
    }

    /// Insert prologue without frame pointer (leaf function optimization).
    pub fn insert_leaf_prologue(&self, mf: &mut MachineFunction, frame_size: i64) {
        if frame_size <= 0 {
            return;
        }

        let mut sub_sp = MachineInstr::new(3); // SUB
        sub_sp.operands.push(MachineOperand::PhysReg(self.sp_reg));
        sub_sp.operands.push(MachineOperand::Imm(frame_size));

        if let Some(entry_block) = mf.blocks.first_mut() {
            entry_block.instructions.insert(0, sub_sp);
        }
    }
}

impl Default for PrologueInserter {
    fn default() -> Self {
        PrologueInserter::for_x86_64()
    }
}

// ============================================================================
// Epilogue Inserter
// ============================================================================

/// Inserts epilogue code before each return instruction.
pub struct EpilogueInserter {
    /// Frame pointer register ID.
    pub fp_reg: u32,
    /// Stack pointer register ID.
    pub sp_reg: u32,
    /// Register size on this architecture.
    pub reg_size: u32,
}

impl EpilogueInserter {
    /// Create a new epilogue inserter.
    pub fn new(fp_reg: u32, sp_reg: u32) -> Self {
        EpilogueInserter {
            fp_reg,
            sp_reg,
            reg_size: 8,
        }
    }

    /// Create an epilogue inserter for x86-64.
    pub fn for_x86_64() -> Self {
        EpilogueInserter {
            fp_reg: 5, // RBP
            sp_reg: 4, // RSP
            reg_size: 8,
        }
    }

    /// Insert epilogue code before each return instruction.
    ///
    /// The epilogue restores the stack pointer, pops saved registers,
    /// restores the frame pointer, and returns.
    pub fn insert_epilogue(
        &self,
        mf: &mut MachineFunction,
        callee_saved: &CalleeSavedAnalysis,
        frame_layout: &FrameLayout,
        has_frame_pointer: bool,
    ) {
        let epilogue = self.build_epilogue(callee_saved, frame_layout, has_frame_pointer);

        // Insert epilogue before each RET instruction
        for bb in &mut mf.blocks {
            // Find RET instructions
            let ret_positions: Vec<usize> = bb
                .instructions
                .iter()
                .enumerate()
                .filter(|(_, mi)| mi.opcode == 13) // RET opcode
                .map(|(i, _)| i)
                .collect();

            // Insert epilogue in reverse order (so positions stay valid)
            for &ret_pos in ret_positions.iter().rev() {
                for instr in epilogue.iter().rev() {
                    bb.instructions.insert(ret_pos, instr.clone());
                }
            }

            // Remove the old RET
            let new_ret_positions: Vec<usize> = bb
                .instructions
                .iter()
                .enumerate()
                .filter(|(_, mi)| mi.opcode == 13)
                .map(|(i, _)| i)
                .collect();
            for &pos in new_ret_positions.iter().rev() {
                bb.instructions.remove(pos);
            }

            // Add fresh RET after each epilogue
            let mut ret_instr = MachineInstr::new(13); // RET
            bb.instructions.push(ret_instr);
        }
    }

    /// Build the epilogue instruction sequence.
    fn build_epilogue(
        &self,
        callee_saved: &CalleeSavedAnalysis,
        frame_layout: &FrameLayout,
        has_frame_pointer: bool,
    ) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        if has_frame_pointer {
            // 1. Restore callee-saved registers (reverse order)
            let mut saved_regs: Vec<u32> = callee_saved.must_save.iter().copied().collect();
            saved_regs.sort();
            saved_regs.reverse();
            for &reg in &saved_regs {
                let mut pop_reg = MachineInstr::new(12); // POP
                pop_reg.operands.push(MachineOperand::PhysReg(reg));
                instrs.push(pop_reg);
            }

            // 2. Deallocate stack space: mov rsp, rbp
            let mut mov_sp = MachineInstr::new(1); // MOV
            mov_sp.operands.push(MachineOperand::PhysReg(self.sp_reg));
            mov_sp.operands.push(MachineOperand::PhysReg(self.fp_reg));
            instrs.push(mov_sp);

            // 3. Pop frame pointer
            let mut pop_fp = MachineInstr::new(12); // POP
            pop_fp.operands.push(MachineOperand::PhysReg(self.fp_reg));
            instrs.push(pop_fp);
        } else {
            // No frame pointer: just restore saved regs and add to SP
            let mut saved_regs: Vec<u32> = callee_saved.must_save.iter().copied().collect();
            saved_regs.sort();
            saved_regs.reverse();
            for &reg in &saved_regs {
                let mut pop_reg = MachineInstr::new(12); // POP
                pop_reg.operands.push(MachineOperand::PhysReg(reg));
                instrs.push(pop_reg);
            }

            if frame_layout.frame_size > 0 {
                let mut add_sp = MachineInstr::new(2); // ADD
                add_sp.operands.push(MachineOperand::PhysReg(self.sp_reg));
                add_sp
                    .operands
                    .push(MachineOperand::Imm(frame_layout.frame_size));
                instrs.push(add_sp);
            }
        }

        let _ = callee_saved;
        let _ = frame_layout;
        instrs
    }
}

impl Default for EpilogueInserter {
    fn default() -> Self {
        EpilogueInserter::for_x86_64()
    }
}

// ============================================================================
// Frame Index Elimination — replace FrameIndex operands with real offsets
// ============================================================================

/// Eliminates FrameIndex machine operands by replacing them with
/// actual SP/FP-relative memory operands.
pub struct FrameIndexEliminator {
    /// Frame pointer register.
    pub fp_reg: u32,
    /// Stack pointer register.
    pub sp_reg: u32,
    /// Current frame layout.
    pub layout: FrameLayout,
}

impl FrameIndexEliminator {
    /// Create a new frame index eliminator.
    pub fn new(fp_reg: u32, sp_reg: u32, layout: FrameLayout) -> Self {
        FrameIndexEliminator {
            fp_reg,
            sp_reg,
            layout,
        }
    }

    /// Eliminate all frame indices in the function.
    ///
    /// Frame indices are placeholders like `FrameIndex(0)` that represent
    /// stack-relative addresses. After frame layout is computed, they
    /// are replaced with actual `[base_reg + offset]` operands.
    pub fn eliminate_frame_indices(&self, mf: &mut MachineFunction) {
        for bb in &mut mf.blocks {
            for mi in &mut bb.instructions {
                self.eliminate_in_instr(mi);
            }
        }
    }

    /// Eliminate frame indices in a single instruction.
    fn eliminate_in_instr(&self, mi: &mut MachineInstr) {
        // FrameIndex operands are represented as special labels or
        // negative immediate values. Replace them with real offsets.
        let mut extras: Vec<MachineOperand> = Vec::new();
        for op in &mut mi.operands {
            if let MachineOperand::Imm(fi) = *op {
                if fi < 0 && fi > -1000000 {
                    let offset = self.layout.callee_saved_offset + fi;
                    *op = MachineOperand::PhysReg(self.fp_reg);
                    extras.push(MachineOperand::Imm(offset));
                }
            }
        }
        for extra in extras {
            mi.operands.push(extra);
        }
    }
}

impl Default for FrameIndexEliminator {
    fn default() -> Self {
        FrameIndexEliminator {
            fp_reg: 5,
            sp_reg: 4,
            layout: FrameLayout::default(),
        }
    }
}

// ============================================================================
// Stack Realignment Prologue/Epilogue — for over-aligned locals
// ============================================================================

/// Stack realignment handler for functions with over-aligned locals.
pub struct StackRealignHandler {
    /// Required stack alignment.
    pub alignment: i64,
    /// Frame pointer register.
    pub fp_reg: u32,
    /// Stack pointer register.
    pub sp_reg: u32,
    /// Scratch register used for alignment computation.
    pub scratch_reg: u32,
}

impl StackRealignHandler {
    /// Create a new stack realignment handler.
    pub fn new(alignment: i64, fp_reg: u32, sp_reg: u32, scratch_reg: u32) -> Self {
        StackRealignHandler {
            alignment,
            fp_reg,
            sp_reg,
            scratch_reg,
        }
    }

    /// Emit realignment prologue.
    ///
    /// ```asm
    /// push rbp
    /// mov  rbp, rsp
    /// push scratch       ; save the scratch register
    /// mov  scratch, rsp  ; remember original RSP
    /// and  rsp, -align   ; align RSP down
    /// sub  rsp, frame_size
    /// ```
    pub fn emit_realign_prologue(&self, frame_size: i64) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        // push rbp
        let mut push_fp = MachineInstr::new(11);
        push_fp.operands.push(MachineOperand::PhysReg(self.fp_reg));
        instrs.push(push_fp);

        // mov rbp, rsp
        let mut mov_fp = MachineInstr::new(1);
        mov_fp.operands.push(MachineOperand::PhysReg(self.fp_reg));
        mov_fp.operands.push(MachineOperand::PhysReg(self.sp_reg));
        instrs.push(mov_fp);

        // push scratch_reg
        let mut push_scratch = MachineInstr::new(11);
        push_scratch
            .operands
            .push(MachineOperand::PhysReg(self.scratch_reg));
        instrs.push(push_scratch);

        // mov scratch_reg, rsp
        let mut mov_scratch = MachineInstr::new(1);
        mov_scratch
            .operands
            .push(MachineOperand::PhysReg(self.scratch_reg));
        mov_scratch
            .operands
            .push(MachineOperand::PhysReg(self.sp_reg));
        instrs.push(mov_scratch);

        // and rsp, -alignment
        let mut and_align = MachineInstr::new(6); // AND
        and_align
            .operands
            .push(MachineOperand::PhysReg(self.sp_reg));
        and_align
            .operands
            .push(MachineOperand::Imm(-self.alignment));
        instrs.push(and_align);

        // sub rsp, frame_size
        if frame_size > 0 {
            let mut sub_frame = MachineInstr::new(3); // SUB
            sub_frame
                .operands
                .push(MachineOperand::PhysReg(self.sp_reg));
            sub_frame.operands.push(MachineOperand::Imm(frame_size));
            instrs.push(sub_frame);
        }

        instrs
    }

    /// Emit realignment epilogue.
    ///
    /// ```asm
    /// lea  rsp, [rbp - scratch_offset]
    /// pop  scratch
    /// pop  rbp
    /// ret
    /// ```
    pub fn emit_realign_epilogue(&self, scratch_offset: i64) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        // lea rsp, [rbp + scratch_offset]
        let mut lea_instr = MachineInstr::new(19); // LEA
        lea_instr
            .operands
            .push(MachineOperand::PhysReg(self.sp_reg));
        lea_instr
            .operands
            .push(MachineOperand::PhysReg(self.fp_reg));
        lea_instr.operands.push(MachineOperand::Imm(scratch_offset));
        instrs.push(lea_instr);

        // pop scratch
        let mut pop_scratch = MachineInstr::new(12); // POP
        pop_scratch
            .operands
            .push(MachineOperand::PhysReg(self.scratch_reg));
        instrs.push(pop_scratch);

        // pop rbp
        let mut pop_fp = MachineInstr::new(12); // POP
        pop_fp.operands.push(MachineOperand::PhysReg(self.fp_reg));
        instrs.push(pop_fp);

        // ret
        let mut ret_instr = MachineInstr::new(13); // RET
        instrs.push(ret_instr);

        instrs
    }
}

impl Default for StackRealignHandler {
    fn default() -> Self {
        StackRealignHandler {
            alignment: 32,
            fp_reg: 5,
            sp_reg: 4,
            scratch_reg: 0, // RAX
        }
    }
}

// ============================================================================
// Packed Stack Layout — minimize stack size by packing small variables
// ============================================================================

/// Packed stack layout optimizer.
///
/// Packs small local variables and spill slots into the same aligned
/// slots to minimize the total stack frame size. Variables smaller than
/// the alignment are co-located when safe.
pub struct PackedStackLayout {
    /// Maximum stack alignment.
    pub alignment: i64,
    /// Slot size (register width).
    pub slot_size: i64,
}

impl PackedStackLayout {
    /// Create a new packed layout optimizer.
    pub fn new(alignment: i64, slot_size: i64) -> Self {
        PackedStackLayout {
            alignment,
            slot_size,
        }
    }

    /// Pack frame objects to minimize stack size.
    ///
    /// Returns a mapping from object ID to its frame offset.
    pub fn pack(&self, objects: &[FrameObject]) -> HashMap<u32, i64> {
        let mut offsets = HashMap::new();
        let mut current_offset: i64 = 0;

        // Sort objects by alignment descending (largest first)
        let mut sorted: Vec<&FrameObject> = objects.iter().collect();
        sorted.sort_by(|a, b| b.alignment.cmp(&a.alignment));

        for obj in sorted {
            // Align the current offset
            let aligned = align_to(-current_offset, obj.alignment as i64);
            current_offset = -aligned;

            offsets.insert(obj.id, current_offset);
            current_offset -= obj.size as i64;
        }

        offsets
    }

    /// Compute the minimum stack size for the given objects.
    pub fn min_stack_size(&self, objects: &[FrameObject]) -> i64 {
        let offsets = self.pack(objects);
        if offsets.is_empty() {
            return 0;
        }

        let min_offset = offsets.values().min().copied().unwrap_or(0);
        let max_size = objects.iter().map(|o| o.size as i64).max().unwrap_or(0);

        align_to((-min_offset) + max_size, self.alignment)
    }
}

/// A frame object to be placed on the stack.
#[derive(Debug, Clone)]
pub struct FrameObject {
    /// Unique identifier for this object.
    pub id: u32,
    /// Size in bytes.
    pub size: u32,
    /// Required alignment.
    pub alignment: u32,
    /// Type of frame object.
    pub obj_type: FrameObjectType,
    /// Whether this object is fixed (cannot be moved).
    pub is_fixed: bool,
}

// ============================================================================
// Prologue/Epilogue Pass — unified pass for the codegen pipeline
// ============================================================================

/// Main prologue/epilogue insertion pass.
pub struct PrologueEpiloguePass {
    /// The prologue inserter.
    pub prologue: PrologueInserter,
    /// The epilogue inserter.
    pub epilogue: EpilogueInserter,
    /// Frame index eliminator.
    pub fi_eliminator: Option<FrameIndexEliminator>,
}

impl PrologueEpiloguePass {
    /// Create a new prologue/epilogue pass for x86-64.
    pub fn for_x86_64() -> Self {
        PrologueEpiloguePass {
            prologue: PrologueInserter::for_x86_64(),
            epilogue: EpilogueInserter::for_x86_64(),
            fi_eliminator: None,
        }
    }

    /// Run the full prologue/epilogue insertion on a machine function.
    pub fn run(&mut self, mf: &mut MachineFunction, target: &TargetRegInfo) {
        // 1. Analyze callee-saved registers
        let callee_saved = CalleeSavedAnalysis::analyze(mf, target, 8);

        // 2. Compute frame layout
        let frame_layout = FrameLayout::compute(callee_saved.total_save_size as i64, 0, 0, 0, 16);

        // 3. Insert prologue
        self.prologue
            .insert_prologue(mf, &callee_saved, &frame_layout);

        // 4. Insert epilogue
        self.epilogue.insert_epilogue(
            mf,
            &callee_saved,
            &frame_layout,
            callee_saved.needs_frame_pointer,
        );

        // 5. Eliminate frame indices
        self.fi_eliminator = Some(FrameIndexEliminator::new(
            self.prologue.fp_reg,
            self.prologue.sp_reg,
            frame_layout,
        ));

        if let Some(ref eliminator) = self.fi_eliminator {
            eliminator.eliminate_frame_indices(mf);
        }
    }
}

impl Default for PrologueEpiloguePass {
    fn default() -> Self {
        PrologueEpiloguePass::for_x86_64()
    }
}

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

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

    fn make_test_mf() -> MachineFunction {
        let mut mf = MachineFunction::new("test_func");
        let mut bb = MachineBasicBlock {
            name: "entry".to_string(),
            instructions: Vec::new(),
            successors: Vec::new(),
        };
        // Add a simple instruction that uses a callee-saved register
        let mut mov = MachineInstr::new(1); // MOV
        mov.operands.push(MachineOperand::PhysReg(3)); // RBX
        mov.operands.push(MachineOperand::Imm(42));
        bb.instructions.push(mov);

        // Add a return
        let ret = MachineInstr::new(13); // RET
        bb.instructions.push(ret);

        mf.blocks.push(bb);
        mf
    }

    #[test]
    fn test_callee_saved_analysis() {
        let mf = make_test_mf();
        let target = TargetRegInfo::for_x86_64();
        let analysis = CalleeSavedAnalysis::analyze(&mf, &target, 8);
        // RBX is callee-saved and used in the test function
        assert!(analysis.must_save.contains(&3));
    }

    #[test]
    fn test_frame_layout_compute() {
        let layout = FrameLayout::compute(8, 16, 32, 0, 16);
        assert!(layout.frame_size > 0);
        assert!(layout.frame_size % 16 == 0); // must be aligned
    }

    #[test]
    fn test_scratch_reg_analysis() {
        let mf = make_test_mf();
        let target = TargetRegInfo::for_x86_64();
        let mut used = HashSet::new();
        used.insert(3); // RBX used
        let analysis = ScratchRegAnalysis::find(&mf, &target, &used);
        // Should find some scratch registers
        assert!(!analysis.available.is_empty());
    }

    #[test]
    fn test_prologue_inserter() {
        let mut mf = make_test_mf();
        let target = TargetRegInfo::for_x86_64();
        let callee_saved = CalleeSavedAnalysis::analyze(&mf, &target, 8);
        let layout = FrameLayout::compute(8, 0, 0, 0, 16);
        let inserter = PrologueInserter::for_x86_64();
        inserter.insert_prologue(&mut mf, &callee_saved, &layout);
        // Should have inserted prologue instructions before the original MOV
        assert!(mf.blocks[0].instructions.len() > 2);
    }

    #[test]
    fn test_epilogue_inserter() {
        let mut mf = make_test_mf();
        let target = TargetRegInfo::for_x86_64();
        let callee_saved = CalleeSavedAnalysis::analyze(&mf, &target, 8);
        let layout = FrameLayout::compute(8, 0, 0, 0, 16);
        let inserter = EpilogueInserter::for_x86_64();
        inserter.insert_epilogue(&mut mf, &callee_saved, &layout, true);
        // Epilogue should have been inserted, then a RET appended
        let last_instr = mf.blocks[0].instructions.last().unwrap();
        assert_eq!(last_instr.opcode, 13); // RET
    }

    #[test]
    fn test_frame_index_eliminator() {
        let layout = FrameLayout::compute(8, 16, 32, 0, 16);
        let eliminator = FrameIndexEliminator::new(5, 4, layout);
        let mut mf = make_test_mf();
        eliminator.eliminate_frame_indices(&mut mf);
        // Should not crash
    }

    #[test]
    fn test_stack_realign_handler() {
        let handler = StackRealignHandler::new(32, 5, 4, 0);
        let prologue = handler.emit_realign_prologue(64);
        assert!(prologue.len() >= 5);
        let epilogue = handler.emit_realign_epilogue(-8);
        assert!(epilogue.len() >= 3);
    }

    #[test]
    fn test_prologue_epilogue_pass() {
        let mut mf = make_test_mf();
        let target = TargetRegInfo::for_x86_64();
        let mut pass = PrologueEpiloguePass::for_x86_64();
        pass.run(&mut mf, &target);
        // After pass, function should have prologue and epilogue
        assert!(mf.blocks[0].instructions.len() > 4);
    }

    #[test]
    fn test_packed_stack_layout() {
        let packer = PackedStackLayout::new(16, 8);
        let objects = vec![
            FrameObject {
                id: 0,
                size: 4,
                alignment: 4,
                obj_type: FrameObjectType::Local,
                is_fixed: false,
            },
            FrameObject {
                id: 1,
                size: 8,
                alignment: 8,
                obj_type: FrameObjectType::Local,
                is_fixed: false,
            },
        ];
        let offsets = packer.pack(&objects);
        assert_eq!(offsets.len(), 2);
    }

    #[test]
    fn test_frame_object() {
        let obj = FrameObject {
            id: 42,
            size: 8,
            alignment: 8,
            obj_type: FrameObjectType::CalleeSaved,
            is_fixed: true,
        };
        assert_eq!(obj.id, 42);
        assert_eq!(obj.size, 8);
    }
}