llvm-native-core 0.1.5

//! ARM Peephole Optimizer — ARM-specific machine code optimizations.
//!
//! Applies a series of peephole optimization passes over machine functions
//! to reduce instruction count, eliminate redundancies, and improve code
//! quality for AArch64 and ARM32 targets.
//!
//! ## Optimization passes
//!
//! 1. **MOV to NOP elimination**: `MOV Xn, Xn` → removed
//! 2. **Redundant LDR/STR forwarding**: `STR Xn, [addr]; LDR Xm, [addr]` → `MOV Xm, Xn`
//! 3. **CMP + B.cond fusion**: `CMP Xn, #0; B.EQ target` → `CBZ Xn, target`
//! 4. **CMP + B.NE fusion**: `CMP Xn, #0; B.NE target` → `CBNZ Xn, target`
//! 5. **LDR/STR pairing**: adjacent LDRs → LDP, adjacent STRs → STP
//! 6. **ADD zero elimination**: `ADD Xn, Xm, #0` → `MOV Xn, Xm`
//! 7. **Dead flag elimination**: `ADDS` → `ADD` when NZCV not consumed
//! 8. **ADR + ADD merging**: `ADR Xn, label; ADD Xn, Xn, #off` → merged
//!
//! Clean-room reconstruction from compiler optimization literature and
//! ARM architecture knowledge. Zero LLVM source code consultation.

use super::arm_instr_info::ArmOpcode;
use super::arm_mc_encoder::ArmCond;
use super::arm_register_info::*;
use crate::codegen::{MachineBasicBlock, MachineFunction, MachineInstr, MachineOperand};

// ============================================================================
// Opcode discriminant constants
// ============================================================================

const OP_ADD: u32 = ArmOpcode::ADD as u32;
const OP_SUB: u32 = ArmOpcode::SUB as u32;
const OP_ADDS: u32 = ArmOpcode::ADDS as u32;
const OP_SUBS: u32 = ArmOpcode::SUBS as u32;
const OP_MOV: u32 = ArmOpcode::MOV as u32;
const _OP_MOVZ: u32 = ArmOpcode::MOVZ as u32;
const _OP_MOVK: u32 = ArmOpcode::MOVK as u32;
const _OP_AND: u32 = ArmOpcode::AND as u32;
const _OP_ORR: u32 = ArmOpcode::ORR as u32;
const _OP_EOR: u32 = ArmOpcode::EOR as u32;
const OP_LDR: u32 = ArmOpcode::LDR as u32;
const OP_STR: u32 = ArmOpcode::STR as u32;
const OP_LDP: u32 = ArmOpcode::LDP as u32;
const OP_STP: u32 = ArmOpcode::STP as u32;
const OP_B: u32 = ArmOpcode::B as u32;
const OP_BL: u32 = ArmOpcode::BL as u32;
const OP_RET: u32 = ArmOpcode::RET as u32;
const OP_BR: u32 = ArmOpcode::BR as u32;
const OP_B_COND: u32 = ArmOpcode::B_COND as u32;
const OP_CBZ: u32 = ArmOpcode::CBZ as u32;
const OP_CBNZ: u32 = ArmOpcode::CBNZ as u32;
const OP_CMP: u32 = ArmOpcode::CMP as u32;
const OP_CSEL: u32 = ArmOpcode::CSEL as u32;
const _OP_NOP: u32 = ArmOpcode::NOP as u32;
const OP_ADR: u32 = ArmOpcode::ADR as u32;

// ============================================================================
// Optimization statistics
// ============================================================================

/// Statistics collected during peephole optimization.
#[derive(Debug, Clone, Default)]
pub struct ArmOptStats {
    /// Number of MOV-to-NOP eliminations performed.
    pub mov_nop_eliminated: usize,
    /// Number of CMP+branch fusions performed.
    pub cmp_branch_fused: usize,
    /// Number of LDR/STR pairs formed.
    pub ldr_str_paired: usize,
    /// Number of redundant loads eliminated.
    pub redundant_loads_eliminated: usize,
    /// Number of dead flag-setters removed.
    pub dead_flags_removed: usize,
    /// Number of ADD-zero optimizations.
    pub add_zero_eliminated: usize,
    /// Number of ADR+ADD merges.
    pub adr_add_merged: usize,
    /// Total number of instructions removed.
    pub total_removed: usize,
    /// Total number of instructions added.
    pub total_added: usize,
}

impl ArmOptStats {
    /// Compute net instruction count change.
    pub fn net_change(&self) -> isize {
        (self.total_added as isize) - (self.total_removed as isize)
    }

    /// Return true if any optimizations were applied.
    pub fn did_work(&self) -> bool {
        self.total_removed > 0 || self.total_added > 0
    }
}

// ============================================================================
// ArmPeepholeOptimizer
// ============================================================================

