llvm-native-core-ext 0.1.0

//! ARM/AArch64 Backend — instruction selection, register allocation,
//! calling convention, and assembly printer for the ARM architecture.
//! Phase 6 — LLVM.TARGET.ARM.1 Court.
//!
//! Clean-room behavioral reconstruction from the ARM Architecture
//! Reference Manual (ARM ARM), AAPCS64/AAPCS calling convention
//! specifications, and AArch64 instruction set encoding. Zero LLVM
//! source code consultation.
//!
//! Supports:
//! - AArch64 (ARMv8-A, 64-bit): full GP register set (x0-x30, sp),
//!   3-operand RISC instructions, AAPCS64 calling convention
//! - ARM32 (ARMv7-A, 32-bit): 16 GP registers (r0-r15), conditional
//!   execution, AAPCS calling convention

use llvm_native_core::codegen::{
    MachineBasicBlock, MachineFunction, MachineInstr, MachineOperand, PhysReg, VirtReg,
};
use llvm_native_core::value::ValueRef;
use std::collections::HashMap;

// ============================================================================
// ARM Opcode Constants
// ============================================================================

/// AArch64 opcode constants for codegen.
#[allow(dead_code)]
pub mod aarch64_opcodes {
    pub const NOP: u32 = 0;
    pub const MOV: u32 = 1;
    pub const ADD: u32 = 2;
    pub const SUB: u32 = 3;
    pub const MUL: u32 = 4;
    pub const AND: u32 = 5;
    pub const ORR: u32 = 6;
    pub const EOR: u32 = 7;
    pub const LSL: u32 = 8;
    pub const LSR: u32 = 9;
    pub const ASR: u32 = 10;
    pub const LDR: u32 = 11;
    pub const STR: u32 = 12;
    pub const STP: u32 = 13;
    pub const LDP: u32 = 14;
    pub const B: u32 = 15;
    pub const BL: u32 = 16;
    pub const RET: u32 = 17;
    pub const CMP: u32 = 18;
    pub const ADRP: u32 = 19;
    pub const ADR: u32 = 20;
    pub const CSEL: u32 = 21;
    pub const CSINV: u32 = 22;
    pub const MOVK: u32 = 23;
    pub const MOVZ: u32 = 24;
    pub const SBFM: u32 = 25;
    pub const UBFM: u32 = 26;
}

/// ARM32 opcode constants for codegen.
#[allow(dead_code)]
pub mod arm32_opcodes {
    pub const NOP: u32 = 100;
    pub const MOV: u32 = 101;
    pub const MOVW: u32 = 102;
    pub const MOVT: u32 = 103;
    pub const ADD: u32 = 104;
    pub const SUB: u32 = 105;
    pub const MUL: u32 = 106;
    pub const AND: u32 = 107;
    pub const ORR: u32 = 108;
    pub const EOR: u32 = 109;
    pub const LDR: u32 = 110;
    pub const STR: u32 = 111;
    pub const B: u32 = 112;
    pub const BL: u32 = 113;
    pub const BX: u32 = 114;
    pub const CMP: u32 = 115;
    pub const PUSH: u32 = 116;
    pub const POP: u32 = 117;
    #[allow(non_upper_case_globals)]
    pub const MOV_rr: u32 = 118;
}

// ============================================================================
// AArch64 Register Definitions (per AAPCS64)
// ============================================================================

/// AArch64 general-purpose register constants.
#[allow(dead_code)]
pub mod aarch64_regs {
    use llvm_native_core::codegen::PhysReg;
    pub const X0: PhysReg = 0;
    pub const X1: PhysReg = 1;
    pub const X2: PhysReg = 2;
    pub const X3: PhysReg = 3;
    pub const X4: PhysReg = 4;
    pub const X5: PhysReg = 5;
    pub const X6: PhysReg = 6;
    pub const X7: PhysReg = 7;
    pub const X8: PhysReg = 8;
    pub const X9: PhysReg = 9;
    pub const X10: PhysReg = 10;
    pub const X11: PhysReg = 11;
    pub const X12: PhysReg = 12;
    pub const X13: PhysReg = 13;
    pub const X14: PhysReg = 14;
    pub const X15: PhysReg = 15;
    pub const X16: PhysReg = 16;
    pub const X17: PhysReg = 17;
    pub const X18: PhysReg = 18;
    pub const X19: PhysReg = 19;
    pub const X20: PhysReg = 20;
    pub const X21: PhysReg = 21;
    pub const X22: PhysReg = 22;
    pub const X23: PhysReg = 23;
    pub const X24: PhysReg = 24;
    pub const X25: PhysReg = 25;
    pub const X26: PhysReg = 26;
    pub const X27: PhysReg = 27;
    pub const X28: PhysReg = 28;
    pub const X29: PhysReg = 29;
    pub const X30: PhysReg = 30;
    pub const FP: PhysReg = 29; // x29 — frame pointer
    pub const LR: PhysReg = 30; // x30 — link register
    pub const SP: PhysReg = 31; // xzr/sp
}

