llvm-native-core-ext 0.1.0

//! LLVM TailDuplication — duplicates blocks to eliminate branches and
//! improve fallthrough, reducing branch penalties and enabling further
//! optimizations.
//! Clean-room behavioral reconstruction.
//!
//! Tail duplication identifies basic blocks that are small enough and
//! have multiple predecessors. By duplicating the tail block into each
//! predecessor, we eliminate unconditional branches and create larger
//! straight-line code regions. This improves I-cache utilization, reduces
//! branch mispredictions, and exposes more opportunities for subsequent
//! optimizations like CSE and instruction scheduling.
//!
//! Algorithm:
//!   1. Scan all blocks to find "tails": blocks with multiple predecessors
//!      that are small enough to be profitably duplicated.
//!   2. For each (predecessor, tail) pair, check duplication criteria:
//!      - The tail block must be small (≤ 8 instructions by default).
//!      - The tail must not contain indirect branches or calls.
//!      - Duplicating should not increase code size beyond a threshold.
//!   3. Duplicate the tail into the predecessor, merging the copied
//!      instructions and fixing up branches.
//!   4. If all predecessors have been updated, remove the original tail.
//!   5. Iterate until convergence or size limit reached.

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

// ============================================================================
// TailDuplication Pass
// ============================================================================

/// TailDuplication — duplicates blocks to eliminate branches and improve
/// fallthrough.
#[derive(Debug, Clone)]
pub struct TailDuplication {
    /// Maximum number of instructions allowed in a tail to be duplicated.
    pub max_tail_size: usize,
    /// Maximum total code size expansion factor (1.0 = no expansion).
    pub max_size_expansion: f64,
    /// Number of tails duplicated.
    pub duplicated: usize,
}

impl TailDuplication {
    /// Create a new TailDuplication pass with default thresholds.
    pub fn new() -> Self {
        Self {
            max_tail_size: 8,
            max_size_expansion: 1.5,
            duplicated: 0,
        }
    }

    /// Run tail duplication on a machine function.
    /// Returns the number of tails duplicated.
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        self.duplicated = 0;
        let original_size = self.count_total_instructions(mf);

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

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

            let tails = self.find_duplicatable_tails(mf);

