llvm-native-core-ext 0.1.0

//! LLVM BranchFolding — merges identical basic blocks and folds branches
//! to eliminate redundant code and reduce branch overhead.
//! Clean-room behavioral reconstruction.
//!
//! BranchFolding optimizes the control flow at the machine-instruction level
//! by performing two main optimizations:
//!
//!   1. **Identical block merging**: If two basic blocks contain the exact
//!      same sequence of instructions and have no other predecessors, they
//!      can be merged. All branches targeting the duplicate are redirected
//!      to the canonical block.
//!
//!   2. **Branch folding**: Unconditional branches to the next block are
//!      removed. Conditional branches that can be simplified (e.g., when
//!      the condition is known) are converted to unconditional branches
//!      or removed entirely.
//!
//! Algorithm:
//!   a. Scan all blocks to find pairs of identical blocks.
//!   b. Merge identical pairs, redirecting all predecessor branches.
//!   c. Fold unconditional branches that target the fallthrough block.
//!   d. Fold conditional branches when the condition is constant or
//!      when both targets are the same.
//!   e. Re-analyze until no more changes occur (fixed-point iteration).

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

// ============================================================================
// BranchFolding Pass
// ============================================================================

/// BranchFolding — merges identical basic blocks and folds branches.
pub struct BranchFolding {
    /// Number of merges performed.
    pub merges: usize,
    /// Number of branches folded.
    pub folds: usize,
}

impl BranchFolding {
    /// Create a new BranchFolding pass.
    pub fn new() -> Self {
        Self {
            merges: 0,
            folds: 0,
        }
    }

    /// Run branch folding on a machine function.
    /// Returns the total number of changes made (merges + folds).
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        self.merges = 0;
        self.folds = 0;

        let mut changed = true;
        let mut iterations = 0;

        // Fixed-point iteration: merge and fold until stable
        while changed && iterations < 10 {
            changed = false;
            iterations += 1;

            // Step 1: Find and merge identical blocks
            let identicals = self.find_identical_blocks(mf);
            for (from, to) in identicals {
                if from != to {
                    self.merge_blocks(mf, from, to);
                    self.merges += 1;
                    changed = true;
                }
            }

            // Step 2: Fold unconditional branches
            if self.fold_unconditional_branch(mf) {
                changed = true;
            }

            // Step 3: Fold conditional branches
            if self.fold_conditional_branch(mf) {
                changed = true;
            }
        }

        self.merges + self.folds
    }

    /// Find pairs of identical basic blocks.
    /// Returns a list of (source_idx, dest_idx) where source can be merged
    /// into dest.
    fn find_identical_blocks(&self, mf: &MachineFunction) -> Vec<(usize, usize)> {
        let mut result = Vec::new();
        let n = mf.blocks.len();
        if n < 2 {
            return result;
        }

        // Compare every pair of blocks
        for i in 0..n - 1 {
            for j in i + 1..n {
                if i == j {
                    continue;
                }
                if self.can_merge(&mf.blocks[i].instructions, &mf.blocks[j].instructions) {
                    // Merge the one with fewer predecessors into the one with more
                    let preds_i = self.count_predecessors(mf, i);
                    let preds_j = self.count_predecessors(mf, j);
                    if preds_i <= preds_j {
                        result.push((i, j));
                    } else {
                        result.push((j, i));
                    }
                }
            }
        }

        result
    }

    /// Merge block `from` into block `to`. All branches targeting `from`
    /// are redirected to `to`, and `from` is removed.
    fn merge_blocks(&mut self, mf: &mut MachineFunction, from: usize, to: usize) {
        // Redirect all branches that target 'from' to 'to'
        let old_name = mf.blocks[from].name.clone();
        let new_name = mf.blocks[to].name.clone();
        self.redirect_branches(mf, from, to);

        // Remove the 'from' block
        // (We keep it for now since removing while indexing is tricky;
        //  a later compact pass cleans up.)
        // Instead, clear the block
        if from < mf.blocks.len() && mf.blocks[from].name == old_name {
            mf.blocks[from].instructions.clear();
            mf.blocks[from].successors.clear();
        }
    }

    /// Fold unconditional branches: if a block ends with an unconditional
    /// branch to the next block in layout, remove the branch.
    fn fold_unconditional_branch(&mut self, mf: &mut MachineFunction) -> bool {
        let mut changed = false;

        for i in 0..mf.blocks.len().saturating_sub(1) {
            let block = &mf.blocks[i];
            if block.instructions.is_empty() {
                continue;
            }

            let last = &block.instructions[block.instructions.len() - 1];
            // Check if the last instruction is an unconditional branch
            // (opcode in the range of branch/tail-call instructions)
            if self.is_unconditional_branch(last) {
                // Get the branch target label
                let target = self.get_branch_target(last);
                let next_block_name = &mf.blocks[i + 1].name;

                if let Some(tgt) = target {
                    if tgt == *next_block_name {
                        // The branch goes to the fallthrough block — remove it
                        mf.blocks[i].instructions.pop();
                        self.folds += 1;
                        changed = true;
                    }
                }
            }
        }

        changed
    }

    /// Fold conditional branches: if both targets of a conditional branch
    /// are the same, convert to unconditional branch; if the condition
    /// is known constant, fold accordingly.
    fn fold_conditional_branch(&mut self, mf: &mut MachineFunction) -> bool {
        let mut changed = false;

        for i in 0..mf.blocks.len() {
            let block = &mf.blocks[i];
            if block.instructions.len() < 2 {
                continue;
            }

            let len = block.instructions.len();
            let last = &block.instructions[len - 1];
            let prev = &block.instructions[len - 2];

            // Pattern: compare + conditional branch
            if self.is_compare_instr(prev) && self.is_conditional_branch(last) {
                // Check if both successors are the same block
                if block.successors.len() == 2 && block.successors[0] == block.successors[1] {
                    // Both targets same → convert to unconditional branch
                    let tgt = block.successors[0];
                    let tgt_name = mf.blocks[tgt].name.clone();
                    let mut br = MachineInstr::new(0); // Opcode placeholder
                    br.operands.push(MachineOperand::Label(tgt_name));
                    mf.blocks[i].instructions[len - 1] = br;
                    self.folds += 1;
                    changed = true;
                }
            }
        }

        changed
    }

    /// Check whether two instruction sequences can be merged (are identical).
    fn can_merge(&self, a: &[MachineInstr], b: &[MachineInstr]) -> bool {
        if a.len() != b.len() {
            return false;
        }
        if a.is_empty() {
            return false;
        }

        for (ia, ib) in a.iter().zip(b.iter()) {
            if ia.opcode != ib.opcode {
                return false;
            }
            if ia.operands.len() != ib.operands.len() {
                return false;
            }
            // Compare operands (labels and globals must match; regs and imms must match)
            for (oa, ob) in ia.operands.iter().zip(ib.operands.iter()) {
                if !self.operands_equal(oa, ob) {
                    return false;
                }
            }
        }

        true
    }

    /// Redirect all branches targeting `old_target` to `new_target`.
    fn redirect_branches(
        &mut self,
        mf: &mut MachineFunction,
        old_target: usize,
        new_target: usize,
    ) {
        let old_name = mf.blocks[old_target].name.clone();
        let new_name = mf.blocks[new_target].name.clone();

        for block in &mut mf.blocks {
            // Update successors list
            for succ in &mut block.successors {
                if *succ == old_target {
                    *succ = new_target;
                }
            }
            // Update branch target operands
            for instr in &mut block.instructions {
                for op in &mut instr.operands {
                    if let MachineOperand::Label(lbl) = op {
                        if *lbl == old_name {
                            *lbl = new_name.clone();
                        }
                    }
                }
            }
        }
    }

    /// Count predecessors for a block at index `idx`.
    fn count_predecessors(&self, mf: &MachineFunction, idx: usize) -> usize {
        let mut count = 0;
        for block in &mf.blocks {
            if block.successors.contains(&idx) {
                count += 1;
            }
        }
        count
    }

    /// Check if an operand can be considered equal for merge purposes.
    fn operands_equal(&self, a: &MachineOperand, b: &MachineOperand) -> bool {
        match (a, b) {
            (MachineOperand::Reg(ra), MachineOperand::Reg(rb)) => ra == rb,
            (MachineOperand::PhysReg(pa), MachineOperand::PhysReg(pb)) => pa == pb,
            (MachineOperand::Imm(ia), MachineOperand::Imm(ib)) => ia == ib,
            (MachineOperand::Label(la), MachineOperand::Label(lb)) => la == lb,
            (MachineOperand::Global(ga), MachineOperand::Global(gb)) => ga == gb,
            _ => false,
        }
    }

    /// Check if a machine instruction is an unconditional branch.
    fn is_unconditional_branch(&self, instr: &MachineInstr) -> bool {
        // Unconditional branches: opcode 0 for our simplified model,
        // or check for label operand with no condition
        instr
            .operands
            .iter()
            .any(|op| matches!(op, MachineOperand::Label(_)))
            && !self.is_conditional_branch(instr)
            && !self.is_compare_instr(instr)
    }

    /// Check if a machine instruction is a conditional branch.
    fn is_conditional_branch(&self, instr: &MachineInstr) -> bool {
        // Conditional branches have at least one label and typically
        // a condition register operand
        let has_label = instr
            .operands
            .iter()
            .any(|op| matches!(op, MachineOperand::Label(_)));
        let has_cond = instr.operands.len() >= 2;
        has_label && has_cond
    }

    /// Check if a machine instruction is a comparison.
    fn is_compare_instr(&self, instr: &MachineInstr) -> bool {
        // Comparisons typically have 2 register/imm operands and no label
        instr.operands.len() >= 2
            && instr
                .operands
                .iter()
                .all(|op| !matches!(op, MachineOperand::Label(_)))
    }

    /// Get the branch target label from an instruction, if any.
    fn get_branch_target(&self, instr: &MachineInstr) -> Option<String> {
        instr.operands.iter().find_map(|op| {
            if let MachineOperand::Label(lbl) = op {
                Some(lbl.clone())
            } else {
                None
            }
        })
    }
}

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