/// ARM32 general-purpose register constants.
#[allow(dead_code)]
pub mod arm32_regs {
    use llvm_native_core::codegen::PhysReg;
    pub const R0: PhysReg = 0;
    pub const R1: PhysReg = 1;
    pub const R2: PhysReg = 2;
    pub const R3: PhysReg = 3;
    pub const R4: PhysReg = 4;
    pub const R5: PhysReg = 5;
    pub const R6: PhysReg = 6;
    pub const R7: PhysReg = 7;
    pub const R8: PhysReg = 8;
    pub const R9: PhysReg = 9;
    pub const R10: PhysReg = 10;
    pub const R11: PhysReg = 11;
    pub const R12: PhysReg = 12;
    pub const R13: PhysReg = 13;
    pub const R14: PhysReg = 14;
    pub const R15: PhysReg = 15;
    pub const SP: PhysReg = 13;
    pub const LR: PhysReg = 14;
    pub const PC: PhysReg = 15;
}

// ============================================================================
// AArch64 Register Allocator
// ============================================================================

/// Linear-scan register allocator for AArch64.
///
/// AAPCS64 register usage:
///   x0-x7:   parameter/result registers (caller-saved)
///   x8-x15:  caller-saved temporaries
///   x16-x17: intra-procedure-call scratch (caller-saved)
///   x18:     platform register (if needed)
///   x19-x28: callee-saved
///   x29:     frame pointer (FP)
///   x30:     link register (LR)
///   x31:     stack pointer / zero register
///
/// For our allocator, we use: x0-x7, x9-x15 (skip x8 for indirect result)
pub struct AArch64RegisterAllocator {
    pub assignments: HashMap<VirtReg, PhysReg>,
    /// Available registers for allocation (caller-saved pool)
    available: Vec<PhysReg>,
}

impl AArch64RegisterAllocator {
    pub fn new() -> Self {
        Self {
            assignments: HashMap::new(),
            // x0-x7 for params/results, x9-x15 for temporaries
            available: vec![
                aarch64_regs::X0,
                aarch64_regs::X1,
                aarch64_regs::X2,
                aarch64_regs::X3,
                aarch64_regs::X4,
                aarch64_regs::X5,
                aarch64_regs::X6,
                aarch64_regs::X7,
                aarch64_regs::X9,
                aarch64_regs::X10,
                aarch64_regs::X11,
                aarch64_regs::X12,
                aarch64_regs::X13,
                aarch64_regs::X14,
                aarch64_regs::X15,
            ],
        }
    }

    /// Allocate physical registers to virtual registers.
    pub fn allocate(&mut self, mf: &mut MachineFunction) {
        let mut next_phys = 0usize;

        for bb in &mut mf.blocks {
            for mi in &mut bb.instructions {
                if let Some(def) = mi.def {
                    let phys = if let Some(&p) = self.assignments.get(&def) {
                        p
                    } else {
                        let p = self.available[next_phys % self.available.len()];
                        self.assignments.insert(def, p);
                        next_phys += 1;
                        p
                    };
                    mi.operands.insert(0, MachineOperand::PhysReg(phys));
                }

                for op in &mut mi.operands {
                    if let MachineOperand::Reg(vr) = *op {
                        if let Some(&phys) = self.assignments.get(&vr) {
                            *op = MachineOperand::PhysReg(phys);
                        }
                    }
                }
            }
        }
    }
}

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

// ============================================================================
// ARM32 Register Allocator
// ============================================================================

/// Linear-scan register allocator for ARM32.
///
/// AAPCS register usage:
///   r0-r3:   parameter/result registers (caller-saved)
///   r4-r11:  callee-saved (r9 may be platform-specific)
///   r12:     IP (intra-procedure-call scratch)
///   r13:     SP
///   r14:     LR
///   r15:     PC
pub struct ARM32RegisterAllocator {
    pub assignments: HashMap<VirtReg, PhysReg>,
    available: Vec<PhysReg>,
}