            for (src, tail) in tails {
                // Check code size expansion
                let current_size = self.count_total_instructions(mf);
                let expansion = current_size as f64 / original_size.max(1) as f64;
                if expansion > self.max_size_expansion {
                    break;
                }

                if self.should_duplicate(src, tail, mf) {
                    self.duplicate_tail(mf, src, tail);
                    self.duplicated += 1;
                    changed = true;
                }
            }
        }

        self.duplicated
    }

    /// Find blocks eligible for tail duplication.
    /// Returns a list of (predecessor_idx, tail_idx) pairs.
    fn find_duplicatable_tails(&self, mf: &MachineFunction) -> Vec<(usize, usize)> {
        let mut candidates = Vec::new();

        for (tail_idx, tail) in mf.blocks.iter().enumerate() {
            // Skip empty blocks
            if tail.instructions.is_empty() {
                continue;
            }

            // Tail must be small enough
            if tail.instructions.len() > self.max_tail_size {
                continue;
            }

            // Tail must have predecessors
            let preds = self.find_predecessors(mf, tail_idx);
            if preds.len() < 2 {
                // Single-predecessor tails: only duplicate if it enables
                // merging (e.g., to expose fallthrough)
                if preds.len() == 1 {
                    let pred_idx = preds[0];
                    // If pred has multiple successors, duplicating the tail
                    // into this pred can eliminate a branch
                    if mf.blocks[pred_idx].successors.len() == 1 && pred_idx + 1 != tail_idx {
                        candidates.push((pred_idx, tail_idx));
                    }
                }
                continue;
            }

            // Multiple predecessors: candidate for duplication
            for &pred_idx in &preds {
                candidates.push((pred_idx, tail_idx));
            }
        }

        candidates
    }

    /// Decide whether to duplicate a tail into a predecessor.
    fn should_duplicate(&self, src: usize, tail: usize, mf: &MachineFunction) -> bool {
        if src >= mf.blocks.len() || tail >= mf.blocks.len() {
            return false;
        }
        if src == tail {
            return false;
        }

        let tail_block = &mf.blocks[tail];
        let src_block = &mf.blocks[src];

        // Don't duplicate if the tail contains indirect branches or calls
        for instr in &tail_block.instructions {
            if self.is_indirect_branch(instr) || self.is_call(instr) {
                return false;
            }
        }

        // Don't duplicate if the predecessor already falls through
        if src + 1 == tail && src_block.successors.len() == 1 && src_block.successors[0] == tail {
            return false;
        }

        // Don't duplicate if the tail has side effects that can't be duplicated
        for instr in &tail_block.instructions {
            if self.has_unsafe_side_effect(instr) {
                return false;
            }
        }

        // Benefit: duplicating eliminates a branch at the end of src
        // Cost: tail instructions are duplicated
        // Heuristic: duplicate if tail ≤ max_tail_size and benefit > cost
        tail_block.instructions.len() <= self.max_tail_size
    }

    /// Duplicate the tail block into the source predecessor.
    fn duplicate_tail(&mut self, mf: &mut MachineFunction, src: usize, tail: usize) {
        // Clone the tail's instructions
        let tail_instrs = mf.blocks[tail].instructions.clone();
        let tail_successors = mf.blocks[tail].successors.clone();

        // Remove the branch targeting the tail from the src block
        // and append the tail's instructions
        let tail_name = mf.blocks[tail].name.clone();
        {
            let src_block = &mut mf.blocks[src];

            // Remove the last instruction if it's a branch to the tail
            if let Some(last) = src_block.instructions.last() {
                if self.is_branch_to(last, &tail_name) {
                    src_block.instructions.pop();
                }
            }

            // Fix up labels in the copied instructions
            for mut instr in tail_instrs {
                // Any label references that pointed to the tail's successors
                // remain valid since they point forward
                self.remap_labels_local(&mut instr);
                src_block.instructions.push(instr);
            }

            // Update src successors to the tail's successors
            src_block.successors = tail_successors;
        }

        // Remove the tail block if it no longer has any predecessors
        let remaining_preds = self.find_predecessors(mf, tail);
        if remaining_preds.is_empty() {
            mf.blocks[tail].instructions.clear();
            mf.blocks[tail].successors.clear();
        }
    }

    /// Fix branches by updating references from old target to new target.
    fn fix_branches(&mut self, mf: &mut MachineFunction, old: usize, new: usize) {
        let old_name = mf.blocks[old].name.clone();
        let new_name = mf.blocks[new].name.clone();

        for block in &mut mf.blocks {
            for succ in &mut block.successors {
                if *succ == old {
                    *succ = new;
                }
            }
            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 total instructions in the function.
    fn count_total_instructions(&self, mf: &MachineFunction) -> usize {
        mf.blocks.iter().map(|b| b.instructions.len()).sum()
    }

    /// Find all predecessor indices for a block.
    fn find_predecessors(&self, mf: &MachineFunction, idx: usize) -> Vec<usize> {
        let mut preds = Vec::new();
        for (i, block) in mf.blocks.iter().enumerate() {
            if block.successors.contains(&idx) {
                preds.push(i);
            }
        }
        preds
    }

    /// Check if an instruction is an indirect branch (can't safely duplicate).
    fn is_indirect_branch(&self, instr: &MachineInstr) -> bool {
        // Indirect branches reference a register for the target
        instr
            .operands
            .iter()
            .any(|op| matches!(op, MachineOperand::Reg(_)))
            && instr
                .operands
                .iter()
                .any(|op| matches!(op, MachineOperand::Label(_)))
    }

    /// Check if an instruction is a call.
    fn is_call(&self, instr: &MachineInstr) -> bool {
        // Calls reference global symbols typically
        instr
            .operands
            .iter()
            .any(|op| matches!(op, MachineOperand::Global(_)))
    }

    /// Check if an instruction has side effects that prevent duplication.
    fn has_unsafe_side_effect(&self, _instr: &MachineInstr) -> bool {
        // For simplicity, assume most instructions are safe to duplicate.
        // In a real backend, check for things like:
        // - Memory barriers / fences
        // - Atomic operations
        // - Exception-causing instructions that rely on precise location
        false
    }

    /// Check if an instruction is a branch to a specific label.
    fn is_branch_to(&self, instr: &MachineInstr, target: &str) -> bool {
        instr.operands.iter().any(|op| {
            if let MachineOperand::Label(lbl) = op {
                lbl == target
            } else {
                false
            }
        })
    }

    /// Remap labels in an instruction to avoid conflicts after duplication.
    fn remap_labels_local(&self, _instr: &mut MachineInstr) {
        // In a full implementation, this would update label references
        // that were local to the tail block. For now, labels to
        // successors are left as-is (they point forward correctly).
    }
}

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

// ============================================================================
// Tail Duplication Cost Model
// ============================================================================

/// DuplicationCost models the cost/benefit analysis for duplicating
/// a tail block into a predecessor.
#[derive(Debug, Clone)]
pub struct DuplicationCost {
    /// Number of instructions that would be duplicated.
    pub instruction_count: usize,
    /// Number of branches eliminated (1 if unconditional, 0 if fallthrough).
    pub branches_eliminated: usize,
    /// Estimated cycles saved per execution.
    pub cycles_saved: f64,
    /// Code size increase in bytes.
    pub size_increase: usize,
    /// Whether duplication is profitable.
    pub is_profitable: bool,
    /// The predecessor block index.
    pub predecessor: usize,
    /// The tail block index.
    pub tail: usize,
}

impl DuplicationCost {
    /// Create a new cost estimate.
    pub fn new(pred: usize, tail: usize) -> Self {
        Self {
            instruction_count: 0,
            branches_eliminated: 0,
            cycles_saved: 0.0,
            size_increase: 0,
            is_profitable: false,
            predecessor: pred,
            tail,
        }
    }

    /// Evaluate profitability given block data.
    pub fn evaluate(&mut self, tail_block: &MachineBasicBlock, pred_has_branch_to_tail: bool) {
        self.instruction_count = tail_block.instructions.len();

        // Branches eliminated: if the predecessor branches to the tail,
        // duplication eliminates that branch (saving ~1-2 cycles).
        if pred_has_branch_to_tail {
            self.branches_eliminated = 1;
        }

        // Cycles saved: branch penalty (~2 cycles) + fallthrough benefit (~1 cycle)
        self.cycles_saved = if pred_has_branch_to_tail {
            2.0 + 1.0
        } else {
            1.0 // Only fallthrough benefit
        };

        // Size increase: duplicated instructions minus removed branch
        let branch_size = if pred_has_branch_to_tail { 4 } else { 0 };
        self.size_increase = self.instruction_count * 4 - branch_size;

        // Profitability heuristic:
        // - Must save at least 1 cycle
        // - Size increase must be acceptable (≤ 2x the branch saved)
        self.is_profitable =
            self.cycles_saved >= 1.0 && self.instruction_count <= 8 && self.size_increase <= 64;
    }
}

// ============================================================================
// Branch Prediction Optimization via Tail Duplication
// ============================================================================

/// BranchPredictionDuplicator uses profile information (or static
/// heuristics) to selectively duplicate tails only for hot predecessors,
/// avoiding code bloat on cold paths.
#[derive(Debug, Clone)]
pub struct BranchPredictionDuplicator {
    /// Base tail duplication pass.
    pub base: TailDuplication,
    /// Estimated block execution frequencies (block_idx -> count).
    pub frequencies: HashMap<usize, f64>,
    /// Minimum frequency ratio to consider a predecessor "hot".
    pub hot_threshold: f64,
    /// Number of hot-path duplications performed.
    pub hot_duplications: usize,
}

impl BranchPredictionDuplicator {
    /// Create a new BPDuplicator.
    pub fn new() -> Self {
        Self {
            base: TailDuplication::new(),
            frequencies: HashMap::new(),
            hot_threshold: 0.1,
            hot_duplications: 0,
        }
    }

    /// Run profile-guided tail duplication.
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        self.hot_duplications = 0;

        // Estimate frequencies (static heuristic: earlier blocks = hotter)
        self.estimate_frequencies(mf);

        // Find duplicatable tails (returns Vec<(pred_idx, tail_idx)>)
        let tail_pairs = self.base.find_duplicatable_tails(mf);

        // Group by tail_idx
        let mut tail_to_preds: HashMap<usize, Vec<usize>> = HashMap::new();
        for &(pred_idx, tail_idx) in &tail_pairs {
            tail_to_preds.entry(tail_idx).or_default().push(pred_idx);
        }

        let max_freq = self.frequencies.values().cloned().fold(0.0f64, f64::max);

        for (tail_idx, pred_indices) in &tail_to_preds {
            let tail_block = mf.blocks[*tail_idx].clone();

            // Only duplicate into hot predecessors
            for &pred_idx in pred_indices {
                let pred_freq = self.frequencies.get(&pred_idx).copied().unwrap_or(0.0);

                let ratio = if max_freq > 0.0 {
                    pred_freq / max_freq
                } else {
                    0.0
                };

                if ratio >= self.hot_threshold {
                    // Hot predecessor — duplicate tail into it
                    let pred_has_branch = self.pred_has_branch_to(mf, pred_idx, *tail_idx);

                    let mut cost = DuplicationCost::new(pred_idx, *tail_idx);
                    cost.evaluate(&tail_block, pred_has_branch);

                    if cost.is_profitable {
                        // Perform duplication
                        self.base.duplicate_tail(mf, pred_idx, *tail_idx);
                        self.base.duplicated += 1;
                        self.hot_duplications += 1;
                    }
                }
            }
        }

        self.base.duplicated
    }

    /// Estimate block execution frequencies using static heuristic:
    /// blocks earlier in the layout are more likely to execute.
    fn estimate_frequencies(&mut self, mf: &MachineFunction) {
        self.frequencies.clear();
        let n = mf.blocks.len();

        if n == 0 {
            return;
        }

        // Entry block has frequency 1.0
        self.frequencies.insert(0, 1.0);

        // Propagate forward: each successor gets a fraction of predecessor's freq
        for i in 0..n {
            let pred_freq = self.frequencies.get(&i).copied().unwrap_or(0.0);
            let succs = &mf.blocks[i].successors;
            if succs.is_empty() {
                continue;
            }

            let share = pred_freq / succs.len() as f64;
            for &succ_idx in succs.iter() {
                *self.frequencies.entry(succ_idx).or_insert(0.0) += share;
            }
        }
    }

    /// Check if a predecessor has an explicit branch to the tail block.
    fn pred_has_branch_to(&self, mf: &MachineFunction, pred_idx: usize, tail_idx: usize) -> bool {
        let pred = &mf.blocks[pred_idx];
        pred.successors.contains(&tail_idx)
    }
}

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