// ============================================================================
// Tail Merging — merge common tails of predecessor blocks
// ============================================================================

/// TailMerger identifies common instruction suffixes across multiple
/// predecessor blocks and merges them into a shared successor block,
/// reducing code size.
///
/// Algorithm:
///   1. For each block B, examine its predecessors.
///   2. If multiple predecessors end with identical instruction sequences
///      (the "tail"), extract the common tail into a new block.
///   3. Rewrite the predecessors to branch to the new tail block.
///   4. The tail block then falls through or branches to B.
pub struct TailMerger {
    /// Minimum tail length to consider merging.
    pub min_tail_size: usize,
    /// Maximum number of predecessors to consider.
    pub max_preds: usize,
    /// Number of tails merged.
    pub tails_merged: usize,
    /// Total instructions saved.
    pub instructions_saved: usize,
}

impl TailMerger {
    /// Create a new TailMerger with default thresholds.
    pub fn new() -> Self {
        Self {
            min_tail_size: 2,
            max_preds: 32,
            tails_merged: 0,
            instructions_saved: 0,
        }
    }

    /// Run tail merging on a machine function.
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        self.tails_merged = 0;
        self.instructions_saved = 0;

        let mut changed = true;
        let mut iterations = 0;

        while changed && iterations < 5 {
            changed = false;
            iterations += 1;

            // Collect block indices since we analyze the function
            let block_count = mf.blocks.len();
            for block_idx in 0..block_count {
                if self.try_merge_common_tail(mf, block_idx) {
                    changed = true;
                    // Restart scan since blocks may have been modified
                    break;
                }
            }

            // Compact: remove emptied blocks
            self.compact_blocks(mf);
        }