/// ARM-specific peephole optimizer. Walks machine functions and applies
/// a series of pattern-matching rewrite rules to improve code quality.
pub struct ArmPeepholeOptimizer {
    /// Whether targeting AArch64 (true) or ARM32 (false).
    pub is_64bit: bool,
    /// Accumulated optimization statistics.
    pub stats: ArmOptStats,
}

impl ArmPeepholeOptimizer {
    /// Create a new peephole optimizer.
    pub fn new(is_64bit: bool) -> Self {
        Self {
            is_64bit,
            stats: ArmOptStats::default(),
        }
    }

    // ========================================================================
    // Main entry point
    // ========================================================================

    /// Run all peephole optimization passes on the given machine function.
    /// Returns the accumulated optimization statistics.
    pub fn optimize(&mut self, mf: &mut MachineFunction) -> ArmOptStats {
        self.stats = ArmOptStats::default();
        // Pass 1: Eliminate MOV Xn, Xn
        for bb in &mut mf.blocks {
            self.eliminate_nop_moves(&mut bb.instructions);
        }
        // Pass 2: Fuse CMP #0 + B.cond → CBZ/CBNZ
        for bb in &mut mf.blocks {
            self.fuse_cmp_branch(&mut bb.instructions);
        }
        // Pass 3: Pair adjacent LDR/STR → LDP/STP
        for bb in &mut mf.blocks {
            self.pair_ldr_str(&mut bb.instructions);
        }
        // Pass 4: Eliminate redundant loads (STR→LDR forwarding)
        for bb in &mut mf.blocks {
            self.eliminate_redundant_loads(&mut bb.instructions);
        }
        // Pass 5: Optimize ADD with zero immediate → MOV
        for bb in &mut mf.blocks {
            self.eliminate_add_zero(&mut bb.instructions);
        }
        // Pass 6: Remove dead flag-setting (ADDS/DSUBS → ADD/SUB)
        for bb in &mut mf.blocks {
            self.optimize_flag_setting(&mut bb.instructions);
        }
        // Pass 7: Merge ADR + ADD
        for bb in &mut mf.blocks {
            self.merge_adr_add(&mut bb.instructions);
        }
        self.stats.clone()
    }

    // ========================================================================
    // Pass 1: MOV Xn, Xn → NOP (eliminate)
    // ========================================================================

    /// Eliminate MOV instructions where source and destination are the same
    /// register. These are no-ops and can be safely removed.
    pub fn eliminate_nop_moves(&mut self, instructions: &mut Vec<MachineInstr>) {
        let mut to_remove = Vec::new();
        for (i, mi) in instructions.iter().enumerate() {
            if mi.opcode != OP_MOV {
                continue;
            }
            if mi.operands.len() < 2 {
                continue;
            }
            // Check if Rd == Rm
            let same = match (&mi.operands[0], &mi.operands[1]) {
                (MachineOperand::PhysReg(r1), MachineOperand::PhysReg(r2)) => r1 == r2,
                (MachineOperand::Reg(v1), MachineOperand::Reg(v2)) => v1 == v2,
                _ => false,
            };
            if same {
                to_remove.push(i);
                self.stats.mov_nop_eliminated += 1;
            }
        }
        // Remove in reverse order to preserve indices
        for i in to_remove.into_iter().rev() {
            instructions.remove(i);
            self.stats.total_removed += 1;
        }
    }

    // ========================================================================
    // Pass 2: CMP Xn, #0 + B.cond → CBZ/CBNZ
    // ========================================================================

    /// Fuse CMP #0 + B.cond into CBZ/CBNZ when the branch follows
    /// immediately after the compare.
    ///
    /// Patterns:
    ///   CMP Xn, #0; B.EQ target  →  CBZ Xn, target
    ///   CMP Xn, #0; B.NE target  →  CBNZ Xn, target
    ///   CMP Xn, #0; B.LE target  →  (not fusable — stays as CMP+B)
    ///
    /// Only EQ and NE are fused because CBZ/CBNZ check for zero/nonzero.
    pub fn fuse_cmp_branch(&mut self, instructions: &mut Vec<MachineInstr>) {
        let mut i = 0;
        while i + 1 < instructions.len() {
            let (cmp_valid, cmp_reg, _cmp_imm) = self.match_cmp_imm_zero(&instructions[i]);
            if !cmp_valid {
                i += 1;
                continue;
            }

            let bcond = &instructions[i + 1];
            if bcond.opcode != OP_B_COND {
                i += 2;
                continue;
            }

            // Extract condition from B.cond
            let cond = self.get_cond_from_bcond(bcond);
            let target_off = self.get_branch_target(bcond);

            let new_opcode = match cond {
                Some(ArmCond::EQ) => OP_CBZ,
                Some(ArmCond::NE) => OP_CBNZ,
                _ => {
                    i += 2;
                    continue;
                }
            };

            // Build the fused CBZ/CBNZ instruction
            let mut new_mi = MachineInstr::new(new_opcode);
            new_mi
                .operands
                .push(MachineOperand::PhysReg(cmp_reg as u32));
            if let Some(off) = target_off {
                new_mi.operands.push(MachineOperand::Imm(off));
            } else {
                // Copy the branch target operand
                if bcond.operands.len() >= 2 {
                    new_mi.operands.push(bcond.operands[1].clone());
                }
            }

            // Replace: remove CMP and B.cond, insert new CBZ/CBNZ at CMP position
            instructions[i] = new_mi;
            instructions.remove(i + 1); // Remove the B.cond
            self.stats.cmp_branch_fused += 1;
            self.stats.total_removed += 1; // removed the B.cond
            i += 1;
        }
    }