// ============================================================================
// Fallthrough-Enabling Tail Duplication
// ============================================================================

/// FallthroughDuplicator specifically duplicates tails to create
/// fallthrough paths, which are optimal for branch prediction and
/// I-cache utilization.
#[derive(Debug, Clone)]
pub struct FallthroughDuplicator {
    /// Base duplication pass.
    pub base: TailDuplication,
    /// Blocks that have been modified.
    pub modified: HashSet<usize>,
}

impl FallthroughDuplicator {
    /// Create a new fallthrough duplicator.
    pub fn new() -> Self {
        Self {
            base: TailDuplication::new(),
            modified: HashSet::new(),
        }
    }

    /// Run fallthrough optimization via tail duplication.
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        self.modified.clear();
        self.base.duplicated = 0;

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

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

            for block_idx in 0..mf.blocks.len().saturating_sub(1) {
                let next_idx = block_idx + 1;

                // Check if the current block has an unconditional branch
                // to a block that is NOT the next block in layout.
                let succs = mf.blocks[block_idx].successors.clone();

                if succs.len() == 1 {
                    let target_idx = succs[0];

                    if target_idx != next_idx {
                        // Branch goes somewhere else — check if the target
                        // can be duplicated into the current block to
                        // create fallthrough.
                        let target_block = &mf.blocks[target_idx];

                        // If the target is small and has few predecessors,
                        // duplicating it here creates fallthrough
                        if target_block.instructions.len() <= self.base.max_tail_size
                            && !self.modified.contains(&target_idx)
                        {
                            self.base.duplicate_tail(mf, block_idx, target_idx);
                            self.modified.insert(block_idx);
                            changed = true;
                        }
                    }
                }

                // Also check: if block ends with conditional branch and
                // the fallthrough target is not the next block, we might
                // duplicate the next block's tail to fix this.
                if succs.len() == 2 {
                    let fallthrough_target = succs[1];

                    if fallthrough_target != next_idx {
                        // The natural fallthrough doesn't go to the next
                        // layout block. Try to duplicate the next block
                        // into the fallthrough path.
                        if mf.blocks[next_idx].instructions.len() <= self.base.max_tail_size {
                            // Duplicate next block into the fallthrough path
                            // (simplified: just mark as opportunity)
                            self.base.duplicated += 1;
                            changed = true;
                        }
                    }
                }
            }
        }

        self.base.duplicated
    }
}

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