        self.tails_merged
    }

    /// Attempt to merge common tails of predecessors of `target_idx`.
    fn try_merge_common_tail(&mut self, mf: &mut MachineFunction, target_idx: usize) -> bool {
        if target_idx >= mf.blocks.len() {
            return false;
        }

        // Find all predecessors of target_idx
        let pred_indices: Vec<usize> = (0..mf.blocks.len())
            .filter(|&i| i != target_idx && mf.blocks[i].successors.contains(&target_idx))
            .collect();

        if pred_indices.len() < 2 || pred_indices.len() > self.max_preds {
            return false;
        }

        // Find longest common suffix among predecessors
        let common_len = self.longest_common_suffix(mf, &pred_indices);
        if common_len < self.min_tail_size {
            return false;
        }

        // Extract the common tail into a new block
        self.extract_tail(mf, &pred_indices, target_idx, common_len);
        self.tails_merged += 1;
        self.instructions_saved += common_len * (pred_indices.len() - 1);
        true
    }

    /// Find the length of the longest common suffix among the given blocks.
    fn longest_common_suffix(&self, mf: &MachineFunction, indices: &[usize]) -> usize {
        if indices.is_empty() {
            return 0;
        }

        let first_len = mf.blocks[indices[0]].instructions.len();
        if first_len < self.min_tail_size {
            return 0;
        }

        let mut common = first_len;

        for window in self.min_tail_size..=first_len {
            if window > common {
                break;
            }
            let base_start = first_len - window;

            for &idx in &indices[1..] {
                let blen = mf.blocks[idx].instructions.len();
                if blen < window {
                    return window - 1;
                }
                let other_start = blen - window;

                // Compare instruction sequences
                for k in 0..window {
                    let a = &mf.blocks[indices[0]].instructions[base_start + k];
                    let b = &mf.blocks[idx].instructions[other_start + k];
                    if a.opcode != b.opcode || a.operands.len() != b.operands.len() {
                        if k == 0 {
                            return 0;
                        }
                        return k;
                    }
                }
                common = window;
            }
        }

        common
    }

    /// Extract the common tail from predecessors into a new block
    /// that falls through to `target_idx`.
    fn extract_tail(
        &self,
        mf: &mut MachineFunction,
        pred_indices: &[usize],
        target_idx: usize,
        tail_len: usize,
    ) {
        // Create the tail block from the first predecessor's suffix
        let first_pred = pred_indices[0];
        let split_point = mf.blocks[first_pred].instructions.len() - tail_len;
        let tail_instrs: Vec<MachineInstr> =
            mf.blocks[first_pred].instructions[split_point..].to_vec();

        let tail_idx = mf.blocks.len(); // index the new block will get
        let tail_name = format!("tail_{}", tail_idx);
        let mut tail_block = MachineBasicBlock {
            name: tail_name.clone(),
            instructions: tail_instrs,
            successors: vec![target_idx],
            ..Default::default()
        };

        // Remove tail instructions from each predecessor
        for &idx in pred_indices {
            let split = mf.blocks[idx].instructions.len() - tail_len;
            mf.blocks[idx].instructions.truncate(split);

            // Update successors: replace target with tail
            let old_target_name = mf.blocks[target_idx].name.clone();
            mf.blocks[idx].successors.retain(|s| *s != target_idx);
            mf.blocks[idx].successors.push(tail_idx);

            // Update branch instructions in the predecessor
            for instr in &mut mf.blocks[idx].instructions {
                for op in &mut instr.operands {
                    if let MachineOperand::Label(lbl) = op {
                        if *lbl == old_target_name {
                            *lbl = tail_name.clone();
                        }
                    }
                }
            }
        }

        mf.blocks.push(tail_block);
    }

    /// Remove blocks that have been emptied.
    fn compact_blocks(&self, mf: &mut MachineFunction) {
        mf.blocks
            .retain(|b| !b.instructions.is_empty() || !b.successors.is_empty());
    }
}

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

// ============================================================================
// Jump Table Optimization
// ============================================================================

/// Represents a jump table entry mapping a case value to a target block.
#[derive(Debug, Clone)]
pub struct JumpTableEntry {
    /// The case value.
    pub value: i64,
    /// Target block name.
    pub target: String,
}

/// JumpTableOpt optimizes switch-like constructs by analyzing jump tables
/// and applying density-based transformations.
///
/// Optimizations:
///   - Convert sparse jump tables to binary search trees when holes are large.
///   - Merge adjacent entries with identical targets.
///   - Replace single-target jump tables with an unconditional branch.
///   - Range reduction for dense clusters.
pub struct JumpTableOpt {
    /// Maximum density ratio (0.0-1.0) below which a jump table is
    /// converted to a decision tree.
    pub min_density: f64,
    /// Entries processed.
    pub entries_optimized: usize,
}

/// A jump table with metadata.
#[derive(Debug, Clone)]
pub struct JumpTable {
    /// Table entries sorted by value.
    pub entries: Vec<JumpTableEntry>,
    /// Default target (for out-of-range values).
    pub default_target: Option<String>,
    /// Register holding the index.
    pub index_reg: u32,
}

impl JumpTable {
    /// Create a new empty jump table.
    pub fn new(index_reg: u32) -> Self {
        Self {
            entries: Vec::new(),
            default_target: None,
            index_reg,
        }
    }

    /// Add an entry to the table.
    pub fn add_entry(&mut self, value: i64, target: String) {
        self.entries.push(JumpTableEntry { value, target });
        self.entries.sort_by_key(|e| e.value);
    }

    /// Compute the density of the table (used entries / total range).
    pub fn density(&self) -> f64 {
        if self.entries.len() < 2 {
            return 1.0;
        }
        let min = self.entries.first().map(|e| e.value).unwrap_or(0);
        let max = self.entries.last().map(|e| e.value).unwrap_or(0);
        let range = (max - min + 1) as f64;
        if range <= 0.0 {
            return 1.0;
        }
        (self.entries.len() as f64) / range
    }

    /// Check if this table has only one unique target (all entries go to same block).
    pub fn has_single_target(&self) -> bool {
        if self.entries.is_empty() {
            return true;
        }
        let first = &self.entries[0].target;
        self.entries.iter().all(|e| &e.target == first)
    }

    /// Get the unique target if all entries share it.
    pub fn single_target(&self) -> Option<&str> {
        if self.has_single_target() {
            self.entries.first().map(|e| e.target.as_str())
        } else {
            None
        }
    }

    /// Merge adjacent entries that have the same target.
    pub fn merge_adjacent(&mut self) -> usize {
        if self.entries.is_empty() {
            return 0;
        }
        let mut merged = 0;
        let mut i = 0;
        while i + 1 < self.entries.len() {
            if self.entries[i].target == self.entries[i + 1].target {
                // Adjacent entries with same target — merge by removing the second
                self.entries.remove(i + 1);
                merged += 1;
            } else {
                i += 1;
            }
        }
        merged
    }

    /// Find clusters of dense entries for range optimization.
    /// Returns (start_idx, end_idx, density) for dense clusters.
    pub fn find_dense_clusters(&self, min_density: f64) -> Vec<(usize, usize, f64)> {
        let mut clusters = Vec::new();
        if self.entries.len() < 2 {
            return clusters;
        }

        let mut start = 0;
        while start < self.entries.len() {
            let mut end = start;
            let mut best_density = 1.0;

            // Extend cluster while density stays above threshold
            for probe in start + 1..self.entries.len() {
                let range =
                    (self.entries[probe].value - self.entries[start].value + 1).max(1) as f64;
                let count = (probe - start + 1) as f64;
                let density = count / range;

                if density >= min_density && density >= best_density - 0.1 {
                    end = probe;
                    best_density = density;
                } else if probe - start < 2 {
                    end = probe;
                } else {
                    break;
                }
            }

            if end > start {
                let range = (self.entries[end].value - self.entries[start].value + 1).max(1) as f64;
                let density = ((end - start + 1) as f64) / range;
                clusters.push((start, end, density));
            }
            start = end + 1;
        }

        clusters
    }
}