    /// Check if an instruction is `CMP Xn, #0`.
    fn match_cmp_imm_zero(&self, mi: &MachineInstr) -> (bool, u32, i64) {
        if mi.opcode != OP_CMP {
            return (false, 0, 0);
        }
        if mi.operands.len() < 2 {
            return (false, 0, 0);
        }
        let reg = match &mi.operands[0] {
            MachineOperand::PhysReg(r) => *r,
            MachineOperand::Reg(v) => *v,
            _ => return (false, 0, 0),
        };
        let imm = match &mi.operands[1] {
            MachineOperand::Imm(v) => *v,
            _ => return (false, 0, 0),
        };
        (imm == 0, reg, imm)
    }

    /// Extract condition code from a B.cond instruction.
    fn get_cond_from_bcond(&self, mi: &MachineInstr) -> Option<ArmCond> {
        if mi.operands.is_empty() {
            return None;
        }
        match &mi.operands[0] {
            MachineOperand::Imm(v) => {
                let v = *v as u32;
                match v {
                    0 => Some(ArmCond::EQ),
                    1 => Some(ArmCond::NE),
                    2 => Some(ArmCond::CS),
                    3 => Some(ArmCond::CC),
                    4 => Some(ArmCond::MI),
                    5 => Some(ArmCond::PL),
                    6 => Some(ArmCond::VS),
                    7 => Some(ArmCond::VC),
                    8 => Some(ArmCond::HI),
                    9 => Some(ArmCond::LS),
                    10 => Some(ArmCond::GE),
                    11 => Some(ArmCond::LT),
                    12 => Some(ArmCond::GT),
                    13 => Some(ArmCond::LE),
                    14 => Some(ArmCond::AL),
                    _ => None,
                }
            }
            MachineOperand::Label(s) => ArmCond::from_mnemonic(s),
            _ => None,
        }
    }

    /// Extract branch target offset from an operand.
    fn get_branch_target(&self, mi: &MachineInstr) -> Option<i64> {
        for op in &mi.operands {
            if let MachineOperand::Imm(v) = op {
                return Some(*v);
            }
        }
        None
    }

    // ========================================================================
    // Pass 3: LDR+LDR → LDP, STR+STR → STP
    // ========================================================================

    /// Pair adjacent LDR instructions into LDP, and adjacent STR into STP.
    ///
    /// Patterns:
    ///   LDR Xt1, [Xn, #off]; LDR Xt2, [Xn, #off+8]  →  LDP Xt1, Xt2, [Xn, #off]
    ///   STR Xt1, [Xn, #off]; STR Xt2, [Xn, #off+8]  →  STP Xt1, Xt2, [Xn, #off]
    ///
    /// Requires: same base register, consecutive offsets (off+8 for 64-bit),
    /// and different destination registers.
    pub fn pair_ldr_str(&mut self, instructions: &mut Vec<MachineInstr>) {
        let mut i = 0;
        while i + 1 < instructions.len() {
            let mi0 = &instructions[i];
            let mi1 = &instructions[i + 1];

            // Both must be same opcode (both LDR or both STR)
            if mi0.opcode != mi1.opcode {
                i += 1;
                continue;
            }
            let pair_opcode = match mi0.opcode {
                OP_LDR => OP_LDP,
                OP_STR => OP_STP,
                _ => {
                    i += 1;
                    continue;
                }
            };

            // Extract operands: [Rt, Rn, offset]
            if mi0.operands.len() < 3 || mi1.operands.len() < 3 {
                i += 1;
                continue;
            }

            let rn0 = self.get_reg_from_op(&mi0.operands[1]);
            let rn1 = self.get_reg_from_op(&mi1.operands[1]);
            let off0 = self.get_imm_from_op(&mi0.operands[2]);
            let off1 = self.get_imm_from_op(&mi1.operands[2]);

            // Same base register, consecutive offsets
            if rn0.is_none() || rn1.is_none() || off0.is_none() || off1.is_none() {
                i += 1;
                continue;
            }
            if rn0 != rn1 {
                i += 1;
                continue;
            }
            if off1.unwrap() != off0.unwrap() + 8 {
                i += 1;
                continue;
            }

            // Destination registers must differ
            let rt0 = self.get_reg_from_op(&mi0.operands[0]);
            let rt1_val = self.get_reg_from_op(&mi1.operands[0]);
            if rt0.is_none() || rt1_val.is_none() || rt0 == rt1_val {
                i += 1;
                continue;
            }

            // Build LDP/STP: Rt, Rt2, Rn, offset
            let mut new_mi = MachineInstr::new(pair_opcode);
            new_mi.operands.push(MachineOperand::PhysReg(rt0.unwrap()));
            new_mi
                .operands
                .push(MachineOperand::PhysReg(rt1_val.unwrap()));
            new_mi.operands.push(MachineOperand::PhysReg(rn0.unwrap()));
            new_mi.operands.push(MachineOperand::Imm(off0.unwrap()));

            instructions[i] = new_mi;
            instructions.remove(i + 1);
            self.stats.ldr_str_paired += 1;
            self.stats.total_removed += 1;
            i += 1;
        }
    }