// ============================================================================
// TailDuplicator — Unified Tail Duplication with Cost Model
// ============================================================================

/// TailDuplicator combines all tail duplication strategies into a
/// single unified pass with cost-based decision making.
pub struct TailDuplicator {
    /// Base duplication pass.
    pub base: TailDuplication,
    /// Branch-prediction-aware duplication.
    pub bp_duplicator: BranchPredictionDuplicator,
    /// Fallthrough-enabling duplication.
    pub fallthrough_duplicator: FallthroughDuplicator,
    /// Accumulated cost records.
    pub costs: Vec<DuplicationCost>,
    /// Total instructions duplicated.
    pub total_instructions_duplicated: usize,
}

impl TailDuplicator {
    /// Create a new unified tail duplicator.
    pub fn new() -> Self {
        Self {
            base: TailDuplication::new(),
            bp_duplicator: BranchPredictionDuplicator::new(),
            fallthrough_duplicator: FallthroughDuplicator::new(),
            costs: Vec::new(),
            total_instructions_duplicated: 0,
        }
    }

    /// Run full tail duplication with all strategies.
    pub fn run_on_function(&mut self, mf: &mut MachineFunction) -> usize {
        self.costs.clear();
        self.total_instructions_duplicated = 0;

        // Phase 1: Profile-guided duplication (hot paths first)
        let bp_count = self.bp_duplicator.run_on_function(mf);

        // Phase 2: Fallthrough optimization
        let ft_count = self.fallthrough_duplicator.run_on_function(mf);

        // Phase 3: Baseline duplication for remaining opportunities
        let base_count = self.base.run_on_function(mf);

        self.total_instructions_duplicated = bp_count + ft_count + base_count;

        self.total_instructions_duplicated
    }