impl ARM32RegisterAllocator {
    pub fn new() -> Self {
        Self {
            assignments: HashMap::new(),
            // r0-r3 for params, r4-r10 for callee-saved temps
            available: vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
        }
    }

    pub fn allocate(&mut self, mf: &mut MachineFunction) {
        let mut next_phys = 0usize;

        for bb in &mut mf.blocks {
            for mi in &mut bb.instructions {
                if let Some(def) = mi.def {
                    let phys = if let Some(&p) = self.assignments.get(&def) {
                        p
                    } else {
                        let p = self.available[next_phys % self.available.len()];
                        self.assignments.insert(def, p);
                        next_phys += 1;
                        p
                    };
                    mi.operands.insert(0, MachineOperand::PhysReg(phys));
                }

                for op in &mut mi.operands {
                    if let MachineOperand::Reg(vr) = *op {
                        if let Some(&phys) = self.assignments.get(&vr) {
                            *op = MachineOperand::PhysReg(phys);
                        }
                    }
                }
            }
        }
    }
}

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

// ============================================================================
// AArch64 Instruction Selector
// ============================================================================

/// Instruction selector for AArch64: lowers LLVM IR to AArch64 MachineInstrs.
///
/// Key differences from X86-64:
/// - 3-operand RISC instructions (dst, src1, src2)
/// - No push/pop; use STP/LDP with SP adjustment
/// - Return via RET (branch to LR)
/// - 31 GP registers (x0-x30, xzr/sp)
pub struct AArch64InstructionSelector;

impl AArch64InstructionSelector {
    /// Select machine instructions for a function's IR, targeting AArch64.
    pub fn select(mf: &mut MachineFunction, func: &ValueRef) {
        let f = func.borrow();
        let mut vreg_map: HashMap<usize, VirtReg> = HashMap::new();

        for op in &f.operands {
            let ir_bb = op.borrow();
            if !ir_bb.is_basic_block() {
                continue;
            }

            let mut mbb = MachineBasicBlock {
                name: ir_bb.name.clone(),
                instructions: Vec::new(),
                successors: Vec::new(),
            };

            for inst_val in &ir_bb.operands {
                let inst = inst_val.borrow();
                if !inst.is_instruction() {
                    continue;
                }

                let num_ops = inst.operands.len();
                let is_void = inst.ty.is_void();

                if is_void && num_ops == 0 {
                    // ret void → RET
                    mbb.instructions
                        .push(MachineInstr::new(aarch64_opcodes::RET));
                } else if num_ops == 1 && inst.operands[0].borrow().is_basic_block() {
                    // br label %dest → B label
                    let dest = inst.operands[0].borrow().name.clone();
                    mbb.successors.push(dest.clone());
                    let mut mi = MachineInstr::new(aarch64_opcodes::B);
                    mi.push_label(&dest);
                    mbb.instructions.push(mi);
                } else if num_ops == 2 && !is_void {
                    // Binary ops: all map to ADD for simple ISel
                    let def = mf.new_vreg();
                    let op0_id = inst.operands[0].borrow().vid as usize;
                    let op1_id = inst.operands[1].borrow().vid as usize;
                    let vr0 = *vreg_map.entry(op0_id).or_insert_with(|| mf.new_vreg());
                    let vr1 = *vreg_map.entry(op1_id).or_insert_with(|| mf.new_vreg());

                    let mut mi = MachineInstr::new(aarch64_opcodes::ADD).with_def(def);
                    mi.push_reg(vr0);
                    if inst.operands[1].borrow().is_constant() {
                        mi.push_imm(inst.operands[1].borrow().name.parse().unwrap_or(0));
                    } else {
                        mi.push_reg(vr1);
                    }
                    vreg_map.insert(inst.vid as usize, def);
                    mbb.instructions.push(mi);
                } else if num_ops == 1 && inst.ty.is_pointer() {
                    // alloca → SUB sp, sp, #size (handled in prologue)
                    let def = mf.new_vreg();
                    let mut mi = MachineInstr::new(aarch64_opcodes::SUB).with_def(def);
                    mi.push_reg(aarch64_regs::SP);
                    mi.push_reg(aarch64_regs::SP);
                    // Placeholder size; actual size handled by frame lowering
                    mi.push_imm(16);
                    vreg_map.insert(inst.vid as usize, def);
                    mbb.instructions.push(mi);
                } else if num_ops == 2 && is_void {
                    // store: STR src, [dst]
                    let src_id = inst.operands[0].borrow().vid as usize;
                    let dst_id = inst.operands[1].borrow().vid as usize;
                    let vr_src = *vreg_map.entry(src_id).or_insert_with(|| mf.new_vreg());
                    let vr_dst = *vreg_map.entry(dst_id).or_insert_with(|| mf.new_vreg());
                    let mut mi = MachineInstr::new(aarch64_opcodes::STR);
                    mi.push_reg(vr_src);
                    mi.push_reg(vr_dst);
                    mbb.instructions.push(mi);
                } else if num_ops == 1 && !is_void {
                    // load: LDR dst, [src]
                    let src_id = inst.operands[0].borrow().vid as usize;
                    let vr_src = *vreg_map.entry(src_id).or_insert_with(|| mf.new_vreg());
                    let def = mf.new_vreg();
                    let mut mi = MachineInstr::new(aarch64_opcodes::LDR).with_def(def);
                    mi.push_reg(vr_src);
                    vreg_map.insert(inst.vid as usize, def);
                    mbb.instructions.push(mi);
                }
            }

            mf.push_block(mbb);
        }
    }
}

