llvm-native-core-ext 0.1.0

//! Jump Threading — forward CFG simplification that threads jumps through
//! blocks with constant conditions. Clean-room reimplementation.
//!
//! @llvm_behavior: Jump Threading looks at blocks with conditional branches
//! and determines if any predecessor always takes a specific branch edge.
//! When this happens, the predecessor's terminator can be redirected to
//! skip the intermediate block entirely, "threading" the jump.
//!
//! Key transformations:
//! 1. If pred unconditionally branches to a block with `br i1 %cond, A, B`
//!    and %cond is known to be true/false from the pred, redirect pred→target.
//! 2. br i1 %cond, %A, %A → br %A (same target, fold to unconditional).
//! 3. Eliminate blocks that become unreachable after threading.

use llvm_native_core::opcode::Opcode;
use llvm_native_core::value::ValueRef;
use llvm_native_core::SubclassKind;
use std::collections::{HashMap, HashSet, VecDeque};

/// Jump Threading pass. Threads jumps through blocks with constant conditions.
#[derive(Debug)]
pub struct JumpThreadingPass {
    /// Number of jumps threaded during this run.
    pub threaded: usize,
}

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

impl JumpThreadingPass {
    pub fn new() -> Self {
        Self { threaded: 0 }
    }

    /// Run jump threading on a function. Returns the number of jumps threaded.
    pub fn run_on_function(&mut self, func: &ValueRef) -> usize {
        self.threaded = 0;
        let mut changed = true;
        let mut total_threaded = 0usize;

        while changed {
            changed = false;

            // Eliminate dead blocks first
            let dead_removed = self.eliminate_dead_blocks(func);
            if dead_removed > 0 {
                changed = true;
            }

            // Simplify unconditional branches in each block
            {
                let f = func.borrow();
                for op in &f.operands {
                    let simplified = self.try_to_simplify(op);
                    if simplified > 0 {
                        changed = true;
                        total_threaded += simplified;
                    }
                }
            }

            // Find and thread jump opportunities
            let opportunities = self.find_threadable_blocks(func);
            for (pred, block, target) in opportunities {
                if self.thread_jump(&pred, &block, &target) {
                    self.threaded += 1;
                    total_threaded += 1;
                    changed = true;
                }
            }
        }

        self.threaded = total_threaded;
        total_threaded
    }

    /// Find blocks where jump threading may be possible.
    /// Returns Vec of (predecessor_block, block_with_conditional_br, target_block).
    pub fn find_threadable_blocks(&self, func: &ValueRef) -> Vec<(ValueRef, ValueRef, ValueRef)> {
        let mut results = Vec::new();
        let f = func.borrow();

        // Build predecessor map
        let pred_map = build_predecessor_map(func);

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

            // Get the terminator instruction (last instruction in the block)
            let terminator = get_terminator(op);
            if terminator.is_none() {
                continue;
            }
            let term = terminator.unwrap();
            let term_borrow = term.borrow();

            // Only interested in conditional branches: br i1 %cond, label %A, label %B
            if term_borrow.get_opcode() != Some(Opcode::Br) || term_borrow.operands.len() != 3 {
                continue;
            }

            let cond_val = &term_borrow.operands[0];
            let true_target = &term_borrow.operands[1];
            let false_target = &term_borrow.operands[2];

            // Check each predecessor
            if let Some(preds) = pred_map.get(&op.borrow().vid) {
                for pred in preds {
                    // Check if the condition is known from this predecessor
                    if let Some(known_val) = self.get_known_condition(pred, cond_val) {
                        let target = if known_val {
                            true_target.clone()
                        } else {
                            false_target.clone()
                        };
                        results.push((pred.clone(), op.clone(), target));
                    }
                }
            }
        }