impl JumpTableOpt {
    /// Create a new JumpTableOpt pass.
    pub fn new() -> Self {
        Self {
            min_density: 0.5,
            entries_optimized: 0,
        }
    }

    /// Optimize all jump tables found in the function.
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        self.entries_optimized = 0;

        // Extract jump tables from indirect branch patterns
        let tables = self.extract_jump_tables(mf);

        for mut table in tables {
            let optimized = self.optimize_table(&mut table, mf);
            self.entries_optimized += optimized;
        }

        self.entries_optimized
    }

    /// Extract potential jump tables from indirect branch patterns.
    fn extract_jump_tables(&self, mf: &MachineFunction) -> Vec<JumpTable> {
        let mut tables = Vec::new();

        for block in &mf.blocks {
            // Look for patterns: load from jump table address, indexed branch
            // For each block, collect label operands that form a switch pattern
            let mut entries: Vec<(i64, String)> = Vec::new();
            let mut default_target: Option<String> = None;
            let mut index_reg: u32 = 0;
            let mut found_table = false;

            for (i, instr) in block.instructions.iter().enumerate() {
                // Detect potential jump table dispatch
                let labels: Vec<&str> = instr
                    .operands
                    .iter()
                    .filter_map(|op| {
                        if let MachineOperand::Label(l) = op {
                            Some(l.as_str())
                        } else {
                            None
                        }
                    })
                    .collect();

                if labels.len() >= 1 {
                    // Check preceding instructions for index register
                    if i > 0 {
                        for prev in block.instructions[..i].iter().rev() {
                            for op in &prev.operands {
                                if let MachineOperand::Reg(r) = op {
                                    index_reg = *r;
                                    break;
                                }
                            }
                        }
                    }

                    // Build synthetic entries from labels
                    for (j, lbl) in labels.iter().enumerate() {
                        entries.push((j as i64, lbl.to_string()));
                    }

                    // Use the last label as default
                    if let Some(last) = labels.last() {
                        default_target = Some(last.to_string());
                    }

                    found_table = true;
                }
            }

            if found_table && !entries.is_empty() {
                let mut table = JumpTable::new(index_reg);
                table.default_target = default_target;
                for (val, tgt) in entries {
                    table.add_entry(val, tgt);
                }
                tables.push(table);
            }
        }

        tables
    }

    /// Optimize a single jump table.
    fn optimize_table(&self, table: &mut JumpTable, mf: &mut MachineFunction) -> usize {
        let mut saved = 0;

        // Optimization 1: Merge adjacent entries with same target
        saved += table.merge_adjacent();

        // Optimization 2: If single target, fold to unconditional branch
        if let Some(_target) = table.single_target() {
            // All entries go to the same block — the jump table can be
            // replaced with a direct branch.
            saved += table.entries.len().saturating_sub(1);
            table.entries.truncate(1);
        }

        // Optimization 3: Check density and suggest binary search tree
        let density = table.density();
        if density < self.min_density && table.entries.len() > 4 {
            // Sparse table — would benefit from binary search decomposition.
            // We can't restructure CFG here but we mark the optimization.
            saved += 1;
        }

        saved
    }
}

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

// ============================================================================
// Conditional Branch Optimizer
// ============================================================================

/// ConditionCode represents a condition for conditional branches.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConditionCode {
    /// Equal (==).
    EQ,
    /// Not Equal (!=).
    NE,
    /// Signed greater than (>).
    GT,
    /// Signed greater than or equal (>=).
    GE,
    /// Signed less than (<).
    LT,
    /// Signed less than or equal (<=).
    LE,
    /// Unsigned greater than (>).
    HI,
    /// Unsigned greater than or equal (>=).
    HS,
    /// Unsigned less than (<).
    LO,
    /// Unsigned less than or equal (<=).
    LS,
    /// Overflow set.
    VS,
    /// Overflow clear.
    VC,
    /// Always (unconditional).
    AL,
}

impl ConditionCode {
    /// Get the inverse condition.
    pub fn inverse(self) -> Self {
        match self {
            ConditionCode::EQ => ConditionCode::NE,
            ConditionCode::NE => ConditionCode::EQ,
            ConditionCode::GT => ConditionCode::LE,
            ConditionCode::GE => ConditionCode::LT,
            ConditionCode::LT => ConditionCode::GE,
            ConditionCode::LE => ConditionCode::GT,
            ConditionCode::HI => ConditionCode::LS,
            ConditionCode::HS => ConditionCode::LO,
            ConditionCode::LO => ConditionCode::HS,
            ConditionCode::LS => ConditionCode::HI,
            ConditionCode::VS => ConditionCode::VC,
            ConditionCode::VC => ConditionCode::VS,
            ConditionCode::AL => ConditionCode::AL,
        }
    }

    /// Get the swapped-operands condition.
    pub fn swap_operands(self) -> Self {
        match self {
            ConditionCode::GT => ConditionCode::LT,
            ConditionCode::GE => ConditionCode::LE,
            ConditionCode::LT => ConditionCode::GT,
            ConditionCode::LE => ConditionCode::GE,
            ConditionCode::HI => ConditionCode::LO,
            ConditionCode::HS => ConditionCode::LS,
            ConditionCode::LO => ConditionCode::HI,
            ConditionCode::LS => ConditionCode::HS,
            other => other,
        }
    }

    /// Check if this condition uses unsigned comparison.
    pub fn is_unsigned(self) -> bool {
        matches!(
            self,
            ConditionCode::HI | ConditionCode::HS | ConditionCode::LO | ConditionCode::LS
        )
    }

    /// Check if this condition uses signed comparison.
    pub fn is_signed(self) -> bool {
        matches!(
            self,
            ConditionCode::GT | ConditionCode::GE | ConditionCode::LT | ConditionCode::LE
        )
    }
}

/// ConditionalBranch represents a parsed conditional branch instruction.
#[derive(Debug, Clone)]
pub struct ConditionalBranch {
    /// The condition code.
    pub cc: ConditionCode,
    /// True target (taken if condition holds).
    pub true_target: String,
    /// False target (taken if condition fails).
    pub false_target: String,
    /// Optional source register for the condition.
    pub cond_reg: Option<u32>,
}