// ============================================================================
// ARM32 Instruction Selector
// ============================================================================

/// Instruction selector for ARM32: lowers LLVM IR to ARM32 MachineInstrs.
///
/// Key differences from AArch64:
/// - Conditional execution on most instructions
/// - Push/pop for stack operations
/// - BX LR for return
/// - MOV/MOVW/MOVT for constant loading
pub struct ARM32InstructionSelector;

impl ARM32InstructionSelector {
    pub fn select(mf: &mut MachineFunction, func: &ValueRef) {
        let f = func.borrow();
        let mut vreg_map: HashMap<usize, VirtReg> = HashMap::new();

        for op in &f.operands {
            let ir_bb = op.borrow();
            if !ir_bb.is_basic_block() {
                continue;
            }

            let mut mbb = MachineBasicBlock {
                name: ir_bb.name.clone(),
                instructions: Vec::new(),
                successors: Vec::new(),
            };

            for inst_val in &ir_bb.operands {
                let inst = inst_val.borrow();
                if !inst.is_instruction() {
                    continue;
                }

                let num_ops = inst.operands.len();
                let is_void = inst.ty.is_void();

                if is_void && num_ops == 0 {
                    // ret void → BX LR
                    let mut mi = MachineInstr::new(arm32_opcodes::BX);
                    mi.push_reg(arm32_regs::LR);
                    mbb.instructions.push(mi);
                } else if num_ops == 1 && inst.operands[0].borrow().is_basic_block() {
                    // br label → B label
                    let dest = inst.operands[0].borrow().name.clone();
                    mbb.successors.push(dest.clone());
                    let mut mi = MachineInstr::new(arm32_opcodes::B);
                    mi.push_label(&dest);
                    mbb.instructions.push(mi);
                } else if num_ops == 2 && !is_void {
                    // Binary ops
                    let def = mf.new_vreg();
                    let op0_id = inst.operands[0].borrow().vid as usize;
                    let op1_id = inst.operands[1].borrow().vid as usize;
                    let vr0 = *vreg_map.entry(op0_id).or_insert_with(|| mf.new_vreg());
                    let vr1 = *vreg_map.entry(op1_id).or_insert_with(|| mf.new_vreg());

                    let mut mi = MachineInstr::new(arm32_opcodes::ADD).with_def(def);
                    mi.push_reg(vr0);
                    if inst.operands[1].borrow().is_constant() {
                        mi.push_imm(inst.operands[1].borrow().name.parse().unwrap_or(0));
                    } else {
                        mi.push_reg(vr1);
                    }
                    vreg_map.insert(inst.vid as usize, def);
                    mbb.instructions.push(mi);
                } else if num_ops == 1 && inst.ty.is_pointer() {
                    // alloca → SUB sp, sp, #size
                    let def = mf.new_vreg();
                    let mut mi = MachineInstr::new(arm32_opcodes::SUB).with_def(def);
                    mi.push_reg(arm32_regs::SP);
                    mi.push_reg(arm32_regs::SP);
                    mi.push_imm(16);
                    vreg_map.insert(inst.vid as usize, def);
                    mbb.instructions.push(mi);
                }
            }

            mf.push_block(mbb);
        }
    }
}