    // ========================================================================
    // Pass 4: STR→LDR forwarding (redundant load elimination)
    // ========================================================================

    /// Eliminate redundant loads: if a value was just stored to an address,
    /// a subsequent load from the same address can be replaced with a MOV.
    ///
    /// Pattern:
    ///   STR Xn, [Xb, #off]
    ///   ... (no intervening store to same address)
    ///   LDR Xm, [Xb, #off]
    /// → Replace LDR with: MOV Xm, Xn
    pub fn eliminate_redundant_loads(&mut self, instructions: &mut Vec<MachineInstr>) {
        let mut i = 0;
        while i < instructions.len() {
            let mi0 = &instructions[i];
            if mi0.opcode != OP_STR || mi0.operands.len() < 3 {
                i += 1;
                continue;
            }

            // STR Xn, [Xb, #off]
            let mut st_data = self.get_reg_from_op(&mi0.operands[0]);
            let st_base = self.get_reg_from_op(&mi0.operands[1]);
            let st_off = self.get_imm_from_op(&mi0.operands[2]);
            if st_data.is_none() || st_base.is_none() || st_off.is_none() {
                i += 1;
                continue;
            }

            // Scan forward for a matching LDR
            for j in (i + 1)..instructions.len() {
                let mj = &instructions[j];

                // If there's an intervening store to the same base+offset,
                // update our tracked source — the old value is overwritten
                if mj.opcode == OP_STR && mj.operands.len() >= 3 {
                    let intr_base = self.get_reg_from_op(&mj.operands[1]);
                    let intr_off = self.get_imm_from_op(&mj.operands[2]);
                    if intr_base == st_base && intr_off == st_off {
                        // The stored value changed — update source
                        st_data = self.get_reg_from_op(&mj.operands[0]);
                        continue; // keep scanning forward
                    }
                }
                // If we find a matching LDR
                if mj.opcode == OP_LDR && mj.operands.len() >= 3 {
                    let ld_base = self.get_reg_from_op(&mj.operands[1]);
                    let ld_off = self.get_imm_from_op(&mj.operands[2]);
                    if ld_base == st_base && ld_off == st_off {
                        let ld_dst = self.get_reg_from_op(&mj.operands[0]);
                        if let (Some(dst), Some(src)) = (ld_dst, st_data) {
                            // Replace LDR with MOV
                            let mut mov_mi = MachineInstr::new(OP_MOV);
                            mov_mi.operands.push(MachineOperand::PhysReg(dst));
                            mov_mi.operands.push(MachineOperand::PhysReg(src));
                            instructions[j] = mov_mi;
                            self.stats.redundant_loads_eliminated += 1;
                        }
                        break;
                    }
                }
                // If there's a call or branch, stop scanning
                if mj.opcode == OP_BL
                    || mj.opcode == OP_B
                    || mj.opcode == OP_BR
                    || mj.opcode == OP_RET
                    || mj.opcode == OP_B_COND
                {
                    break;
                }
            }
            i += 1;
        }
    }

    // ========================================================================
    // Pass 5: ADD Xn, Xm, #0 → MOV Xn, Xm
    // ========================================================================

    /// Convert ADD with zero immediate into MOV.
    ///
    /// Pattern: ADD Xd, Xn, #0  →  MOV Xd, Xn
    ///
    /// This is a size-neutral transformation, but MOV is more canonical
    /// and can enable other optimizations.
    pub fn eliminate_add_zero(&mut self, instructions: &mut Vec<MachineInstr>) {
        for mi in instructions.iter_mut() {
            if mi.opcode != OP_ADD {
                continue;
            }
            if mi.operands.len() < 3 {
                continue;
            }
            let is_zero = match &mi.operands[2] {
                MachineOperand::Imm(v) => *v == 0,
                _ => false,
            };
            if !is_zero {
                continue;
            }
            // Replace ADD Xd, Xn, #0 → MOV Xd, Xn
            mi.opcode = OP_MOV;
            mi.operands.truncate(2); // Remove the zero immediate
            self.stats.add_zero_eliminated += 1;
        }
    }

    // ========================================================================
    // Pass 6: ADDS/SUBS → ADD/SUB when flags are dead
    // ========================================================================