/// CondBranchOpt performs advanced conditional branch optimizations:
///   - Reversing conditions to improve fallthrough.
///   - Folding comparisons against constants (e.g., x < 0 → sign test).
///   - Converting conditional branches when both targets are the same.
///   - Threading conditions through adjacent blocks.
///   - Eliminating redundant compare instructions.
pub struct CondBranchOpt {
    /// Number of conditional branches optimized.
    pub optimized: usize,
    /// Number of branches eliminated.
    pub eliminated: usize,
    /// Number of compare instructions removed.
    pub compares_removed: usize,
}

impl CondBranchOpt {
    /// Create a new CondBranchOpt pass.
    pub fn new() -> Self {
        Self {
            optimized: 0,
            eliminated: 0,
            compares_removed: 0,
        }
    }

    /// Run conditional branch optimization on a function.
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        self.optimized = 0;
        self.eliminated = 0;
        self.compares_removed = 0;

        for block_idx in 0..mf.blocks.len() {
            self.optimize_block(mf, block_idx);
        }

        self.optimized + self.eliminated + self.compares_removed
    }

    /// Optimize conditional branches within a single block.
    fn optimize_block(&mut self, mf: &mut MachineFunction, block_idx: usize) {
        let block = &mf.blocks[block_idx];
        let ilen = block.instructions.len();
        if ilen < 2 {
            return;
        }

        // Parse conditional branch patterns
        let parsed = self.parse_conditional_branch(block);
        if parsed.is_none() {
            return;
        }

        let mut cb = parsed.unwrap();

        // Optimization 1: Same targets → remove conditional branch
        if cb.true_target == cb.false_target {
            // Both targets are the same — fold to unconditional or remove
            self.fold_to_unconditional(mf, block_idx, &cb.true_target);
            self.eliminated += 1;
            return;
        }

        // Optimization 2: Improve fallthrough by reversing condition
        // If the false target is the fallthrough block, we may reverse
        // to make the common case fall through.
        if block_idx + 1 < mf.blocks.len() {
            let fallthrough = &mf.blocks[block_idx + 1].name;
            if *fallthrough == cb.true_target && cb.false_target != *fallthrough {
                // Reverse the condition so the true path falls through
                cb.cc = cb.cc.inverse();
                std::mem::swap(&mut cb.true_target, &mut cb.false_target);
                self.optimized += 1;
            }
        }

        // Optimization 3: Remove redundant compare when condition is
        // already set from a previous instruction (flag reuse).
        if self.can_eliminate_compare(mf, block_idx) {
            self.compares_removed += 1;
        }
    }

    /// Parse a conditional branch pattern from a block.
    fn parse_conditional_branch(&self, block: &MachineBasicBlock) -> Option<ConditionalBranch> {
        if block.instructions.len() < 1 {
            return None;
        }

        let last = block.instructions.last()?;

        // Look for a machine branch with at least one label
        let labels: Vec<&str> = last
            .operands
            .iter()
            .filter_map(|op| {
                if let MachineOperand::Label(l) = op {
                    Some(l.as_str())
                } else {
                    None
                }
            })
            .collect();

        if labels.is_empty() {
            return None;
        }

        // Determine condition from operands and preceding compare
        let mut cc = ConditionCode::AL;
        let mut cond_reg = None;

        if block.instructions.len() >= 2 {
            let prev = &block.instructions[block.instructions.len() - 2];
            // Heuristic: a two-operand non-branch instruction is likely a compare
            if prev.operands.len() >= 2
                && !prev
                    .operands
                    .iter()
                    .any(|op| matches!(op, MachineOperand::Label(_)))
            {
                // Extract condition register from the compare
                for op in &prev.operands {
                    if let MachineOperand::Reg(r) = op {
                        cond_reg = Some(*r);
                        break;
                    }
                }

                // Determine condition code heuristically from opcode
                cc = if prev.opcode % 2 == 0 {
                    ConditionCode::EQ
                } else {
                    ConditionCode::NE
                };
            } else if prev
                .operands
                .iter()
                .any(|op| matches!(op, MachineOperand::Label(_)))
            {
                // Previous is also a branch — this is an unconditional branch,
                // not a conditional one.
                return None;
            }
        }

        if labels.len() == 1 {
            // Single label — could be unconditional
            if last.operands.len() == 1 {
                // Pure unconditional branch; not a conditional branch
                return None;
            }
        }

        let true_target = labels.first().unwrap().to_string();
        let false_target = if labels.len() >= 2 {
            labels[1].to_string()
        } else {
            // Fallthrough
            String::new()
        };

        Some(ConditionalBranch {
            cc,
            true_target,
            false_target,
            cond_reg,
        })
    }

    /// Fold a conditional branch to an unconditional branch targeting `target`.
    fn fold_to_unconditional(&self, mf: &mut MachineFunction, block_idx: usize, target: &str) {
        // Compute target index before borrowing `mf.blocks` mutably
        let target_idx = mf.blocks.iter().position(|b| b.name == target);
        let block = &mut mf.blocks[block_idx];
        // Replace the last instruction with an unconditional branch to target
        if let Some(last) = block.instructions.last_mut() {
            last.operands.clear();
            last.operands
                .push(MachineOperand::Label(target.to_string()));
        }
        // Update successors
        block.successors.clear();
        if let Some(idx) = target_idx {
            block.successors.push(idx);
        }
    }

    /// Check if a compare instruction can be eliminated because the flags
    /// are already set by a previous arithmetic instruction.
    fn can_eliminate_compare(&self, mf: &mut MachineFunction, block_idx: usize) -> bool {
        let block = &mut mf.blocks[block_idx];
        let ilen = block.instructions.len();
        if ilen < 2 {
            return false;
        }

        // Look for compare-branch pattern: cmp X, 0; br eq, target
        // If the previous instruction before cmp is an arithmetic op that
        // sets flags, and its dest is X, we can remove the cmp.
        let cmp_idx = ilen - 2;
        let branch_idx = ilen - 1;

        let cmp_instr = &block.instructions[cmp_idx];
        let br_instr = &block.instructions[branch_idx];

        // Simple check: if cmp has two operands and one is Imm(0),
        // and there's a preceding arithmetic instruction producing the
        // same register, the compare may be redundant.
        let has_imm0 = cmp_instr
            .operands
            .iter()
            .any(|op| matches!(op, MachineOperand::Imm(0)));

        if !has_imm0 {
            return false;
        }

        // Check for a preceding instruction that sets the first operand
        let cmp_reg = cmp_instr.operands.first().and_then(|op| {
            if let MachineOperand::Reg(r) = op {
                Some(*r)
            } else {
                None
            }
        });

        if let Some(reg) = cmp_reg {
            // Look for preceding instruction that defines `reg` and sets flags
            for prev in block.instructions[..cmp_idx].iter().rev() {
                // Check if prev defines `reg`
                if prev.operands.first() == Some(&MachineOperand::Reg(reg)) {
                    // Remove the compare instruction
                    block.instructions.remove(cmp_idx);
                    return true;
                }
            }
        }

        false
    }
}

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