// ============================================================================
// AArch64 Assembly Printer
// ============================================================================

/// Emits AArch64 assembly text from a MachineFunction.
///
/// AArch64 assembly syntax (GNU as):
/// - Labels: `name:`
/// - Instructions: `    mnemonic op1, op2, op3`
/// - Registers: x0-x30, sp, lr, fp (no % prefix)
/// - Immediates: #42 (hash prefix)
/// - Comments: // or /* */
/// - Directives: .globl, .text, .align
pub struct AArch64AsmPrinter {
    pub output: String,
}

impl AArch64AsmPrinter {
    pub fn new() -> Self {
        Self {
            output: String::new(),
        }
    }

    pub fn print_function(&mut self, mf: &MachineFunction) {
        // Function prologue (AAPCS64)
        self.output.push_str(&format!(".globl {}\n", mf.name));
        self.output
            .push_str(&format!(".type {}, @function\n", mf.name));
        self.output.push_str(&format!("{}:\n", mf.name));

        // Standard prologue: stp x29, x30, [sp, #-16]!  ;  mov x29, sp
        self.output.push_str("    stp x29, x30, [sp, #-16]!\n");
        self.output.push_str("    mov x29, sp\n");

        for bb in &mf.blocks {
            if !bb.name.is_empty() && bb.name != "entry" {
                self.output.push_str(&format!(".L{}:\n", bb.name));
            }
            for mi in &bb.instructions {
                self.print_instr(mi);
            }
        }

        // Epilogue: ldp x29, x30, [sp], #16  ;  ret
        self.output.push_str(".Lfunc_end:\n");
        self.output.push_str("    ldp x29, x30, [sp], #16\n");
        self.output.push_str("    ret\n");
        self.output
            .push_str(&format!(".size {}, .Lfunc_end-{}\n", mf.name, mf.name));
        self.output.push('\n');
    }

    fn print_instr(&mut self, mi: &MachineInstr) {
        let mnemonic = llvm_native_core::target_info::aarch64_mnemonic(mi.opcode);
        self.output.push_str(&format!("    {}", mnemonic));

        let mut first = true;
        for op in &mi.operands {
            if first {
                self.output.push(' ');
                first = false;
            } else {
                self.output.push_str(", ");
            }
            match op {
                MachineOperand::PhysReg(r) => self
                    .output
                    .push_str(&llvm_native_core::target_info::aarch64_reg_name(*r).to_string()),
                MachineOperand::Reg(r) => self.output.push_str(&format!("v{}", r)),
                MachineOperand::Imm(i) => self.output.push_str(&format!("#{}", i)),
                MachineOperand::Label(l) => self.output.push_str(&format!(".L{}", l)),
                MachineOperand::Global(g) => self.output.push_str(g),
            }
        }
        self.output.push('\n');
    }
}

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

// ============================================================================
// ARM32 Assembly Printer
// ============================================================================

/// Emits ARM32 assembly text from a MachineFunction.
///
/// ARM32 assembly syntax (GNU as):
/// - Labels: `name:`
/// - Instructions: `    mnemonic op1, op2, op3`
/// - Registers: r0-r15, sp, lr, pc (no % prefix)
/// - Immediates: #42 (hash prefix)
/// - Directives: .globl, .text, .align, .arm / .thumb
pub struct ARM32AsmPrinter {
    pub output: String,
}

impl ARM32AsmPrinter {
    pub fn new() -> Self {
        Self {
            output: String::new(),
        }
    }

    pub fn print_function(&mut self, mf: &MachineFunction) {
        self.output.push_str(&format!(".globl {}\n", mf.name));
        self.output
            .push_str(&format!(".type {}, %function\n", mf.name));
        self.output.push_str(&format!("{}:\n", mf.name));

        // Standard prologue: push {fp, lr}  ;  mov fp, sp
        self.output.push_str("    push {fp, lr}\n");
        self.output.push_str("    mov fp, sp\n");

        for bb in &mf.blocks {
            if !bb.name.is_empty() && bb.name != "entry" {
                self.output.push_str(&format!(".L{}:\n", bb.name));
            }
            for mi in &bb.instructions {
                self.print_instr(mi);
            }
        }

        // Epilogue: pop {fp, pc}
        self.output.push_str("    pop {fp, pc}\n");
        self.output.push('\n');
    }