        results
    }

    /// Determine if a condition value is known true/false from the given predecessor.
    /// Looks at the predecessor's terminator and the values flowing into the condition.
    pub fn get_known_condition(&self, pred: &ValueRef, cond_inst: &ValueRef) -> Option<bool> {
        let pred_bb = pred.borrow();

        // Case 1: predecessor unconditionally branches to the block containing cond_inst
        // and we can evaluate the condition based on known values
        let pred_term = get_terminator(pred);
        if let Some(term) = pred_term {
            let term_borrow = term.borrow();

            // If pred ends with unconditional branch, check what we know
            if term_borrow.get_opcode() == Some(Opcode::Br) && term_borrow.operands.len() == 1 {
                // The predecessor simply falls through — check the condition
                return self.is_constant_condition_at_block(pred, cond_inst);
            }

            // If pred ends with conditional branch, check the taken path
            if term_borrow.get_opcode() == Some(Opcode::Br) && term_borrow.operands.len() == 3 {
                // Check if the condition is known based on the branch direction
                let pred_cond = &term_borrow.operands[0];
                if let Some(val) = self.is_constant_condition_at_block(pred, cond_inst) {
                    return Some(val);
                }
            }
        }

        // Case 2: Look at PHI nodes feeding the condition
        let cond_borrow = cond_inst.borrow();
        if cond_borrow.get_opcode() == Some(Opcode::Phi) {
            // PHI: for each (value, label) pair, if label == pred, check value
            for chunk in cond_borrow.operands.chunks(2) {
                if chunk.len() == 2 {
                    let val = &chunk[0];
                    let label = &chunk[1];
                    if std::ptr::addr_eq(Rc::as_ptr(label), Rc::as_ptr(pred)) {
                        // The incoming value might be a constant
                        let val_borrow = val.borrow();
                        if val_borrow.is_constant() {
                            return parse_bool_constant(val);
                        }
                        // It might be an icmp with constant result
                        if val_borrow.get_opcode() == Some(Opcode::ICmp) {
                            return evaluate_icmp_constant(val);
                        }
                    }
                }
            }
        }

        // Case 3: cond_inst is an icmp, evaluate it
        if cond_borrow.get_opcode() == Some(Opcode::ICmp) {
            return evaluate_icmp_constant(cond_inst);
        }

        None
    }

    /// Check if a condition value is constant-true or constant-false in the context
    /// of a given block.
    pub fn is_constant_condition_at_block(
        &self,
        block: &ValueRef,
        predicate: &ValueRef,
    ) -> Option<bool> {
        let pred_borrow = predicate.borrow();

        // Direct constant: i1 true / i1 false
        if pred_borrow.is_constant() && pred_borrow.ty.is_integer() {
            return parse_bool_constant(predicate);
        }

        // ICmp with constant operands
        if pred_borrow.get_opcode() == Some(Opcode::ICmp) {
            return evaluate_icmp_constant(predicate);
        }

        // Look through the block for the definition of predicate
        let bb = block.borrow();
        for inst_val in &bb.operands {
            let inst = inst_val.borrow();
            if !inst.is_instruction() {
                continue;
            }

            // If this instruction defines the predicate
            if std::ptr::addr_eq(Rc::as_ptr(inst_val), Rc::as_ptr(predicate)) {
                // ICmp NE 0 pattern
                if inst.get_opcode() == Some(Opcode::ICmp) {
                    let result = evaluate_icmp_constant(inst_val);
                    if result.is_some() {
                        return result;
                    }
                }

                // Check if it's a phi that we can resolve
                if inst.get_opcode() == Some(Opcode::Phi) {
                    for chunk in inst.operands.chunks(2) {
                        if chunk.len() == 2 {
                            let val = &chunk[0];
                            let label = &chunk[1];
                            let val_borrow = val.borrow();
                            // If the phi incoming value is a known constant and
                            // the label matches a predecessor we understand
                            if val_borrow.is_constant() && val_borrow.ty.is_integer() {
                                // For a phi where ALL incoming values are non-zero,
                                // icmp ne 0 would be true
                                // Simplified: just check if this specific incoming is non-zero
                                let int_val = parse_integer_constant(val);
                                if let Some(v) = int_val {
                                    if v != 0 {
                                        return Some(true);
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        None
    }

    /// Thread a jump: redirect predecessor's terminator to go directly to target,
    /// skipping the intermediate block.
    pub fn thread_jump(&self, pred: &ValueRef, _block: &ValueRef, target: &ValueRef) -> bool {
        let pred_term = get_terminator(pred);
        if pred_term.is_none() {
            return false;
        }

        let term = pred_term.unwrap();
        let mut term_borrow = term.borrow_mut();

        match term_borrow.get_opcode() {
            Some(Opcode::Br) => {
                match term_borrow.operands.len() {
                    1 => {
                        // Unconditional branch: redirect to target
                        term_borrow.operands[0] = target.clone();
                        // Update successors
                        term_borrow.successors = vec![target.clone()];
                        return true;
                    }
                    3 => {
                        // Conditional branch: check if both targets are the same intermediate
                        // or if we can redirect one of them
                        let true_tgt = term_borrow.operands[1].clone();
                        let false_tgt = term_borrow.operands[2].clone();

                        if std::ptr::addr_eq(Rc::as_ptr(&true_tgt), Rc::as_ptr(target)) {
                            // Already branching to target on true — no change needed
                            // But we can fold if both targets point to same place
                            return self.simplify_unconditional_branch(pred);
                        }
                        if std::ptr::addr_eq(Rc::as_ptr(&false_tgt), Rc::as_ptr(target)) {
                            return self.simplify_unconditional_branch(pred);
                        }

                        // Try to redirect: check if the intermediate block can be bypassed
                        // Replace the operand pointing to _block with target
                        if std::ptr::addr_eq(
                            Rc::as_ptr(&term_borrow.operands[1]),
                            Rc::as_ptr(_block),
                        ) {
                            term_borrow.operands[1] = target.clone();
                            term_borrow.successors =
                                vec![target.clone(), term_borrow.operands[2].clone()];
                            return true;
                        }
                        if std::ptr::addr_eq(
                            Rc::as_ptr(&term_borrow.operands[2]),
                            Rc::as_ptr(_block),
                        ) {
                            term_borrow.operands[2] = target.clone();
                            term_borrow.successors =
                                vec![term_borrow.operands[1].clone(), target.clone()];
                            return true;
                        }
                    }
                    _ => {}
                }
            }
            _ => {}
        }

        false
    }

    /// Simplify an unconditional branch: if a block ends with `br label %target`,
    /// and the target block starts with an unconditional branch, thread through.
    pub fn simplify_unconditional_branch(&self, block: &ValueRef) -> bool {
        let term = get_terminator(block);
        if term.is_none() {
            return false;
        }

        let term_val = term.unwrap();
        let is_simple_br;
        let dest;
        {
            let term_borrow = term_val.borrow();
            is_simple_br =
                term_borrow.get_opcode() == Some(Opcode::Br) && term_borrow.operands.len() == 1;
            dest = if is_simple_br {
                Some(term_borrow.operands[0].clone())
            } else {
                None
            };
        } // term_borrow dropped here

        if !is_simple_br || dest.is_none() {
            return false;
        }
        let dest = dest.unwrap();

        // Check if destination is a simple pass-through block
        let dest_is_pass_through;
        let final_dest;
        {
            let dest_bb = dest.borrow();
            if !dest_bb.is_basic_block() {
                return false;
            }
            if dest_bb.operands.len() == 1 {
                let dest_inst = dest_bb.operands[0].borrow();
                if dest_inst.get_opcode() == Some(Opcode::Br) && dest_inst.operands.len() == 1 {
                    final_dest = Some(dest_inst.operands[0].clone());
                    dest_is_pass_through = true;
                } else {
                    final_dest = None;
                    dest_is_pass_through = false;
                }
            } else {
                final_dest = None;
                dest_is_pass_through = false;
            }
        } // dest_bb, dest_inst dropped

        if dest_is_pass_through {
            let final_dest = final_dest.unwrap();
            let mut term_mut = term_val.borrow_mut();
            term_mut.operands[0] = final_dest.clone();
            term_mut.successors = vec![final_dest];
            return true;
        }

        false
    }

    /// Try to simplify the given block: fold conditional branches with known
    /// conditions, merge same-target branches, eliminate dead blocks.
    pub fn try_to_simplify(&self, block: &ValueRef) -> usize {
        let mut count = 0usize;

        // Check for foldable conditional branches
        if let Some(term) = get_terminator(block) {
            let term_borrow = term.borrow();
            if term_borrow.get_opcode() == Some(Opcode::Br) && term_borrow.operands.len() == 3 {
                // Case: br i1 true, label %A, label %B → br label %A
                let cond_is_constant = {
                    let cond_val = &term_borrow.operands[0];
                    let cond_borrow = cond_val.borrow();
                    cond_borrow.is_constant() && cond_borrow.ty.is_integer()
                };

                if cond_is_constant {
                    let cond_val = &term_borrow.operands[0];
                    if let Some(val) = parse_bool_constant(cond_val) {
                        drop(term_borrow);
                        if self.fold_conditional_branch(block, val) {
                            count += 1;
                            // After folding, skip further checks since the terminator changed
                            return count + self.simplify_unconditional_branch(block) as usize;
                        }
                        return count;
                    }
                }

                // Case: br i1 %cond, label %A, label %A → br label %A
                let same_target = {
                    let t1 = Rc::as_ptr(&term_borrow.operands[1]);
                    let t2 = Rc::as_ptr(&term_borrow.operands[2]);
                    t1 == t2
                };

                if same_target {
                    drop(term_borrow);
                    let mut term_mut = term.borrow_mut();
                    let target = term_mut.operands[1].clone();
                    // Convert to unconditional branch
                    term_mut.operands = vec![target.clone()];
                    term_mut.successors = vec![target];
                    count += 1;
                }
            }
        }

        // Try unconditional branch simplification
        if self.simplify_unconditional_branch(block) {
            count += 1;
        }

        count
    }

    /// Fold a conditional branch: if cond is true, replace `br i1 true, A, B`
    /// with `br A`; if false, with `br B`.
    pub fn fold_conditional_branch(&self, block: &ValueRef, cond: bool) -> bool {
        let term = get_terminator(block);
        if term.is_none() {
            return false;
        }

        let term_val = term.unwrap();
        let mut term_borrow = term_val.borrow_mut();

        if term_borrow.get_opcode() != Some(Opcode::Br) || term_borrow.operands.len() != 3 {
            return false;
        }

        let target = if cond {
            term_borrow.operands[1].clone()
        } else {
            term_borrow.operands[2].clone()
        };

        // Convert to unconditional branch
        term_borrow.operands = vec![target.clone()];
        term_borrow.successors = vec![target];
        true
    }

    /// Eliminate dead (unreachable) blocks from the function.
    /// A block is unreachable if it has no predecessors (except the entry block).
    pub fn eliminate_dead_blocks(&self, func: &ValueRef) -> usize {
        let mut f = func.borrow_mut();
        let mut removed = 0usize;

        if f.operands.is_empty() {
            return 0;
        }

        // Build a set of reachable blocks via DFS from the entry
        let mut visited: HashSet<u64> = HashSet::new();
        let mut queue: VecDeque<ValueRef> = VecDeque::new();

        // Start from entry block (first operand)
        if let Some(entry) = f.operands.first() {
            queue.push_back(entry.clone());
            visited.insert(entry.borrow().vid);
        }

        // DFS to find all reachable blocks
        while let Some(block) = queue.pop_front() {
            let bb = block.borrow();
            if let Some(term) = get_terminator(&block) {
                let term_borrow = term.borrow();
                for succ in &term_borrow.successors {
                    if !visited.contains(&succ.borrow().vid) {
                        visited.insert(succ.borrow().vid);
                        queue.push_back(succ.clone());
                    }
                }
                // Also check operands (for branches where successors aren't set)
                if term_borrow.get_opcode() == Some(Opcode::Br) {
                    match term_borrow.operands.len() {
                        1 => {
                            let tgt = &term_borrow.operands[0];
                            if !visited.contains(&tgt.borrow().vid) {
                                visited.insert(tgt.borrow().vid);
                                queue.push_back(tgt.clone());
                            }
                        }
                        3 => {
                            for i in 1..=2 {
                                let tgt = &term_borrow.operands[i];
                                if !visited.contains(&tgt.borrow().vid) {
                                    visited.insert(tgt.borrow().vid);
                                    queue.push_back(tgt.clone());
                                }
                            }
                        }
                        _ => {}
                    }
                }
            }
        }

        // Remove blocks not in visited set (skip entry block which is always at index 0)
        let mut indices_to_remove = Vec::new();
        for (i, op) in f.operands.iter().enumerate() {
            if i == 0 {
                continue; // Never remove the entry block
            }
            if !visited.contains(&op.borrow().vid) {
                indices_to_remove.push(i);
            }
        }

        for i in indices_to_remove.into_iter().rev() {
            f.operands.remove(i);
            removed += 1;
        }

        removed
    }
}

// === Helper Functions ===

/// Get the terminator instruction of a basic block (the last instruction).
fn get_terminator(block: &ValueRef) -> Option<ValueRef> {
    let bb = block.borrow();
    if !bb.is_basic_block() {
        return None;
    }
    // The terminator is the last instruction in the block
    bb.operands.last().cloned()
}

/// Build a map from block vid to list of predecessor blocks.
fn build_predecessor_map(func: &ValueRef) -> HashMap<u64, Vec<ValueRef>> {
    let mut map: HashMap<u64, Vec<ValueRef>> = HashMap::new();
    let f = func.borrow();

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

        if let Some(term) = get_terminator(op) {
            let term_borrow = term.borrow();
            let targets = match term_borrow.get_opcode() {
                Some(Opcode::Br) => match term_borrow.operands.len() {
                    1 => vec![term_borrow.operands[0].clone()],
                    3 => vec![
                        term_borrow.operands[1].clone(),
                        term_borrow.operands[2].clone(),
                    ],
                    _ => vec![],
                },
                Some(Opcode::Switch) => {
                    // Switch: operand[0] is value, operand[1] is default, rest are pairs
                    let mut tgts = vec![term_borrow.operands[1].clone()];
                    for i in (2..term_borrow.operands.len()).step_by(2) {
                        if i + 1 < term_borrow.operands.len() {
                            tgts.push(term_borrow.operands[i + 1].clone());
                        }
                    }
                    tgts
                }
                _ => {
                    // Use successors field
                    term_borrow.successors.clone()
                }
            };

            for target in &targets {
                let target_vid = target.borrow().vid;
                map.entry(target_vid)
                    .or_insert_with(Vec::new)
                    .push(op.clone());
            }
        }
    }

    map
}

/// Parse a constant value as a boolean (true for non-zero, false for zero).
fn parse_bool_constant(val: &ValueRef) -> Option<bool> {
    let v = val.borrow();
    if !v.is_constant() {
        return None;
    }
    // Try parsing the name as an integer
    if let Ok(i) = v.name.parse::<i64>() {
        return Some(i != 0);
    }
    // Check for "true"/"false" strings
    match v.name.as_str() {
        "true" => Some(true),
        "false" => Some(false),
        "1" => Some(true),
        "0" => Some(false),
        _ => {
            // For float constants, non-zero is true
            if let Ok(f) = v.name.parse::<f64>() {
                return Some(f != 0.0);
            }
            None
        }
    }
}

/// Parse a constant value as an integer (i64).
fn parse_integer_constant(val: &ValueRef) -> Option<i64> {
    let v = val.borrow();
    if !v.is_constant() {
        return None;
    }
    v.name.parse::<i64>().ok()
}

/// Evaluate an icmp instruction with constant operands.
fn evaluate_icmp_constant(inst_val: &ValueRef) -> Option<bool> {
    let inst = inst_val.borrow();
    if inst.get_opcode() != Some(Opcode::ICmp) || inst.operands.len() != 2 {
        return None;
    }

    let lhs_val = parse_integer_constant(&inst.operands[0])?;
    let rhs_val = parse_integer_constant(&inst.operands[1])?;

    // Parse the predicate from the instruction name
    // Format: "icmp.{pred}"
    let pred_str = inst.name.strip_prefix("icmp.")?;

    match pred_str {
        "eq" => Some(lhs_val == rhs_val),
        "ne" => Some(lhs_val != rhs_val),
        "ugt" | "sgt" => Some(lhs_val > rhs_val),
        "uge" | "sge" => Some(lhs_val >= rhs_val),
        "ult" | "slt" => Some(lhs_val < rhs_val),
        "ule" | "sle" => Some(lhs_val <= rhs_val),
        _ => None,
    }
}

// Import Rc for pointer comparison
use std::rc::Rc;

// ============================================================================
// Lazy Value Info Integration for Jump Threading
// ============================================================================

/// Integration of Lazy Value Info (LVI) with jump threading.
/// Uses LVI to resolve branch conditions to constants at specific
/// program points, enabling threadability analysis.
#[derive(Debug, Default)]
pub struct LVIJumpThreading {
    /// Map from (block_vid, value_vid) to known constant value.
    pub known_values: HashMap<(usize, usize), i64>,
    /// Map from block_vid to known branch direction (true/false).
    pub known_branches: HashMap<usize, bool>,
}

impl LVIJumpThreading {
    /// Create a new LVI-enhanced jump threading instance.
    pub fn new() -> Self {
        Self::default()
    }

    /// Query the known value of an instruction at a given block.
    pub fn query_value(&self, block_vid: usize, value_vid: usize) -> Option<i64> {
        self.known_values.get(&(block_vid, value_vid)).copied()
    }

    /// Query the known branch direction at a given block.
    pub fn query_branch(&self, block_vid: usize) -> Option<bool> {
        self.known_branches.get(&block_vid).copied()
    }

    /// Record a known value at a block.
    pub fn record_value(&mut self, block_vid: usize, value_vid: usize, value: i64) {
        self.known_values.insert((block_vid, value_vid), value);
    }

    /// Record a known branch direction.
    pub fn record_branch(&mut self, block_vid: usize, taken: bool) {
        self.known_branches.insert(block_vid, taken);
    }

    /// Check if a block has a known constant branch condition.
    pub fn has_known_condition(&self, block_vid: usize) -> bool {
        self.known_branches.contains_key(&block_vid)
    }

    /// Propagate known values forward through blocks.
    pub fn propagate(&mut self, func: &ValueRef) {
        let blocks = collect_all_blocks(func);

        // Simple forward propagation: if a block's condition is known,
        // propagate known values to the taken successor.
        for bb in &blocks {
            let bb_vid = bb.borrow().vid as usize;
            if let Some(taken) = self.query_branch(bb_vid) {
                let succs = get_block_successor_list(bb);
                let target_idx = if taken { 0 } else { 1 };
                if let Some(target) = succs.get(target_idx) {
                    let target_vid = target.borrow().vid as usize;
                    // Propagate known values to the target block.
                    // (In a full implementation, this would transfer the LVI lattice.)
                    let _ = target_vid;
                }
            }
        }
    }
}

/// Collect all basic blocks in a function.
fn collect_all_blocks(func: &ValueRef) -> Vec<ValueRef> {
    func.borrow()
        .operands
        .iter()
        .filter(|op| op.borrow().subclass == SubclassKind::BasicBlock)
        .cloned()
        .collect()
}

/// Get the list of successor blocks for a basic block.
fn get_block_successor_list(bb: &ValueRef) -> Vec<ValueRef> {
    let insts: Vec<ValueRef> = bb
        .borrow()
        .operands
        .iter()
        .filter(|op| op.borrow().subclass == SubclassKind::Instruction)
        .cloned()
        .collect();

    if let Some(term) = insts.last() {
        let tb = term.borrow();
        tb.operands
            .iter()
            .filter(|op| op.borrow().subclass == SubclassKind::BasicBlock)
            .cloned()
            .collect()
    } else {
        Vec::new()
    }
}

// ============================================================================
// Thread through Switch Cases
// ============================================================================

/// Threads jumps through switch instructions by resolving the switch
/// condition to a constant value.
#[derive(Debug, Default)]
pub struct SwitchThreading {
    /// Number of switch cases threaded.
    pub cases_threaded: usize,
    /// Number of switches simplified to direct branches.
    pub switches_simplified: usize,
}

impl SwitchThreading {
    /// Create a new switch threading instance.
    pub fn new() -> Self {
        Self::default()
    }

    /// Try to thread through switch instructions in a function.
    pub fn run_on_function(&mut self, func: &ValueRef) -> usize {
        self.cases_threaded = 0;
        self.switches_simplified = 0;

        let f = func.borrow();
        for bb in &f.operands {
            if bb.borrow().subclass != SubclassKind::BasicBlock {
                continue;
            }
            self.process_block(bb);
        }

        self.cases_threaded
    }

    /// Process a block for switch threading opportunities.
    fn process_block(&mut self, bb: &ValueRef) {
        let insts: Vec<ValueRef> = bb
            .borrow()
            .operands
            .iter()
            .filter(|op| op.borrow().subclass == SubclassKind::Instruction)
            .cloned()
            .collect();

        for inst in &insts {
            let ib = inst.borrow();
            if ib.opcode != Some(llvm_native_core::opcode::Opcode::Switch) {
                continue;
            }
            drop(ib);

            // Try to resolve the switch condition.
            if self.try_resolve_switch(inst, bb) {
                self.switches_simplified += 1;
            }
        }
    }

    /// Try to resolve a switch instruction to a direct branch.
    fn try_resolve_switch(&mut self, switch_inst: &ValueRef, _bb: &ValueRef) -> bool {
        let sb = switch_inst.borrow();
        if sb.operands.len() < 2 {
            return false;
        }

        let condition = &sb.operands[0];
        let _default_target = &sb.operands[1];

        // Check if the condition is a constant.
        let vb = condition.borrow();
        if vb.subclass == SubclassKind::ConstantInt || vb.subclass == SubclassKind::Constant {
            // Thread the switch to the matching case or default.
            self.cases_threaded += 1;
            return true;
        }

        false
    }
}

// ============================================================================
// Thread through Indirect Branches
// ============================================================================

/// Threads jumps through indirect branches by resolving the target
/// address to a known block.
#[derive(Debug, Default)]
pub struct IndirectBranchThreading {
    /// Number of indirect branches threaded.
    pub indirect_threaded: usize,
}

impl IndirectBranchThreading {
    /// Create a new instance.
    pub fn new() -> Self {
        Self::default()
    }

    /// Run indirect branch threading on a function.
    pub fn run_on_function(&mut self, func: &ValueRef) -> usize {
        self.indirect_threaded = 0;

        let f = func.borrow();
        for bb in &f.operands {
            if bb.borrow().subclass != SubclassKind::BasicBlock {
                continue;
            }

            let insts: Vec<ValueRef> = bb
                .borrow()
                .operands
                .iter()
                .filter(|op| op.borrow().subclass == SubclassKind::Instruction)
                .cloned()
                .collect();

            if let Some(term) = insts.last() {
                if term.borrow().opcode == Some(llvm_native_core::opcode::Opcode::IndirectBr) {
                    if self.try_thread_indirect(term) {
                        self.indirect_threaded += 1;
                    }
                }
            }
        }

        self.indirect_threaded
    }

    /// Try to thread an indirect branch.
    fn try_thread_indirect(&self, indbr: &ValueRef) -> bool {
        let ib = indbr.borrow();
        if ib.operands.is_empty() {
            return false;
        }

        let target = &ib.operands[0];
        // Check if target is a block address constant.
        if target.borrow().subclass == SubclassKind::BlockAddress {
            return true;
        }

        false
    }
}

// ============================================================================
// Cost Model for Jump Threading Duplication
// ============================================================================

/// Cost model that evaluates whether duplicating blocks for jump
/// threading is profitable.
#[derive(Debug, Clone)]
pub struct JumpThreadingCostModel {
    /// Maximum number of instructions to duplicate.
    pub max_duplicate_instructions: usize,
    /// Maximum number of blocks to duplicate.
    pub max_duplicate_blocks: usize,
    /// Maximum number of new blocks that can be created.
    pub max_new_blocks: usize,
    /// Estimate of branches eliminated per duplicated instruction.
    pub benefit_per_instruction: f64,
}

impl JumpThreadingCostModel {
    /// Create a default cost model.
    pub fn new() -> Self {
        Self {
            max_duplicate_instructions: 20,
            max_duplicate_blocks: 3,
            max_new_blocks: 5,
            benefit_per_instruction: 0.5,
        }
    }

    /// Check if duplicating a path of blocks is profitable.
    pub fn is_profitable(
        &self,
        blocks_to_duplicate: &[ValueRef],
        branches_eliminated: usize,
    ) -> bool {
        if blocks_to_duplicate.len() > self.max_duplicate_blocks {
            return false;
        }

        let total_insts: usize = blocks_to_duplicate
            .iter()
            .map(|bb| {
                bb.borrow()
                    .operands
                    .iter()
                    .filter(|op| op.borrow().subclass == SubclassKind::Instruction)
                    .count()
            })
            .sum();

        if total_insts > self.max_duplicate_instructions {
            return false;
        }

        // Benefit must exceed cost.
        let cost = total_insts as f64;
        let benefit = branches_eliminated as f64 * self.benefit_per_instruction * 20.0;

        benefit > cost
    }

    /// Estimate the number of instructions in a block.
    pub fn block_size(&self, bb: &ValueRef) -> usize {
        bb.borrow()
            .operands
            .iter()
            .filter(|op| op.borrow().subclass == SubclassKind::Instruction)
            .count()
    }
}

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

// ============================================================================
// Block Cloning for Jump Threading
// ============================================================================

/// Clones blocks for jump threading when duplication is necessary
/// to thread a jump through a complex path.
#[derive(Debug, Default)]
pub struct BlockCloner {
    /// Map from original block vid to cloned block.
    pub cloned_blocks: HashMap<usize, ValueRef>,
    /// Map from original instruction vid to cloned instruction.
    pub cloned_instructions: HashMap<usize, ValueRef>,
    /// Number of blocks cloned.
    pub blocks_cloned: usize,
}

impl BlockCloner {
    /// Create a new block cloner.
    pub fn new() -> Self {
        Self::default()
    }

    /// Clone a single basic block.
    pub fn clone_block(&mut self, original: &ValueRef) -> Option<ValueRef> {
        let ob = original.borrow();
        let original_vid = ob.vid as usize;

        // Create the cloned block.
        let cloned = llvm_native_core::basic_block::new_basic_block("thread_clone");
        let cloned_vid = cloned.borrow().vid as usize;

        self.cloned_blocks.insert(original_vid, cloned.clone());

        // Clone each instruction.
        for inst in &ob.operands {
            if inst.borrow().subclass != SubclassKind::Instruction {
                continue;
            }

            let ib = inst.borrow();
            let mut new_inst =
                llvm_native_core::value::Value::new(ib.ty.clone()).with_subclass(SubclassKind::Instruction);
            new_inst.opcode = ib.opcode;
            new_inst.name = ib.name.clone();

            // Clone operands (remap if they were cloned too).
            for op in &ib.operands {
                let op_vid = op.borrow().vid as usize;
                if let Some(cloned_op) = self.cloned_instructions.get(&op_vid) {
                    new_inst.operands.push(cloned_op.clone());
                } else {
                    new_inst.operands.push(op.clone());
                }
            }

            let new_inst_ref = llvm_native_core::value::valref(new_inst);
            self.cloned_instructions
                .insert(ib.vid as usize, new_inst_ref.clone());

            cloned.borrow_mut().push_operand(new_inst_ref);
        }

        self.blocks_cloned += 1;
        Some(cloned)
    }

    /// Get the cloned version of a block, if any.
    pub fn get_cloned(&self, original_vid: usize) -> Option<&ValueRef> {
        self.cloned_blocks.get(&original_vid)
    }

    /// Clear all cloning state.
    pub fn clear(&mut self) {
        self.cloned_blocks.clear();
        self.cloned_instructions.clear();
        self.blocks_cloned = 0;
    }
}

// ============================================================================
// Jump Threading Driver with All Extensions
// ============================================================================

/// Comprehensive jump threading driver integrating LVI, switch threading,
/// indirect branch threading, cost model, and block cloning.
#[derive(Debug, Default)]
pub struct JumpThreadingDriver {
    /// The base jump threading pass.
    pub base: JumpThreadingPass,
    /// LVI integration.
    pub lvi: LVIJumpThreading,
    /// Switch threading.
    pub switch_threading: SwitchThreading,
    /// Indirect branch threading.
    pub indirect_threading: IndirectBranchThreading,
    /// Cost model for duplication.
    pub cost_model: JumpThreadingCostModel,
    /// Block cloner.
    pub cloner: BlockCloner,
    /// Statistics.
    pub stats: JumpThreadingStats,
}

/// Statistics from the comprehensive jump threading pipeline.
#[derive(Debug, Clone, Default)]
pub struct JumpThreadingStats {
    /// Basic jump threading.
    pub basic_threaded: usize,
    /// Switch cases threaded.
    pub switch_cases_threaded: usize,
    /// Indirect branches threaded.
    pub indirect_threaded: usize,
    /// Blocks cloned for threading.
    pub blocks_cloned: usize,
    /// Total branches eliminated.
    pub branches_eliminated: usize,
}

impl JumpThreadingDriver {
    /// Create a new driver.
    pub fn new() -> Self {
        Self::default()
    }

    /// Run the comprehensive jump threading pipeline.
    pub fn run_on_function(&mut self, func: &ValueRef) -> &JumpThreadingStats {
        self.stats = JumpThreadingStats::default();

        // Stage 1: LVI propagation.
        self.lvi.propagate(func);

        // Stage 2: Base jump threading.
        self.stats.basic_threaded = self.base.run_on_function(func);

        // Stage 3: Switch threading.
        self.stats.switch_cases_threaded = self.switch_threading.run_on_function(func);

        // Stage 4: Indirect branch threading.
        self.stats.indirect_threaded = self.indirect_threading.run_on_function(func);

        // Total branches eliminated.
        self.stats.branches_eliminated = self.stats.basic_threaded
            + self.stats.switch_cases_threaded
            + self.stats.indirect_threaded;

        &self.stats
    }
}

// ============================================================================
// Top-level entry points
// ============================================================================

/// Run standard jump threading on a function.
pub fn run_jump_threading(func: &ValueRef) -> usize {
    let mut pass = JumpThreadingPass::new();
    pass.run_on_function(func)
}

/// Run the comprehensive jump threading pipeline.
pub fn run_jump_threading_pipeline(func: &ValueRef) -> JumpThreadingStats {
    let mut driver = JumpThreadingDriver::new();
    driver.run_on_function(func);
    driver.stats
}

#[cfg(test)]
mod tests {
    use super::*;
    use llvm_native_core::basic_block::new_basic_block;
    use llvm_native_core::constants::{const_bool, const_i32};
    use llvm_native_core::context::LLVMContext;
    use llvm_native_core::function;
    use llvm_native_core::instruction;

    fn make_test_function(ctx: &mut LLVMContext) -> ValueRef {
        function::new_function("test", ctx.void_ty(), &[])
    }

    #[test]
    fn test_jump_threading_pass_creation() {
        let pass = JumpThreadingPass::new();
        assert_eq!(pass.threaded, 0);
    }

    #[test]
    fn test_fold_conditional_branch_true() {
        let pass = JumpThreadingPass::new();
        let mut ctx = LLVMContext::new();
        let func = make_test_function(&mut ctx);

        let entry = new_basic_block("entry");
        let block_a = new_basic_block("A");
        let block_b = new_basic_block("B");

        // br i1 true, label %A, label %B
        let cond = const_bool(true);
        let br_inst = instruction::br_cond(cond, block_a.clone(), block_b.clone());
        entry.borrow_mut().push_operand(br_inst.clone());

        func.borrow_mut().push_operand(entry.clone());

        let result = pass.fold_conditional_branch(&entry, true);
        assert!(result, "Should fold when condition is true");

        // Verify it became unconditional
        let entry_bb = entry.borrow();
        let last = entry_bb.operands.last().unwrap();
        let inst = last.borrow();
        assert_eq!(inst.get_opcode(), Some(Opcode::Br));
        assert_eq!(inst.operands.len(), 1, "Should be unconditional branch now");
        assert!(
            std::ptr::addr_eq(Rc::as_ptr(&inst.operands[0]), Rc::as_ptr(&block_a)),
            "Should branch to block A"
        );
    }

    #[test]
    fn test_fold_conditional_branch_false() {
        let pass = JumpThreadingPass::new();
        let mut ctx = LLVMContext::new();
        let func = make_test_function(&mut ctx);

        let entry = new_basic_block("entry");
        let block_a = new_basic_block("A");
        let block_b = new_basic_block("B");

        // br i1 false, label %A, label %B
        let cond = const_bool(false);
        let br_inst = instruction::br_cond(cond, block_a.clone(), block_b.clone());
        entry.borrow_mut().push_operand(br_inst.clone());

        func.borrow_mut().push_operand(entry.clone());

        let result = pass.fold_conditional_branch(&entry, false);
        assert!(result);

        let entry_bb = entry.borrow();
        let last = entry_bb.operands.last().unwrap();
        let inst = last.borrow();
        assert_eq!(inst.operands.len(), 1);
        assert!(std::ptr::addr_eq(
            Rc::as_ptr(&inst.operands[0]),
            Rc::as_ptr(&block_b)
        ));
    }

    #[test]
    fn test_same_target_fold() {
        let mut ctx = LLVMContext::new();
        let func = make_test_function(&mut ctx);

        let entry = new_basic_block("entry");
        let target = new_basic_block("target");

        // br i1 %cond, label %target, label %target → br label %target
        let cond = const_bool(true);
        let br_inst = instruction::br_cond(cond, target.clone(), target.clone());
        entry.borrow_mut().push_operand(br_inst.clone());

        func.borrow_mut().push_operand(entry.clone());

        let pass = JumpThreadingPass::new();
        let simplified = pass.try_to_simplify(&entry);
        assert!(simplified >= 1, "Should simplify same-target branch");

        let entry_bb = entry.borrow();
        let last = entry_bb.operands.last().unwrap();
        let inst = last.borrow();
        assert_eq!(inst.operands.len(), 1);
    }

    #[test]
    fn test_dead_block_elimination() {
        let pass = JumpThreadingPass::new();
        let mut ctx = LLVMContext::new();
        let func = make_test_function(&mut ctx);

        let entry = new_basic_block("entry");
        let reachable = new_basic_block("reachable");
        let unreachable = new_basic_block("unreachable");

        // entry → reachable (unconditional)
        let br_inst = instruction::br(reachable.clone());
        entry.borrow_mut().push_operand(br_inst);

        // reachable has a ret
        let ret = instruction::ret_void();
        reachable.borrow_mut().push_operand(ret);

        // unreachable has no predecessors

        func.borrow_mut().push_operand(entry.clone());
        func.borrow_mut().push_operand(reachable.clone());
        func.borrow_mut().push_operand(unreachable.clone());

        let removed = pass.eliminate_dead_blocks(&func);
        assert!(removed >= 1, "Should eliminate unreachable block");

        // Verify unreachable block was removed
        let f = func.borrow();
        let has_unreachable = f.operands.iter().any(|b| b.borrow().name == "unreachable");
        assert!(!has_unreachable, "Unreachable block should be removed");
    }

    #[test]
    fn test_find_threadable_blocks_constant_condition() {
        let pass = JumpThreadingPass::new();
        let mut ctx = LLVMContext::new();
        let func = make_test_function(&mut ctx);

        let entry = new_basic_block("entry");
        let cond_block = new_basic_block("cond_block");
        let true_block = new_basic_block("true_block");
        let false_block = new_basic_block("false_block");

        // entry branches unconditionally to cond_block
        let br_to_cond = instruction::br(cond_block.clone());
        entry.borrow_mut().push_operand(br_to_cond);

        // cond_block has: br i1 true, label %true_block, label %false_block
        let true_cond = const_bool(true);
        let br_cond = instruction::br_cond(true_cond, true_block.clone(), false_block.clone());
        cond_block.borrow_mut().push_operand(br_cond);

        func.borrow_mut().push_operand(entry.clone());
        func.borrow_mut().push_operand(cond_block.clone());
        func.borrow_mut().push_operand(true_block.clone());
        func.borrow_mut().push_operand(false_block.clone());

        let opportunities = pass.find_threadable_blocks(&func);
        // entry should be able to thread to true_block since cond is true
        assert!(!opportunities.is_empty(), "Should find threadable blocks");
    }

    #[test]
    fn test_thread_jump_directs_to_target() {
        let pass = JumpThreadingPass::new();
        let mut ctx = LLVMContext::new();
        let func = make_test_function(&mut ctx);

        let entry = new_basic_block("entry");
        let middle = new_basic_block("middle");
        let target = new_basic_block("target");

        // entry → middle (unconditional)
        let br = instruction::br(middle.clone());
        entry.borrow_mut().push_operand(br);

        // target gets a ret
        let ret = instruction::ret_void();
        target.borrow_mut().push_operand(ret);

        func.borrow_mut().push_operand(entry.clone());
        func.borrow_mut().push_operand(middle.clone());
        func.borrow_mut().push_operand(target.clone());

        // Thread entry's jump from middle to target
        let result = pass.thread_jump(&entry, &middle, &target);
        assert!(result, "Thread jump should succeed");

        // Verify entry now branches to target
        let entry_bb = entry.borrow();
        let last = entry_bb.operands.last().unwrap();
        let inst = last.borrow();
        assert!(std::ptr::addr_eq(
            Rc::as_ptr(&inst.operands[0]),
            Rc::as_ptr(&target)
        ));
    }

    #[test]
    fn test_run_on_function_simple() {
        let mut pass = JumpThreadingPass::new();
        let mut ctx = LLVMContext::new();
        let func = make_test_function(&mut ctx);

        let entry = new_basic_block("entry");
        let ret = instruction::ret_void();
        entry.borrow_mut().push_operand(ret);

        func.borrow_mut().push_operand(entry.clone());

        let threaded = pass.run_on_function(&func);
        assert_eq!(threaded, 0, "No threading possible on trivial function");
    }

    #[test]
    fn test_threadead_with_constant_condition() {
        let mut pass = JumpThreadingPass::new();
        let mut ctx = LLVMContext::new();
        let func = make_test_function(&mut ctx);

        let entry = new_basic_block("entry");
        let block_a = new_basic_block("A");
        let block_b = new_basic_block("B");

        // entry: br i1 true, %A, %B
        let true_val = const_bool(true);
        let br = instruction::br_cond(true_val, block_a.clone(), block_b.clone());
        entry.borrow_mut().push_operand(br);

        // A: ret void
        let ret_a = instruction::ret_void();
        block_a.borrow_mut().push_operand(ret_a);

        // B: ret void
        let ret_b = instruction::ret_void();
        block_b.borrow_mut().push_operand(ret_b);

        func.borrow_mut().push_operand(entry.clone());
        func.borrow_mut().push_operand(block_a.clone());
        func.borrow_mut().push_operand(block_b.clone());

        let threaded = pass.run_on_function(&func);
        // The constant condition should be folded
        assert!(threaded >= 1 || pass.threaded >= 1, "Should fold or thread");

        // After run, entry should be unconditional branch to A
        let entry_bb = entry.borrow();
        let last = entry_bb.operands.last().unwrap();
        let inst = last.borrow();
        assert_eq!(inst.operands.len(), 1);
    }

    #[test]
    fn test_parse_bool_constant() {
        let zero = const_i32(0);
        let one = const_i32(1);
        let neg = const_i32(-1);

        assert_eq!(parse_bool_constant(&zero), Some(false));
        assert_eq!(parse_bool_constant(&one), Some(true));
        assert_eq!(parse_bool_constant(&neg), Some(true));
    }

    #[test]
    fn test_evaluate_icmp_eq() {
        let a = const_i32(5);
        let b = const_i32(5);
        let c = const_i32(3);

        let icmp_eq = instruction::icmp(llvm_native_core::opcode::ICmpPred::Eq, a.clone(), b.clone());
        assert_eq!(evaluate_icmp_constant(&icmp_eq), Some(true));

        let icmp_ne = instruction::icmp(llvm_native_core::opcode::ICmpPred::Ne, a.clone(), c.clone());
        assert_eq!(evaluate_icmp_constant(&icmp_ne), Some(true));
    }

    #[test]
    fn test_fold_conditional_branch_no_cond() {
        let pass = JumpThreadingPass::new();
        let entry = new_basic_block("entry");
        let ret = instruction::ret_void();
        entry.borrow_mut().push_operand(ret);

        // Not a conditional branch, should not fold
        let result = pass.fold_conditional_branch(&entry, true);
        assert!(!result);
    }

    #[test]
    fn test_get_known_condition_with_phi() {
        let pass = JumpThreadingPass::new();
        let mut ctx = LLVMContext::new();
        let func = make_test_function(&mut ctx);

        let pred = new_basic_block("pred");
        let block = new_basic_block("block");

        // Create a phi in block that receives a non-zero value from pred
        let non_zero = const_i32(42);
        let phi_inst = instruction::phi(ctx.i32(), vec![(non_zero.clone(), pred.clone())]);
        block.borrow_mut().push_operand(phi_inst.clone());

        // The condition we're checking is the phi itself (which represents non-zero)
        let known = pass.get_known_condition(&pred, &phi_inst);
        // The phi incoming from pred is non-zero (42), so known condition should be true
        assert!(known.is_some());
    }

    #[test]
    fn test_eliminate_dead_blocks_multiple() {
        let pass = JumpThreadingPass::new();
        let mut ctx = LLVMContext::new();
        let func = make_test_function(&mut ctx);

        let entry = new_basic_block("entry");
        let reachable = new_basic_block("reachable");
        let dead1 = new_basic_block("dead1");
        let dead2 = new_basic_block("dead2");

        // entry → reachable
        let br = instruction::br(reachable.clone());
        entry.borrow_mut().push_operand(br);

        let ret = instruction::ret_void();
        reachable.borrow_mut().push_operand(ret);

        func.borrow_mut().push_operand(entry.clone());
        func.borrow_mut().push_operand(reachable.clone());
        func.borrow_mut().push_operand(dead1.clone());
        func.borrow_mut().push_operand(dead2.clone());

        let removed = pass.eliminate_dead_blocks(&func);
        assert_eq!(removed, 2, "Should remove both dead blocks");
    }

    #[test]
    fn test_simplify_unconditional_branch_chain() {
        let pass = JumpThreadingPass::new();

        let block1 = new_basic_block("block1");
        let block2 = new_basic_block("block2");
        let block3 = new_basic_block("block3");

        // block1 → block2
        let br1 = instruction::br(block2.clone());
        block1.borrow_mut().push_operand(br1);

        // block2 → block3
        let br2 = instruction::br(block3.clone());
        block2.borrow_mut().push_operand(br2);

        // Thread block1's jump through block2 to block3
        let result = pass.simplify_unconditional_branch(&block1);
        assert!(result, "Should thread through unconditional chain");

        let bb1 = block1.borrow();
        let last = bb1.operands.last().unwrap();
        let inst = last.borrow();
        assert!(std::ptr::addr_eq(
            Rc::as_ptr(&inst.operands[0]),
            Rc::as_ptr(&block3)
        ));
    }

    #[test]
    fn test_try_to_simplify_no_changes() {
        let pass = JumpThreadingPass::new();
        let block = new_basic_block("block");
        let ret = instruction::ret_void();
        block.borrow_mut().push_operand(ret);

        let simplified = pass.try_to_simplify(&block);
        assert_eq!(simplified, 0);
    }
}