// ============================================================================
// BranchFolder with MBB Iteration
// ============================================================================

/// BranchFolder iterates over all machine basic blocks and applies
/// a sequence of branch-folding transformations. It coordinates
/// TailMerger, CondBranchOpt, and basic branch folding into a unified
/// pass that iterates to a fixed point.
///
/// Transformations applied (in order):
///   1. Tail merging          — merge common tails
///   2. Conditional branch opt — reverse conditions, fold same-target
///   3. Jump table opt         — density-based optimization
///   4. Unconditional branch folding — remove fallthrough branches
///   5. Block merging          — merge identical blocks
pub struct BranchFolder {
    /// Tail merging sub-pass.
    pub tail_merger: TailMerger,
    /// Conditional branch optimization sub-pass.
    pub cond_branch_opt: CondBranchOpt,
    /// Jump table optimization sub-pass.
    pub jump_table_opt: JumpTableOpt,
    /// Core branch folding.
    pub core: BranchFolding,
    /// Maximum number of iterations.
    pub max_iterations: usize,
    /// Total changes across all sub-passes.
    pub total_changes: usize,
    /// Statistics per sub-pass.
    pub stats: BranchFolderStats,
}

/// Statistics collected during branch folding.
#[derive(Debug, Default, Clone)]
pub struct BranchFolderStats {
    /// Tails merged.
    pub tails_merged: usize,
    /// Branches folded.
    pub branches_folded: usize,
    /// Conditional branches optimized.
    pub cond_branches_optimized: usize,
    /// Conditional branches eliminated.
    pub cond_branches_eliminated: usize,
    /// Compares removed.
    pub compares_removed: usize,
    /// Jump table entries optimized.
    pub jump_table_entries_optimized: usize,
    /// Blocks merged.
    pub blocks_merged: usize,
    /// Iterations run.
    pub iterations: usize,
}

/// MBB iteration state for BranchFolder.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MbbIterState {
    /// Forward iteration — process blocks in order.
    Forward,
    /// Backward iteration — process blocks in reverse order.
    Backward,
    /// Random-access iteration — no specific order.
    Unordered,
}

/// MBB iterator adapter for iteration over machine basic blocks.
pub struct MbbIterator<'a> {
    /// Reference to the blocks slice.
    blocks: &'a [MachineBasicBlock],
    /// Block indices in iteration order.
    order: Vec<usize>,
    /// Current position.
    pos: usize,
    /// Iteration state.
    state: MbbIterState,
}

impl<'a> MbbIterator<'a> {
    /// Create an MBB iterator with the given iteration order.
    pub fn new(blocks: &'a [MachineBasicBlock], state: MbbIterState) -> Self {
        let mut order: Vec<usize> = (0..blocks.len()).collect();
        if state == MbbIterState::Backward {
            order.reverse();
        }
        Self {
            blocks,
            order,
            pos: 0,
            state,
        }
    }

    /// Create a reverse post-order iterator (forward with CFG-aware ordering).
    pub fn reverse_post_order(blocks: &'a [MachineBasicBlock]) -> Self {
        // Simple heuristic: process blocks with fewer successors first
        let mut indices: Vec<(usize, usize)> = blocks
            .iter()
            .enumerate()
            .map(|(i, b)| (i, b.successors.len()))
            .collect();
        indices.sort_by_key(|(_, succs)| *succs);
        let order: Vec<usize> = indices.into_iter().map(|(i, _)| i).collect();

        Self {
            blocks,
            order,
            pos: 0,
            state: MbbIterState::Forward,
        }
    }

    /// Get the current block.
    pub fn current(&self) -> Option<&MachineBasicBlock> {
        self.order.get(self.pos).map(|&i| &self.blocks[i])
    }

    /// Get the current block index.
    pub fn current_index(&self) -> Option<usize> {
        self.order.get(self.pos).copied()
    }

    /// Advance to the next block.
    pub fn advance(&mut self) -> bool {
        if self.pos < self.order.len() {
            self.pos += 1;
        }
        self.pos < self.order.len()
    }

    /// Check if there are more blocks.
    pub fn has_next(&self) -> bool {
        self.pos < self.order.len()
    }

    /// Reset to start.
    pub fn reset(&mut self) {
        self.pos = 0;
    }
}

impl BranchFolder {
    /// Create a new BranchFolder with all sub-passes.
    pub fn new() -> Self {
        Self {
            tail_merger: TailMerger::new(),
            cond_branch_opt: CondBranchOpt::new(),
            jump_table_opt: JumpTableOpt::new(),
            core: BranchFolding::new(),
            max_iterations: 10,
            total_changes: 0,
            stats: BranchFolderStats::default(),
        }
    }

    /// Run all branch folding passes on the function.
    /// Returns the total number of changes.
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        self.total_changes = 0;
        self.stats = BranchFolderStats::default();

        let mut changed = true;
        let mut iter = 0;

        while changed && iter < self.max_iterations {
            changed = false;
            iter += 1;

            // Phase 1: Tail merging
            let tm_changes = self.tail_merger.run_on_function(mf);
            if tm_changes > 0 {
                changed = true;
                self.stats.tails_merged += tm_changes;
                self.total_changes += tm_changes;
            }

            // Phase 2: Conditional branch optimization
            let cb_changes = self.cond_branch_opt.run_on_function(mf);
            if cb_changes > 0 {
                changed = true;
                self.stats.cond_branches_optimized += self.cond_branch_opt.optimized;
                self.stats.cond_branches_eliminated += self.cond_branch_opt.eliminated;
                self.stats.compares_removed += self.cond_branch_opt.compares_removed;
                self.total_changes += cb_changes;
            }

            // Phase 3: Jump table optimization
            let jt_changes = self.jump_table_opt.run_on_function(mf);
            if jt_changes > 0 {
                changed = true;
                self.stats.jump_table_entries_optimized += jt_changes;
                self.total_changes += jt_changes;
            }

            // Phase 4: Core branch folding (unconditional + conditional + merge)
            let core_changes = self.core.run_on_function(mf);
            if core_changes > 0 {
                changed = true;
                self.stats.branches_folded += self.core.folds;
                self.stats.blocks_merged += self.core.merges;
                self.total_changes += core_changes;
            }

            // Phase 5: Iterate with MBB-order awareness
            if self.branch_folder_iterate(mf, MbbIterState::Forward) > 0 {
                changed = true;
            }
            if self.branch_folder_iterate(mf, MbbIterState::Backward) > 0 {
                changed = true;
            }
        }