    fn print_instr(&mut self, mi: &MachineInstr) {
        let mnemonic = Self::mnemonic(mi.opcode);
        self.output.push_str(&format!("    {}", mnemonic));

        let mut first = true;
        for op in &mi.operands {
            if first {
                self.output.push(' ');
                first = false;
            } else {
                self.output.push_str(", ");
            }
            match op {
                MachineOperand::PhysReg(r) => self
                    .output
                    .push_str(&llvm_native_core::target_info::arm32_reg_name(*r)),
                MachineOperand::Reg(r) => self.output.push_str(&format!("v{}", r)),
                MachineOperand::Imm(i) => self.output.push_str(&format!("#{}", i)),
                MachineOperand::Label(l) => self.output.push_str(&format!(".L{}", l)),
                MachineOperand::Global(g) => self.output.push_str(g),
            }
        }
        self.output.push('\n');
    }

    fn mnemonic(opcode: u32) -> &'static str {
        match opcode {
            100 => "nop",
            101 => "mov",
            102 => "movw",
            103 => "movt",
            104 => "add",
            105 => "sub",
            106 => "mul",
            107 => "and",
            108 => "orr",
            109 => "eor",
            110 => "ldr",
            111 => "str",
            112 => "b",
            113 => "bl",
            114 => "bx",
            115 => "cmp",
            116 => "push",
            117 => "pop",
            118 => "mov",
            _ => "unknown",
        }
    }
}

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

// ============================================================================
// High-Level ARM Backend Driver
// ============================================================================

/// Compile an IR function to AArch64 assembly.
pub fn compile_aarch64(func: &ValueRef) -> String {
    let f = func.borrow();
    let mut mf = MachineFunction::new(&f.name);

    // Instruction selection
    AArch64InstructionSelector::select(&mut mf, func);

    // Register allocation
    let mut ra = AArch64RegisterAllocator::new();
    ra.allocate(&mut mf);

    // Assembly printing
    let mut printer = AArch64AsmPrinter::new();
    printer.print_function(&mf);
    printer.output
}