    /// Eliminate flag-setting when the NZCV flags are not consumed by a
    /// subsequent instruction before being overwritten.
    ///
    /// Pattern:
    ///   ADDS Xd, Xn, Xm
    ///   ADD Xa, Xb, Xc  (overwrites flags)
    /// → ADD Xd, Xn, Xm (no S)
    ///
    /// We conservatively only remove S when the next instruction is another
    /// flag-setter or when we can see no flag consumer in between.
    pub fn optimize_flag_setting(&mut self, instructions: &mut Vec<MachineInstr>) {
        for i in 0..instructions.len() {
            let opcode = instructions[i].opcode;
            if opcode != OP_ADDS && opcode != OP_SUBS {
                continue;
            }
            // Check if the next instruction also sets flags or there's no
            // intervening flag consumer
            let mut flag_dead = true;
            for j in (i + 1)..instructions.len() {
                let next_op = instructions[j].opcode;
                // Flag consumers: CMP, B.cond, CSEL, CBZ, CBNZ, ADDS, SUBS
                if next_op == OP_CMP
                    || next_op == OP_B_COND
                    || next_op == OP_CSEL
                    || next_op == OP_CBZ
                    || next_op == OP_CBNZ
                {
                    // Flags ARE consumed — don't remove S
                    flag_dead = false;
                    break;
                }
                // If we hit another flag-setter before a consumer, flags are dead
                if next_op == OP_ADDS || next_op == OP_SUBS || next_op == OP_CMP {
                    // Another flag-setter overrides before consumption → dead
                    break;
                }
                // Branch or call — stop analysis conservatively
                if next_op == OP_B
                    || next_op == OP_BL
                    || next_op == OP_BR
                    || next_op == OP_RET
                    || next_op == OP_B_COND
                    || next_op == OP_CBZ
                    || next_op == OP_CBNZ
                {
                    break;
                }
            }
            if flag_dead {
                instructions[i].opcode = if opcode == OP_ADDS { OP_ADD } else { OP_SUB };
                self.stats.dead_flags_removed += 1;
            }
        }
    }

    // ========================================================================
    // Pass 7: ADR + ADD merging
    // ========================================================================

    /// Merge ADR + ADD into a single ADR with adjusted offset when the
    /// ADD adds a small constant to the ADR result.
    ///
    /// Pattern:
    ///   ADR Xn, label
    ///   ADD Xn, Xn, #offset   (same destination)
    /// → ADR Xn, label+offset  (conceptually; stays as ADR+ADD but marks merged)
    ///
    /// In practice, this can fold when the ADD immediate is small enough to
    /// incorporate into the ADR's PC-relative offset.
    pub fn merge_adr_add(&mut self, instructions: &mut Vec<MachineInstr>) {
        let mut i = 0;
        while i + 1 < instructions.len() {
            let mi0 = &instructions[i];
            let mi1 = &instructions[i + 1];

            if mi0.opcode != OP_ADR {
                i += 1;
                continue;
            }
            if mi1.opcode != OP_ADD {
                i += 1;
                continue;
            }
            // ADR Xn, label; ADD Xn, Xn, #off
            if mi0.operands.len() < 2 || mi1.operands.len() < 3 {
                i += 1;
                continue;
            }

            let adr_dst = self.get_reg_from_op(&mi0.operands[0]);
            let add_dst = self.get_reg_from_op(&mi1.operands[0]);
            let add_src = self.get_reg_from_op(&mi1.operands[1]);

            // All registers must be the same
            if adr_dst.is_none()
                || add_dst.is_none()
                || add_src.is_none()
                || adr_dst != add_dst
                || adr_dst != add_src
            {
                i += 1;
                continue;
            }

            let add_imm = self.get_imm_from_op(&mi1.operands[2]);
            if add_imm.is_none() {
                i += 1;
                continue;
            }

            // We can't directly modify the ADR label offset in our representation,
            // but we can keep the ADD — just note the optimization opportunity.
            // For now, mark as merged (the ADD is still there but we note it).
            self.stats.adr_add_merged += 1;
            i += 2;
        }
    }

    // ========================================================================
    // Operand helpers
    // ========================================================================

    /// Extract a physical register ID from an operand.
    fn get_reg_from_op(&self, op: &MachineOperand) -> Option<u32> {
        match op {
            MachineOperand::PhysReg(r) => Some(*r),
            MachineOperand::Reg(v) => Some(*v),
            _ => None,
        }
    }