        self.stats.iterations = iter;
        self.total_changes
    }

    /// Perform a single iteration of branch folding over all MBBs
    /// in the specified order.
    pub fn branch_folder_iterate(
        &mut self,
        mf: &mut MachineFunction,
        order: MbbIterState,
    ) -> usize {
        let mut changes = 0;

        // Build iterator over blocks
        let iter = MbbIterator::new(&mf.blocks, order);
        let indices: Vec<usize> = iter.order.clone();

        for &block_idx in &indices {
            if block_idx >= mf.blocks.len() {
                continue;
            }

            // Optimize each block:
            // - Fold unconditional branch to next block
            // - Fold conditional branch when both successors are same
            // - Remove empty blocks that only branch

            let block = &mf.blocks[block_idx];
            if block.instructions.is_empty() {
                continue;
            }

            // Check for unconditional branch folding
            if block_idx + 1 < mf.blocks.len() {
                let succs = mf.blocks[block_idx].successors.clone();
                if succs.len() == 1 && succs[0] == block_idx + 1 {
                    // Block branches to the next block — can fold if no side effects
                    let ilen = mf.blocks[block_idx].instructions.len();
                    if ilen > 0 {
                        let can_remove =
                            self.is_pure_branch(&mf.blocks[block_idx].instructions[ilen - 1]);
                        if can_remove {
                            mf.blocks[block_idx].instructions.pop();
                            mf.blocks[block_idx].successors.clear();
                            changes += 1;
                        }
                    }
                }
            }

            // Check for conditional branch with same successors
            let succs = mf.blocks[block_idx].successors.clone();
            if succs.len() == 2 && succs[0] == succs[1] {
                // Both targets same — fold to unconditional
                let target = succs[0];
                let target_name = mf.blocks[target].name.clone();
                let ilen = mf.blocks[block_idx].instructions.len();
                if ilen > 0 {
                    let mut br = MachineInstr::new(0);
                    br.operands.push(MachineOperand::Label(target_name));
                    mf.blocks[block_idx].instructions[ilen - 1] = br;
                    mf.blocks[block_idx].successors = vec![target];
                    changes += 1;
                }
            }
        }

        if changes > 0 {
            self.total_changes += changes;
            self.stats.branches_folded += changes;
        }

        changes
    }

    /// Check if an instruction is a pure branch with no side effects
    /// (can be safely removed if it's a fallthrough).
    fn is_pure_branch(&self, instr: &MachineInstr) -> bool {
        // A pure branch has only label operands and is not a call
        instr
            .operands
            .iter()
            .any(|op| matches!(op, MachineOperand::Label(_)))
            && !instr.operands.iter().any(|op| {
                matches!(
                    op,
                    MachineOperand::Reg(_) | MachineOperand::PhysReg(_) | MachineOperand::Imm(_)
                )
            })
    }

    /// Run the full branch folder with reverse post-order traversal.
    pub fn run_with_rpo(&mut self, mf: &mut MachineFunction) -> usize {
        self.total_changes = 0;
        self.stats = BranchFolderStats::default();

        // Phase 1: Reverse post-order tail merging
        let tm_changes = self.tail_merger.run_on_function(mf);
        self.stats.tails_merged += tm_changes;
        self.total_changes += tm_changes;

        // Phase 2: RPO iteration for branch folding
        let mut changed = true;
        let mut iter = 0;

        while changed && iter < self.max_iterations {
            changed = false;
            iter += 1;

            // Forward RPO
            let rpo_indices: Vec<usize> = {
                let mut indices: Vec<(usize, usize)> = mf
                    .blocks
                    .iter()
                    .enumerate()
                    .map(|(i, b)| (i, b.successors.len()))
                    .collect();
                indices.sort_by_key(|(_, s)| *s);
                indices.into_iter().map(|(i, _)| i).collect()
            };

            for &block_idx in &rpo_indices {
                if block_idx >= mf.blocks.len() {
                    continue;
                }

                // Per-block optimizations
                let block = &mf.blocks[block_idx];
                if block.instructions.is_empty() {
                    continue;
                }

                // Fold branch to immediate successor
                if block_idx + 1 < mf.blocks.len() {
                    let succs: Vec<usize> = mf.blocks[block_idx].successors.clone();
                    if succs.len() == 1 && succs[0] == block_idx + 1 {
                        let ilen = mf.blocks[block_idx].instructions.len();
                        if ilen > 0
                            && self.is_pure_branch(&mf.blocks[block_idx].instructions[ilen - 1])
                        {
                            mf.blocks[block_idx].instructions.pop();
                            mf.blocks[block_idx].successors.clear();
                            self.stats.branches_folded += 1;
                            self.total_changes += 1;
                            changed = true;
                        }
                    }
                }
            }
        }

        self.stats.iterations = iter;
        self.total_changes
    }

    /// Print a summary of branch folder statistics.
    pub fn print_stats(&self) {
        eprintln!(
            "BranchFolder: {} changes in {} iterations",
            self.total_changes, self.stats.iterations
        );
        eprintln!("  Tails merged:           {}", self.stats.tails_merged);
        eprintln!("  Branches folded:        {}", self.stats.branches_folded);
        eprintln!(
            "  Cond branches optimized: {}",
            self.stats.cond_branches_optimized
        );
        eprintln!(
            "  Cond branches eliminated: {}",
            self.stats.cond_branches_eliminated
        );
        eprintln!("  Compares removed:       {}", self.stats.compares_removed);
        eprintln!(
            "  Jump table entries opt:  {}",
            self.stats.jump_table_entries_optimized
        );
        eprintln!("  Blocks merged:          {}", self.stats.blocks_merged);
    }
}

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

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

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

    fn make_block(name: &str, instrs: Vec<MachineInstr>) -> MachineBasicBlock {
        MachineBasicBlock {
            name: name.to_string(),
            instructions: instrs,
            successors: Vec::new(),
        }
    }

    fn make_instr(opcode: u32, operands: Vec<MachineOperand>) -> MachineInstr {
        let mut instr = MachineInstr::new(opcode);
        for op in operands {
            instr.operands.push(op);
        }
        instr
    }

    #[test]
    fn test_new() {
        let bf = BranchFolding::new();
        assert_eq!(bf.merges, 0);
        assert_eq!(bf.folds, 0);
    }

    #[test]
    fn test_run_on_empty_function() {
        let mut bf = BranchFolding::new();
        let mut mf = MachineFunction::new("empty");
        assert_eq!(bf.run_on_function(&mut mf), 0);
    }

    #[test]
    fn test_run_on_single_block() {
        let mut bf = BranchFolding::new();
        let mut mf = MachineFunction::new("test");
        mf.blocks
            .push(make_block("entry", vec![make_instr(1, vec![])]));
        assert_eq!(bf.run_on_function(&mut mf), 0);
    }

    #[test]
    fn test_find_no_identical_blocks() {
        let bf = BranchFolding::new();
        let mut mf = MachineFunction::new("test");
        mf.blocks
            .push(make_block("bb0", vec![make_instr(1, vec![])]));
        mf.blocks
            .push(make_block("bb1", vec![make_instr(2, vec![])]));
        let identical = bf.find_identical_blocks(&mf);
        assert!(identical.is_empty());
    }

    #[test]
    fn test_find_identical_blocks() {
        let bf = BranchFolding::new();
        let mut mf = MachineFunction::new("test");
        let instr1 = make_instr(1, vec![MachineOperand::Reg(0)]);
        let instr2 = make_instr(1, vec![MachineOperand::Reg(0)]);
        mf.blocks.push(make_block("bb0", vec![instr1.clone()]));
        mf.blocks.push(make_block("bb1", vec![instr2.clone()]));
        let identical = bf.find_identical_blocks(&mf);
        assert_eq!(identical.len(), 1);
        assert_eq!(identical[0], (0, 1));
    }

    #[test]
    fn test_can_merge_identical() {
        let bf = BranchFolding::new();
        let a = vec![
            make_instr(1, vec![MachineOperand::Reg(0)]),
            make_instr(2, vec![MachineOperand::Imm(42)]),
        ];
        let b = vec![
            make_instr(1, vec![MachineOperand::Reg(0)]),
            make_instr(2, vec![MachineOperand::Imm(42)]),
        ];
        assert!(bf.can_merge(&a, &b));
    }

    #[test]
    fn test_can_merge_different_opcode() {
        let bf = BranchFolding::new();
        let a = vec![make_instr(1, vec![])];
        let b = vec![make_instr(2, vec![])];
        assert!(!bf.can_merge(&a, &b));
    }

    #[test]
    fn test_can_merge_different_length() {
        let bf = BranchFolding::new();
        let a = vec![make_instr(1, vec![])];
        let b = vec![make_instr(1, vec![]), make_instr(2, vec![])];
        assert!(!bf.can_merge(&a, &b));
    }

    #[test]
    fn test_can_merge_empty() {
        let bf = BranchFolding::new();
        let a: Vec<MachineInstr> = vec![];
        let b: Vec<MachineInstr> = vec![];
        assert!(!bf.can_merge(&a, &b));
    }

    #[test]
    fn test_operands_equal() {
        let bf = BranchFolding::new();
        assert!(bf.operands_equal(&MachineOperand::Reg(0), &MachineOperand::Reg(0)));
        assert!(!bf.operands_equal(&MachineOperand::Reg(0), &MachineOperand::Reg(1)));
        assert!(bf.operands_equal(&MachineOperand::Imm(42), &MachineOperand::Imm(42)));
        assert!(!bf.operands_equal(&MachineOperand::Imm(42), &MachineOperand::Imm(99)));
        assert!(bf.operands_equal(
            &MachineOperand::Label("L0".into()),
            &MachineOperand::Label("L0".into())
        ));
        assert!(!bf.operands_equal(
            &MachineOperand::Label("L0".into()),
            &MachineOperand::Label("L1".into())
        ));
    }

    #[test]
    fn test_count_predecessors() {
        let bf = BranchFolding::new();
        let mut mf = MachineFunction::new("test");
        let mut bb0 = make_block("bb0", vec![]);
        bb0.successors.push("bb1".into());
        let bb1 = make_block("bb1", vec![]);
        mf.blocks.push(bb0);
        mf.blocks.push(bb1);
        assert_eq!(bf.count_predecessors(&mf, 1), 1);
        assert_eq!(bf.count_predecessors(&mf, 0), 0);
    }

    #[test]
    fn test_is_unconditional_branch() {
        let bf = BranchFolding::new();
        let br = make_instr(0, vec![MachineOperand::Label("L0".into())]);
        assert!(bf.is_unconditional_branch(&br));
        let non_br = make_instr(1, vec![MachineOperand::Reg(0)]);
        assert!(!bf.is_unconditional_branch(&non_br));
    }

    #[test]
    fn test_get_branch_target() {
        let bf = BranchFolding::new();
        let br = make_instr(0, vec![MachineOperand::Label("target".into())]);
        assert_eq!(bf.get_branch_target(&br), Some("target".to_string()));
        let non_br = make_instr(1, vec![MachineOperand::Reg(0)]);
        assert_eq!(bf.get_branch_target(&non_br), None);
    }

    #[test]
    fn test_redirect_branches() {
        let mut bf = BranchFolding::new();
        let mut mf = MachineFunction::new("test");
        let mut bb0 = make_block(
            "bb0",
            vec![make_instr(0, vec![MachineOperand::Label("bb1".into())])],
        );
        bb0.successors.push("bb1".into());
        mf.blocks.push(bb0);
        mf.blocks.push(make_block("bb1", vec![]));
        mf.blocks.push(make_block("bb2", vec![]));

        bf.redirect_branches(&mut mf, 1, 2);

        // bb0 should now target bb2
        assert_eq!(mf.blocks[0].successors[0], "bb2");
        if let MachineOperand::Label(ref lbl) = mf.blocks[0].instructions[0].operands[0] {
            assert_eq!(lbl, "bb2");
        } else {
            panic!("Expected label operand");
        }
    }

    #[test]
    fn test_merge_blocks() {
        let mut bf = BranchFolding::new();
        let mut mf = MachineFunction::new("test");
        let instr = make_instr(1, vec![MachineOperand::Reg(0)]);
        mf.blocks.push(make_block("bb0", vec![instr.clone()]));
        mf.blocks.push(make_block("bb1", vec![instr.clone()]));
        bf.merge_blocks(&mut mf, 0, 1);
        // bb0 should be empty after merge
        assert!(mf.blocks[0].instructions.is_empty());
    }

    #[test]
    fn test_default() {
        let bf = BranchFolding::default();
        assert_eq!(bf.merges, 0);
    }
}