/// Compile an IR function to ARM32 assembly.
pub fn compile_arm32(func: &ValueRef) -> String {
    let f = func.borrow();
    let mut mf = MachineFunction::new(&f.name);

    // Instruction selection
    ARM32InstructionSelector::select(&mut mf, func);

    // Register allocation
    let mut ra = ARM32RegisterAllocator::new();
    ra.allocate(&mut mf);

    // Assembly printing
    let mut printer = ARM32AsmPrinter::new();
    printer.print_function(&mf);
    printer.output
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use llvm_native_core::basic_block::new_basic_block;
    use llvm_native_core::function::new_function;
    use llvm_native_core::instruction::ret_void;
    use llvm_native_core::types::Type;

    /// Helper: build a simple IR function (add two constants, return void).
    fn build_simple_ir() -> ValueRef {
        let func = new_function("test_fn", Type::void(), &[]);
        let entry = new_basic_block("entry");
        let ret = ret_void();
        entry.borrow_mut().push_operand(ret);
        func.borrow_mut().push_operand(entry.clone());
        func
    }

    /// Helper: build a two-block IR function with a branch.
    fn build_branch_ir() -> ValueRef {
        let func = new_function("branch_fn", Type::void(), &[]);
        let entry = new_basic_block("entry");
        let target = new_basic_block("target");
        let ret = ret_void();

        // br label %target
        let br = llvm_native_core::instruction::br(target.clone());

        entry.borrow_mut().push_operand(br);
        target.borrow_mut().push_operand(ret);
        func.borrow_mut().push_operand(entry.clone());
        func.borrow_mut().push_operand(target.clone());
        func
    }

    // === AArch64 Register Allocator Tests ===

    #[test]
    fn test_aarch64_reg_allocator_create() {
        let ra = AArch64RegisterAllocator::new();
        assert!(!ra.available.is_empty());
        assert!(ra.available.contains(&aarch64_regs::X0));
        assert!(ra.available.contains(&aarch64_regs::X15));
    }

    #[test]
    fn test_aarch64_reg_allocator_allocate() {
        let mut mf = MachineFunction::new("test");
        let vr = mf.new_vreg();

        let mut mbb = MachineBasicBlock {
            name: "entry".into(),
            instructions: Vec::new(),
            successors: Vec::new(),
        };
        let mi = MachineInstr::new(aarch64_opcodes::ADD).with_def(vr);
        mbb.instructions.push(mi);
        mf.push_block(mbb);

        let mut ra = AArch64RegisterAllocator::new();
        ra.allocate(&mut mf);

        // Verify the def was assigned a physical register
        assert!(!ra.assignments.is_empty());
    }

    // === ARM32 Register Allocator Tests ===

    #[test]
    fn test_arm32_reg_allocator_create() {
        let ra = ARM32RegisterAllocator::new();
        assert!(!ra.available.is_empty());
        assert!(ra.available.contains(&arm32_regs::R0));
    }

    // === AArch64 Instruction Selection Tests ===

    #[test]
    fn test_aarch64_isel_simple_function() {
        let ir_func = build_simple_ir();
        let f = ir_func.borrow();
        let mut mf = MachineFunction::new(&f.name);
        AArch64InstructionSelector::select(&mut mf, &ir_func);

        assert!(!mf.blocks.is_empty());
        // Should have at least the entry block with a RET
        let entry = &mf.blocks[0];
        assert!(!entry.instructions.is_empty());
    }

    #[test]
    fn test_aarch64_isel_branch_function() {
        let ir_func = build_branch_ir();
        let f = ir_func.borrow();
        let mut mf = MachineFunction::new(&f.name);
        AArch64InstructionSelector::select(&mut mf, &ir_func);

        assert!(!mf.blocks.is_empty());
        // Entry block should have a B instruction targeting "target"
        let entry = &mf.blocks[0];
        let has_branch = entry
            .instructions
            .iter()
            .any(|mi| mi.opcode == aarch64_opcodes::B);
        assert!(has_branch, "Entry block should contain a branch");
    }

    // === ARM32 Instruction Selection Tests ===

    #[test]
    fn test_arm32_isel_simple_function() {
        let ir_func = build_simple_ir();
        let f = ir_func.borrow();
        let mut mf = MachineFunction::new(&f.name);
        ARM32InstructionSelector::select(&mut mf, &ir_func);

        assert!(!mf.blocks.is_empty());
        let entry = &mf.blocks[0];
        assert!(!entry.instructions.is_empty());
    }

    #[test]
    fn test_arm32_isel_branch_function() {
        let ir_func = build_branch_ir();
        let f = ir_func.borrow();
        let mut mf = MachineFunction::new(&f.name);
        ARM32InstructionSelector::select(&mut mf, &ir_func);

        assert!(!mf.blocks.is_empty());
        let entry = &mf.blocks[0];
        let has_branch = entry
            .instructions
            .iter()
            .any(|mi| mi.opcode == arm32_opcodes::B);
        assert!(has_branch, "Entry block should contain a branch");
    }

    // === AArch64 Assembly Printer Tests ===

    #[test]
    fn test_aarch64_asm_printer_simple() {
        let ir_func = build_simple_ir();
        let asm = compile_aarch64(&ir_func);

        assert!(asm.contains(".globl test_fn"));
        assert!(asm.contains("test_fn:"));
        // Check for AArch64 prologue
        assert!(asm.contains("stp x29, x30"));
        assert!(asm.contains("mov x29, sp"));
        // Check for ret
        assert!(asm.contains("ret"));
        // Check for epilogue
        assert!(asm.contains("ldp x29, x30"));
    }

    #[test]
    fn test_aarch64_asm_printer_branch() {
        let ir_func = build_branch_ir();
        let asm = compile_aarch64(&ir_func);

        assert!(asm.contains(".globl branch_fn"));
        assert!(asm.contains(".Ltarget:"));
    }

    #[test]
    fn test_aarch64_asm_printer_has_size_directive() {
        let ir_func = build_simple_ir();
        let asm = compile_aarch64(&ir_func);

        assert!(asm.contains(".size test_fn"));
    }

    // === ARM32 Assembly Printer Tests ===

    #[test]
    fn test_arm32_asm_printer_simple() {
        let ir_func = build_simple_ir();
        let asm = compile_arm32(&ir_func);

        assert!(asm.contains(".globl test_fn"));
        assert!(asm.contains("test_fn:"));
        // Check for ARM32 prologue
        assert!(asm.contains("push {fp, lr}"));
        assert!(asm.contains("mov fp, sp"));
        // Check for epilogue
        assert!(asm.contains("pop {fp, pc}"));
    }

    #[test]
    fn test_arm32_asm_printer_branch() {
        let ir_func = build_branch_ir();
        let asm = compile_arm32(&ir_func);

        assert!(asm.contains(".globl branch_fn"));
        assert!(asm.contains(".Ltarget:"));
    }

    // === Full Compilation Pipeline Tests ===

    #[test]
    fn test_full_aarch64_pipeline() {
        // Build a simple function, compile for AArch64, verify output structure
        let ir_func = build_simple_ir();
        let asm = compile_aarch64(&ir_func);

        // Should have proper structure
        assert!(!asm.is_empty());
        assert!(asm.contains(".globl"));
        assert!(asm.contains(":"));
        assert!(asm.contains("ret"));
    }

    #[test]
    fn test_full_arm32_pipeline() {
        let ir_func = build_simple_ir();
        let asm = compile_arm32(&ir_func);

        assert!(!asm.is_empty());
        assert!(asm.contains(".globl"));
        assert!(asm.contains(":"));
    }

    #[test]
    fn test_aarch64_register_constants() {
        // Verify register numbering
        assert_eq!(aarch64_regs::X0, 0);
        assert_eq!(aarch64_regs::X30, 30);
        assert_eq!(aarch64_regs::SP, 31);
        assert_eq!(aarch64_regs::FP, 29);
        assert_eq!(aarch64_regs::LR, 30);
    }

    #[test]
    fn test_arm32_register_constants() {
        assert_eq!(arm32_regs::R0, 0);
        assert_eq!(arm32_regs::R15, 15);
        assert_eq!(arm32_regs::SP, 13);
        assert_eq!(arm32_regs::LR, 14);
        assert_eq!(arm32_regs::PC, 15);
    }

    #[test]
    fn test_aarch64_opcode_constants() {
        assert_eq!(aarch64_opcodes::NOP, 0);
        assert_eq!(aarch64_opcodes::RET, 17);
        assert_eq!(aarch64_opcodes::B, 15);
        assert_eq!(aarch64_opcodes::BL, 16);
        assert_eq!(aarch64_opcodes::ADD, 2);
        assert_eq!(aarch64_opcodes::SUB, 3);
        assert_eq!(aarch64_opcodes::STR, 12);
        assert_eq!(aarch64_opcodes::LDR, 11);
    }

    #[test]
    fn test_arm32_opcode_constants() {
        assert_eq!(arm32_opcodes::NOP, 100);
        assert_eq!(arm32_opcodes::B, 112);
        assert_eq!(arm32_opcodes::BL, 113);
        assert_eq!(arm32_opcodes::BX, 114);
        assert_eq!(arm32_opcodes::ADD, 104);
        assert_eq!(arm32_opcodes::SUB, 105);
    }

    // === Cross-architecture: verify different codegen paths ===

    #[test]
    fn test_aarch64_vs_arm32_output_differ() {
        let ir_func = build_simple_ir();
        let a64_asm = compile_aarch64(&ir_func);
        let a32_asm = compile_arm32(&ir_func);

        // The two backends should produce different assembly
        assert_ne!(a64_asm, a32_asm);
    }

    #[test]
    fn test_aarch64_machineinstr_count() {
        let ir_func = build_branch_ir();
        let f = ir_func.borrow();
        let mut mf = MachineFunction::new(&f.name);
        AArch64InstructionSelector::select(&mut mf, &ir_func);

        // Should have 2 blocks (entry + target)
        assert_eq!(mf.blocks.len(), 2);
        // Entry block should have 1 instruction (B)
        assert_eq!(mf.blocks[0].instructions.len(), 1);
        // Target block should have 1 instruction (RET)
        assert_eq!(mf.blocks[1].instructions.len(), 1);
    }

    #[test]
    fn test_arm32_machineinstr_count() {
        let ir_func = build_branch_ir();
        let f = ir_func.borrow();
        let mut mf = MachineFunction::new(&f.name);
        ARM32InstructionSelector::select(&mut mf, &ir_func);

        assert_eq!(mf.blocks.len(), 2);
        assert_eq!(mf.blocks[0].instructions.len(), 1);
        assert_eq!(mf.blocks[1].instructions.len(), 1);
    }
}