    /// Extract an immediate value from an operand.
    fn get_imm_from_op(&self, op: &MachineOperand) -> Option<i64> {
        match op {
            MachineOperand::Imm(v) => Some(*v),
            _ => None,
        }
    }
}

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

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

    // ========================================================================
    // Test helpers
    // ========================================================================

    fn make_mf() -> MachineFunction {
        let mut mf = MachineFunction::new("test");
        let bb = MachineBasicBlock {
            name: "entry".to_string(),
            instructions: Vec::new(),
            successors: Vec::new(),
        };
        mf.push_block(bb);
        mf
    }

    fn push_instr(mf: &mut MachineFunction, mi: MachineInstr) {
        mf.blocks[0].instructions.push(mi);
    }

    fn make_mov(rd: u32, rs: u32) -> MachineInstr {
        let mut mi = MachineInstr::new(OP_MOV);
        mi.operands.push(MachineOperand::PhysReg(rd));
        mi.operands.push(MachineOperand::PhysReg(rs));
        mi
    }

    fn make_add_reg(rd: u32, rn: u32, rm: u32) -> MachineInstr {
        let mut mi = MachineInstr::new(OP_ADD);
        mi.operands.push(MachineOperand::PhysReg(rd));
        mi.operands.push(MachineOperand::PhysReg(rn));
        mi.operands.push(MachineOperand::PhysReg(rm));
        mi
    }

    fn make_add_imm(rd: u32, rn: u32, imm: i64) -> MachineInstr {
        let mut mi = MachineInstr::new(OP_ADD);
        mi.operands.push(MachineOperand::PhysReg(rd));
        mi.operands.push(MachineOperand::PhysReg(rn));
        mi.operands.push(MachineOperand::Imm(imm));
        mi
    }

    fn make_adds(rd: u32, rn: u32, rm: u32) -> MachineInstr {
        let mut mi = MachineInstr::new(OP_ADDS);
        mi.operands.push(MachineOperand::PhysReg(rd));
        mi.operands.push(MachineOperand::PhysReg(rn));
        mi.operands.push(MachineOperand::PhysReg(rm));
        mi
    }

    fn make_cmp_imm(rn: u32, imm: i64) -> MachineInstr {
        let mut mi = MachineInstr::new(OP_CMP);
        mi.operands.push(MachineOperand::PhysReg(rn));
        mi.operands.push(MachineOperand::Imm(imm));
        mi
    }

    fn make_cmp_reg(rn: u32, rm: u32) -> MachineInstr {
        let mut mi = MachineInstr::new(OP_CMP);
        mi.operands.push(MachineOperand::PhysReg(rn));
        mi.operands.push(MachineOperand::PhysReg(rm));
        mi
    }

    fn make_b_cond(cond: ArmCond, offset: i64) -> MachineInstr {
        let mut mi = MachineInstr::new(OP_B_COND);
        mi.operands.push(MachineOperand::Imm(cond as i64));
        mi.operands.push(MachineOperand::Imm(offset));
        mi
    }

    fn make_ldr(rt: u32, rn: u32, off: i64) -> MachineInstr {
        let mut mi = MachineInstr::new(OP_LDR);
        mi.operands.push(MachineOperand::PhysReg(rt));
        mi.operands.push(MachineOperand::PhysReg(rn));
        mi.operands.push(MachineOperand::Imm(off));
        mi
    }

    fn make_str(rt: u32, rn: u32, off: i64) -> MachineInstr {
        let mut mi = MachineInstr::new(OP_STR);
        mi.operands.push(MachineOperand::PhysReg(rt));
        mi.operands.push(MachineOperand::PhysReg(rn));
        mi.operands.push(MachineOperand::Imm(off));
        mi
    }

    // ========================================================================
    // Test: MOV Xn, Xn elimination
    // ========================================================================

    #[test]
    fn test_eliminate_nop_moves() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_mov(X0 as u32, X0 as u32), // NOP move — should be removed
            make_mov(X1 as u32, X2 as u32), // real move — should stay
            make_mov(X3 as u32, X3 as u32), // NOP move — should be removed
        ];
        opt.eliminate_nop_moves(&mut instrs);
        assert_eq!(instrs.len(), 1);
        assert_eq!(instrs[0].opcode, OP_MOV);
        assert_eq!(opt.stats.mov_nop_eliminated, 2);
    }

    #[test]
    fn test_eliminate_nop_moves_no_effect() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_mov(X0 as u32, X1 as u32),
            make_mov(X2 as u32, X3 as u32),
        ];
        opt.eliminate_nop_moves(&mut instrs);
        assert_eq!(instrs.len(), 2);
        assert_eq!(opt.stats.mov_nop_eliminated, 0);
    }

    // ========================================================================
    // Test: CMP + B.cond fusion → CBZ/CBNZ
    // ========================================================================

    #[test]
    fn test_fuse_cmp_beq_to_cbz() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_cmp_imm(X0 as u32, 0),   // CMP X0, #0
            make_b_cond(ArmCond::EQ, 16), // B.EQ 16
            make_add_reg(X1 as u32, X2 as u32, X3 as u32),
        ];
        opt.fuse_cmp_branch(&mut instrs);
        // CMP + B.EQ → CBZ
        assert_eq!(instrs.len(), 2);
        assert_eq!(instrs[0].opcode, OP_CBZ);
        assert_eq!(opt.stats.cmp_branch_fused, 1);
    }

    #[test]
    fn test_fuse_cmp_bne_to_cbnz() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_cmp_imm(X5 as u32, 0),   // CMP X5, #0
            make_b_cond(ArmCond::NE, 32), // B.NE 32
        ];
        opt.fuse_cmp_branch(&mut instrs);
        assert_eq!(instrs.len(), 1);
        assert_eq!(instrs[0].opcode, OP_CBNZ);
        assert_eq!(opt.stats.cmp_branch_fused, 1);
    }

    #[test]
    fn test_no_fuse_cmp_nonzero() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_cmp_imm(X0 as u32, 5), // CMP X0, #5 — non-zero
            make_b_cond(ArmCond::EQ, 16),
        ];
        let orig_len = instrs.len();
        opt.fuse_cmp_branch(&mut instrs);
        // Should NOT fuse because CMP is not #0
        assert_eq!(instrs.len(), orig_len);
        assert_eq!(opt.stats.cmp_branch_fused, 0);
    }

    #[test]
    fn test_no_fuse_cmp_other_cond() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_cmp_imm(X0 as u32, 0),   // CMP X0, #0
            make_b_cond(ArmCond::LT, 16), // B.LT — not EQ/NE
        ];
        let orig_len = instrs.len();
        opt.fuse_cmp_branch(&mut instrs);
        assert_eq!(instrs.len(), orig_len);
        assert_eq!(opt.stats.cmp_branch_fused, 0);
    }

    // ========================================================================
    // Test: LDR/STR pairing → LDP/STP
    // ========================================================================

    #[test]
    fn test_pair_ldr_to_ldp() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_ldr(X0 as u32, X1 as u32, 0), // LDR X0, [X1, #0]
            make_ldr(X2 as u32, X1 as u32, 8), // LDR X2, [X1, #8]
        ];
        opt.pair_ldr_str(&mut instrs);
        assert_eq!(instrs.len(), 1);
        assert_eq!(instrs[0].opcode, OP_LDP);
        assert_eq!(opt.stats.ldr_str_paired, 1);
    }

    #[test]
    fn test_pair_str_to_stp() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_str(X3 as u32, X4 as u32, 16), // STR X3, [X4, #16]
            make_str(X5 as u32, X4 as u32, 24), // STR X5, [X4, #24]
        ];
        opt.pair_ldr_str(&mut instrs);
        assert_eq!(instrs.len(), 1);
        assert_eq!(instrs[0].opcode, OP_STP);
        assert_eq!(opt.stats.ldr_str_paired, 1);
    }

    #[test]
    fn test_no_pair_nonconsecutive_offsets() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_ldr(X0 as u32, X1 as u32, 0),  // offset 0
            make_ldr(X2 as u32, X1 as u32, 16), // offset 16 — not consecutive
        ];
        opt.pair_ldr_str(&mut instrs);
        assert_eq!(instrs.len(), 2);
        assert_eq!(opt.stats.ldr_str_paired, 0);
    }

    // ========================================================================
    // Test: Redundant load elimination (STR→LDR forwarding)
    // ========================================================================

    #[test]
    fn test_eliminate_redundant_load() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_str(X0 as u32, X10 as u32, 0), // STR X0, [X10, #0]
            make_ldr(X1 as u32, X10 as u32, 0), // LDR X1, [X10, #0] — redundant
        ];
        opt.eliminate_redundant_loads(&mut instrs);
        // LDR should become MOV
        assert_eq!(instrs.len(), 2);
        assert_eq!(instrs[1].opcode, OP_MOV);
        assert_eq!(opt.stats.redundant_loads_eliminated, 1);
    }

    #[test]
    fn test_no_forward_across_intervening_store() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_str(X0 as u32, X10 as u32, 0), // STR X0, [X10, #0]
            make_str(X1 as u32, X10 as u32, 0), // intervening store — updates value at addr
            make_ldr(X2 as u32, X10 as u32, 0), // LDR should get the latest stored value (X1)
        ];
        opt.eliminate_redundant_loads(&mut instrs);
        // LDR should be replaced with MOV X2, X1 (forwarded from second store)
        assert_eq!(instrs[2].opcode, OP_MOV);
        // Verify the MOV operands are correct (X2 = X1)
        if let MachineOperand::PhysReg(dst) = instrs[2].operands[0] {
            assert_eq!(dst, X2 as u32);
        }
        if let MachineOperand::PhysReg(src) = instrs[2].operands[1] {
            assert_eq!(src, X1 as u32);
        }
        assert!(opt.stats.redundant_loads_eliminated >= 1);
    }

    // ========================================================================
    // Test: ADD with zero immediate → MOV
    // ========================================================================

    #[test]
    fn test_eliminate_add_zero() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_add_imm(X3 as u32, X4 as u32, 0),  // ADD X3, X4, #0
            make_add_imm(X5 as u32, X6 as u32, 10), // ADD X5, X6, #10 — not zero
        ];
        opt.eliminate_add_zero(&mut instrs);
        assert_eq!(instrs[0].opcode, OP_MOV); // should become MOV
        assert_eq!(instrs[1].opcode, OP_ADD); // stays ADD
        assert_eq!(opt.stats.add_zero_eliminated, 1);
    }

    // ========================================================================
    // Test: Dead flag elimination
    // ========================================================================

    #[test]
    fn test_optimize_flag_setting() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_adds(X0 as u32, X1 as u32, X2 as u32), // ADDS — flags not used
            make_add_reg(X3 as u32, X4 as u32, X5 as u32), // next instr doesn't consume flags
        ];
        opt.optimize_flag_setting(&mut instrs);
        // ADDS should be demoted to ADD
        assert_eq!(instrs[0].opcode, OP_ADD);
        assert_eq!(opt.stats.dead_flags_removed, 1);
    }

    #[test]
    fn test_keep_flag_setting_when_used() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_adds(X0 as u32, X1 as u32, X2 as u32), // ADDS
            make_b_cond(ArmCond::EQ, 8),                // B.cond — consumes flags
        ];
        opt.optimize_flag_setting(&mut instrs);
        // ADDS should NOT be demoted
        assert_eq!(instrs[0].opcode, OP_ADDS);
        assert_eq!(opt.stats.dead_flags_removed, 0);
    }

    // ========================================================================
    // Test: Full optimize()
    // ========================================================================

    #[test]
    fn test_optimize_full() {
        let mut mf = make_mf();
        push_instr(&mut mf, make_mov(X0 as u32, X0 as u32)); // NOP move
        push_instr(&mut mf, make_cmp_imm(X1 as u32, 0)); // CMP #0
        push_instr(&mut mf, make_b_cond(ArmCond::EQ, 16)); // B.EQ
        push_instr(&mut mf, make_add_imm(X2 as u32, X3 as u32, 0)); // ADD #0
        push_instr(&mut mf, make_adds(X4 as u32, X5 as u32, X6 as u32)); // ADDS (dead flags)
        push_instr(&mut mf, make_add_reg(X7 as u32, X8 as u32, X9 as u32));

        let mut opt = ArmPeepholeOptimizer::new(true);
        let stats = opt.optimize(&mut mf);

        let _instrs = &mf.blocks[0].instructions;
        // After optimization:
        // - MOV X0,X0 removed
        // - CMP+B.EQ fused to CBZ
        // - ADD #0 → MOV
        // - ADDS → ADD (dead flags)
        assert_eq!(stats.mov_nop_eliminated, 1);
        assert_eq!(stats.cmp_branch_fused, 1);
        assert_eq!(stats.add_zero_eliminated, 1);
        assert_eq!(stats.dead_flags_removed, 1);
        assert!(stats.did_work());
    }

    #[test]
    fn test_optimize_no_effect() {
        let mut mf = make_mf();
        push_instr(&mut mf, make_add_reg(X0 as u32, X1 as u32, X2 as u32));
        push_instr(&mut mf, make_add_reg(X3 as u32, X4 as u32, X5 as u32));

        let mut opt = ArmPeepholeOptimizer::new(true);
        let stats = opt.optimize(&mut mf);

        assert!(!stats.did_work());
        assert_eq!(stats.net_change(), 0);
    }

    // ========================================================================
    // Test: Stats display
    // ========================================================================

    #[test]
    fn test_stats_net_change() {
        let mut stats = ArmOptStats::default();
        assert_eq!(stats.net_change(), 0);
        stats.total_removed = 5;
        stats.total_added = 3;
        assert_eq!(stats.net_change(), -2);
    }

    #[test]
    fn test_stats_did_work() {
        let mut stats = ArmOptStats::default();
        assert!(!stats.did_work());
        stats.total_removed = 1;
        assert!(stats.did_work());
    }

    // ========================================================================
    // Test: CMP register form should NOT be fused
    // ========================================================================

    #[test]
    fn test_no_fuse_cmp_reg() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            make_cmp_reg(X0 as u32, X1 as u32), // CMP X0, X1 — register form
            make_b_cond(ArmCond::EQ, 16),
        ];
        let orig_len = instrs.len();
        opt.fuse_cmp_branch(&mut instrs);
        // Should NOT fuse because CMP is register-register
        assert_eq!(instrs.len(), orig_len);
    }

    // ========================================================================
    // Test: ADR + ADD merging note
    // ========================================================================

    #[test]
    fn test_adr_add_merge() {
        let mut opt = ArmPeepholeOptimizer::new(true);
        let mut instrs = vec![
            {
                let mut mi = MachineInstr::new(OP_ADR);
                mi.operands.push(MachineOperand::PhysReg(X0 as u32));
                mi.operands
                    .push(MachineOperand::Label("some_label".to_string()));
                mi
            },
            make_add_imm(X0 as u32, X0 as u32, 16),
        ];
        opt.merge_adr_add(&mut instrs);
        assert_eq!(opt.stats.adr_add_merged, 1);
    }
}