    /// Evaluate the cost of duplicating a tail into a predecessor.
    pub fn evaluate_duplication(
        &mut self,
        mf: &MachineFunction,
        pred_idx: usize,
        tail_idx: usize,
    ) -> DuplicationCost {
        let tail_block = &mf.blocks[tail_idx];
        let pred_has_branch = mf.blocks[pred_idx].successors.contains(&tail_idx);

        let mut cost = DuplicationCost::new(pred_idx, tail_idx);
        cost.evaluate(&tail_block, pred_has_branch);
        self.costs.push(cost.clone());
        cost
    }

    /// Print duplication statistics.
    pub fn print_stats(&self) {
        eprintln!(
            "TailDuplicator: {} total instructions duplicated",
            self.total_instructions_duplicated
        );
        eprintln!(
            "  Base: {}, BPDuplicator: {}, Fallthrough: {}",
            self.base.duplicated,
            self.bp_duplicator.hot_duplications,
            self.fallthrough_duplicator.base.duplicated
        );
        let profitable = self.costs.iter().filter(|c| c.is_profitable).count();
        eprintln!("  Profitable decisions: {}", profitable);
    }
}

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

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

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

    fn make_block(name: &str, instrs: Vec<MachineInstr>, succs: Vec<&str>) -> MachineBasicBlock {
        // In tests, we map string successor names to numeric IDs.
        let id = match name {
            "bb0" => 0,
            "bb1" => 1,
            "bb2" => 2,
            "bb3" => 3,
            "entry" => 0,
            _ => 0,
        };
        let succ_ids: Vec<usize> = succs
            .iter()
            .map(|s| match *s {
                "bb0" => 0,
                "bb1" => 1,
                "bb2" => 2,
                "bb3" => 3,
                "entry" => 0,
                _ => 0,
            })
            .collect();
        MachineBasicBlock {
            id,
            name: name.to_string(),
            instructions: instrs,
            successors: succ_ids,
            predecessors: Vec::new(),
            is_entry: false,
        }
    }

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

    fn make_br(target: &str) -> MachineInstr {
        make_instr(0, vec![MachineOperand::Label(target.to_string())])
    }

    #[test]
    fn test_new() {
        let td = TailDuplication::new();
        assert_eq!(td.max_tail_size, 8);
        assert!((td.max_size_expansion - 1.5).abs() < 0.001);
        assert_eq!(td.duplicated, 0);
    }

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

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

    #[test]
    fn test_find_duplicatable_tails_no_multi_pred() {
        let td = TailDuplication::new();
        let mut mf = MachineFunction::new("test");
        mf.blocks
            .push(make_block("bb0", vec![make_br("bb1")], vec!["bb1"]));
        mf.blocks.push(make_block(
            "bb1",
            vec![make_instr(1, vec![MachineOperand::Reg(0)])],
            vec![],
        ));
        let tails = td.find_duplicatable_tails(&mf);
        // bb1 has 1 predecessor; tail duplication only if it eliminates a branch
        // and src=bb0 already has only 1 successor pointing to bb1
        assert!(tails.is_empty());
    }

    #[test]
    fn test_find_duplicatable_tails_multi_pred() {
        let td = TailDuplication::new();
        let mut mf = MachineFunction::new("test");
        mf.blocks
            .push(make_block("bb0", vec![make_br("bb2")], vec!["bb2"]));
        mf.blocks
            .push(make_block("bb1", vec![make_br("bb2")], vec!["bb2"]));
        mf.blocks.push(make_block(
            "bb2",
            vec![make_instr(1, vec![MachineOperand::Reg(0)])],
            vec![],
        ));
        let tails = td.find_duplicatable_tails(&mf);
        // bb2 has 2 predecessors → should find 2 candidates
        assert_eq!(tails.len(), 2);
    }

    #[test]
    fn test_find_duplicatable_tails_too_large() {
        let mut td = TailDuplication::new();
        td.max_tail_size = 2;
        let mut mf = MachineFunction::new("test");
        mf.blocks
            .push(make_block("bb0", vec![make_br("bb2")], vec!["bb2"]));
        mf.blocks
            .push(make_block("bb1", vec![make_br("bb2")], vec!["bb2"]));
        mf.blocks.push(make_block(
            "bb2",
            vec![
                make_instr(1, vec![MachineOperand::Reg(0)]),
                make_instr(2, vec![MachineOperand::Reg(1)]),
                make_instr(3, vec![MachineOperand::Reg(2)]),
            ],
            vec![],
        ));
        let tails = td.find_duplicatable_tails(&mf);
        // bb2 has 3 instructions > max_tail_size=2
        assert!(tails.is_empty());
    }

    #[test]
    fn test_should_duplicate() {
        let td = TailDuplication::new();
        let mut mf = MachineFunction::new("test");
        mf.blocks
            .push(make_block("bb0", vec![make_br("bb2")], vec!["bb2"]));
        mf.blocks
            .push(make_block("bb1", vec![make_br("bb2")], vec!["bb2"]));
        mf.blocks.push(make_block(
            "bb2",
            vec![make_instr(1, vec![MachineOperand::Reg(0)])],
            vec!["bb3"],
        ));
        mf.blocks.push(make_block("bb3", vec![], vec![]));
        // bb2 has 2 predecessors; bb0 is not adjacent to bb2, good candidate
        assert!(td.should_duplicate(0, 2, &mf));
        assert!(!td.should_duplicate(0, 0, &mf));
        assert!(!td.should_duplicate(100, 200, &mf));
    }

    #[test]
    fn test_should_not_duplicate_invalid() {
        let td = TailDuplication::new();
        let mf = MachineFunction::new("test");
        assert!(!td.should_duplicate(0, 0, &mf));
        assert!(!td.should_duplicate(100, 200, &mf));
    }

    #[test]
    fn test_duplicate_tail() {
        let mut td = TailDuplication::new();
        let mut mf = MachineFunction::new("test");
        mf.blocks
            .push(make_block("bb0", vec![make_br("bb1")], vec!["bb1"]));
        mf.blocks.push(make_block(
            "bb1",
            vec![make_instr(1, vec![MachineOperand::Reg(0)])],
            vec!["bb2"],
        ));
        mf.blocks.push(make_block("bb2", vec![], vec![]));

        td.duplicate_tail(&mut mf, 0, 1);

        // bb0 should now contain the tail's instructions
        assert!(mf.blocks[0].instructions.len() >= 1);
        // bb0's successors should be bb2 (tail's successor)
        assert_eq!(mf.blocks[0].successors, vec![2usize]);
        // bb1 should be cleared
        assert!(mf.blocks[1].instructions.is_empty());
    }

    #[test]
    fn test_fix_branches() {
        let mut td = TailDuplication::new();
        let mut mf = MachineFunction::new("test");
        mf.blocks
            .push(make_block("bb0", vec![make_br("bb1")], vec!["bb1"]));
        mf.blocks.push(make_block("bb1", vec![], vec![]));
        mf.blocks.push(make_block("bb2", vec![], vec![]));

        td.fix_branches(&mut mf, 1, 2);

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

    #[test]
    fn test_count_total_instructions() {
        let td = TailDuplication::new();
        let mut mf = MachineFunction::new("test");
        mf.blocks.push(make_block(
            "bb0",
            vec![make_instr(1, vec![]), make_instr(2, vec![])],
            vec![],
        ));
        mf.blocks
            .push(make_block("bb1", vec![make_instr(3, vec![])], vec![]));
        assert_eq!(td.count_total_instructions(&mf), 3);
    }

    #[test]
    fn test_find_predecessors() {
        let td = TailDuplication::new();
        let mut mf = MachineFunction::new("test");
        mf.blocks.push(make_block("bb0", vec![], vec!["bb2"]));
        mf.blocks.push(make_block("bb1", vec![], vec!["bb2"]));
        mf.blocks.push(make_block("bb2", vec![], vec![]));
        let preds = td.find_predecessors(&mf, 2);
        assert_eq!(preds.len(), 2);
        assert!(preds.contains(&0));
        assert!(preds.contains(&1));
    }

    #[test]
    fn test_is_indirect_branch() {
        let td = TailDuplication::new();
        let ind_br = make_instr(
            0,
            vec![MachineOperand::Reg(0), MachineOperand::Label("L0".into())],
        );
        assert!(td.is_indirect_branch(&ind_br));
        let dir_br = make_br("L0");
        assert!(!td.is_indirect_branch(&dir_br));
    }

    #[test]
    fn test_is_call() {
        let td = TailDuplication::new();
        let call = make_instr(0, vec![MachineOperand::Global("func".into())]);
        assert!(td.is_call(&call));
        let not_call = make_instr(1, vec![MachineOperand::Reg(0)]);
        assert!(!td.is_call(&not_call));
    }

    #[test]
    fn test_is_branch_to() {
        let td = TailDuplication::new();
        let br = make_br("target");
        assert!(td.is_branch_to(&br, "target"));
        assert!(!td.is_branch_to(&br, "other"));
    }

    #[test]
    fn test_default() {
        let td = TailDuplication::default();
        assert_eq!(td.max_tail_size, 8);
    }
}