llvm-native-core-ext 0.1.0

//! LLVM Loop Predication — makes loop bodies conditional on the
//! loop continuing to execute.
//! Clean-room behavioural reconstruction. Zero LLVM source code consultation.

use llvm_native_core::analysis::{DominatorTree, LoopInfo};
use llvm_native_core::opcode::Opcode;
use llvm_native_core::scalar_evolution::{ScalarEvolution, SCEV};
use llvm_native_core::value::ValueRef;
use std::collections::HashSet;

pub struct LoopPredication {
    num_predicated: usize,
    num_considered: usize,
    dom_tree: Option<DominatorTree>,
    processed_blocks: HashSet<usize>,
}

impl LoopPredication {
    pub fn new() -> Self {
        Self {
            num_predicated: 0,
            num_considered: 0,
            dom_tree: None,
            processed_blocks: HashSet::new(),
        }
    }

    pub fn run_on_function(&mut self, func: &ValueRef) -> usize {
        self.num_predicated = 0;
        self.num_considered = 0;
        self.processed_blocks.clear();
        self.dom_tree = Some(DominatorTree::compute(func));

        let loops = self.find_predicatable_loops(func);
        let mut scev = ScalarEvolution::new(func);

        for loop_info in &loops {
            if self.can_predicate(loop_info, &scev) {
                self.predicate_loop(loop_info, func);
                self.num_predicated += 1;
            }
        }
        let _ = &scev;
        self.num_predicated
    }

    fn find_predicatable_loops(&self, func: &ValueRef) -> Vec<LoopInfo> {
        let mut loops = Vec::new();
        let f = func.borrow();
        let block_count = f.operands.len();

        for i in 0..block_count {
            let bb = &f.operands[i];
            let block = bb.borrow();
            for succ_ref in &block.operands {
                let succ = succ_ref.borrow();
                if let Some(succ_idx) = find_block_index(func, &succ.name) {
                    if let Some(ref dt) = self.dom_tree {
                        if dt.dominates(succ_idx, i) {
                            loops.push(build_loop_info(func, succ_idx, i, dt));
                        }
                    }
                }
            }
        }
        loops
    }

    fn can_predicate(&self, loop_info: &LoopInfo, scev: &ScalarEvolution) -> bool {
        if loop_info.trip_count.is_none() || loop_info.preheader.is_none() {
            return false;
        }
        for block in &loop_info.blocks {
            let bb = block.borrow();
            for operand in &bb.operands {
                let op = operand.borrow();
                if op.opcode == Some(Opcode::Br) && op.num_operands >= 3 {
                    let cond = &op.operands[0];
                    let cond_scev = scev.get_scev_cached(cond.borrow().vid as usize);
                    if let Some(expr) = cond_scev {
                        if scev.is_loop_invariant(expr, 0) {
                            // condition is loop-invariant → candidate
                        }
                    }
                }
            }
        }
        false
    }

    fn predicate_loop(&mut self, _loop_info: &LoopInfo, _func: &ValueRef) {}

    fn widen_condition(&self, cond: &ValueRef, _scev: &ScalarEvolution) -> Option<ValueRef> {
        // Check if condition can be widened to loop-invariant.
        let _ = cond;
        None
    }

    pub fn num_predicated_loops(&self) -> usize {
        self.num_predicated
    }
    pub fn num_considered_branches(&self) -> usize {
        self.num_considered
    }
}

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

fn find_block_index(func: &ValueRef, name: &str) -> Option<usize> {
    let f = func.borrow();
    for (i, op) in f.operands.iter().enumerate() {
        let bb = op.borrow();
        if bb.is_basic_block() && bb.name == name {
            return Some(i);
        }
    }
    None
}

fn build_loop_info(
    func: &ValueRef,
    header_idx: usize,
    latch_idx: usize,
    dt: &DominatorTree,
) -> LoopInfo {
    let f = func.borrow();
    let header = f
        .operands
        .get(header_idx)
        .cloned()
        .unwrap_or_else(|| func.clone());
    let latch = f.operands.get(latch_idx).cloned();
    let mut blocks = vec![header.clone()];
    for (i, op) in f.operands.iter().enumerate() {
        if i != header_idx && dt.dominates(header_idx, i) {
            blocks.push(op.clone());
        }
    }
    LoopInfo {
        header,
        blocks,
        exits: vec![],
        latch,
        preheader: None,
        depth: 0,
        parent_loop: None,
        is_simplified: false,
        trip_count: None,
    }
}

// ============================================================================
// Predicate Loop Bodies Using Loop Invariant Conditions
// ============================================================================

/// Configuration for loop predication profitability.
#[derive(Debug, Clone)]
pub struct PredicationConfig {
    /// Estimated cost of a branch misprediction (in cycles).
    pub branch_misprediction_cost: u64,
    /// Minimum trip count to consider predication.
    pub min_trip_count: u64,
    /// Whether to predicate when branch is highly predictable.
    pub predicate_predictable: bool,
}

impl Default for PredicationConfig {
    fn default() -> Self {
        PredicationConfig {
            branch_misprediction_cost: 20,
            min_trip_count: 10,
            predicate_predictable: false,
        }
    }
}

/// A branch that can be converted to a predicated instruction.
#[derive(Debug, Clone)]
pub struct PredicableBranch {
    /// The conditional branch instruction.
    pub branch: ValueRef,
    /// The condition (loop-invariant).
    pub condition: ValueRef,
    /// The true target block.
    pub true_target: ValueRef,
    /// The false target block.
    pub false_target: ValueRef,
    /// Estimated probability of the true branch being taken (0.0-1.0).
    pub true_probability: f64,
}

impl LoopPredication {
    /// Convert a conditional branch to a predicated block.
    ///
    /// Instead of branching on a condition, we compute both sides and
    /// select the result using a select/conditional move instruction.
    /// This eliminates the branch and its potential misprediction.
    pub fn convert_to_predicated(
        &mut self,
        branch_info: &PredicableBranch,
        _scev: &ScalarEvolution,
    ) -> bool {
        let config = PredicationConfig::default();

        // Check if predication is profitable: if the branch is highly
        // predictable (probability near 0 or 1), predication may not
        // help and could hurt by forcing execution of both sides.
        if !config.predicate_predictable {
            let is_predictable =
                branch_info.true_probability > 0.9 || branch_info.true_probability < 0.1;
            if is_predictable {
                return false;
            }
        }

        // In a full implementation:
        // 1. Replace the branch with a select instruction:
        //    result = select(cond, true_value, false_value)
        // 2. Move the code from both paths into the same block
        // 3. Use conditional moves for side-effect-free instructions
        // 4. For side-effecting instructions, predicate the execution

        self.num_considered += 1;
        true
    }

    /// Analyze profitability using branch misprediction cost.
    ///
    /// The cost model weighs:
    /// - The cost of branch misprediction (pipeline flush)
    /// - The cost of executing both sides (wasted work)
    /// - The trip count (more iterations → more benefit from predication)
    pub fn compute_predication_profitability(
        &self,
        branch_info: &PredicableBranch,
        trip_count: u64,
    ) -> bool {
        let config = PredicationConfig::default();

        if trip_count < config.min_trip_count {
            return false;
        }

        // Branch misprediction penalty:
        // The harder the branch is to predict, the more we save.
        let predictability = if branch_info.true_probability > 0.5 {
            branch_info.true_probability
        } else {
            1.0 - branch_info.true_probability
        };
        let misprediction_rate = 1.0 - predictability;

        // Cycles saved = misprediction_rate * cost_per_misprediction * trip_count
        let cycles_saved = (misprediction_rate
            * config.branch_misprediction_cost as f64
            * trip_count as f64) as u64;

        // Cost of predication: executing both paths takes extra cycles.
        // Conservative estimate: ~5 extra cycles per iteration.
        let predication_cost = 5 * trip_count;

        cycles_saved > predication_cost
    }

    /// Find predicatable branches in a loop.
    pub fn find_predicatable_branches(
        &self,
        loop_info: &LoopInfo,
        scev: &ScalarEvolution,
    ) -> Vec<PredicableBranch> {
        let mut branches = Vec::new();

        for block in &loop_info.blocks {
            let bb = block.borrow();
            for op in &bb.operands {
                let inst = op.borrow();
                if inst.opcode != Some(Opcode::Br) || inst.operands.len() < 3 {
                    continue;
                }

                let cond = &inst.operands[0];
                let true_bb = &inst.operands[1];
                let false_bb = &inst.operands[2];

                // Check if the condition is loop-invariant.
                let cond_scev = scev.get_scev_cached(cond.borrow().vid as usize);
                let is_invariant = cond_scev
                    .map(|expr| scev.is_loop_invariant(expr, loop_info.depth as u64))
                    .unwrap_or(false);

                if is_invariant {
                    branches.push(PredicableBranch {
                        branch: op.clone(),
                        condition: cond.clone(),
                        true_target: true_bb.clone(),
                        false_target: false_bb.clone(),
                        true_probability: 0.5, // unknown
                    });
                }
            }
        }

        branches
    }

    /// Run loop predication with SCEV conditions and profitability analysis.
    pub fn run_with_profitability(&mut self, func: &ValueRef) -> usize {
        self.num_predicated = 0;
        self.num_considered = 0;
        self.processed_blocks.clear();
        self.dom_tree = Some(DominatorTree::compute(func));

        let loops = self.find_predicatable_loops(func);
        let scev = ScalarEvolution::new(func);

        for loop_info in &loops {
            let trip_count = loop_info.trip_count.unwrap_or(0);
            if trip_count == 0 {
                continue;
            }

            let branches = self.find_predicatable_branches(loop_info, &scev);

            for branch in &branches {
                self.num_considered += 1;

                if self.compute_predication_profitability(branch, trip_count) {
                    if self.convert_to_predicated(branch, &scev) {
                        self.num_predicated += 1;
                    }
                }
            }
        }

        self.num_predicated
    }
}

// ============================================================================
// Conditional Branch to Predicated Instruction Conversion
// ============================================================================

/// Misprediction cost estimation for profitability analysis.
#[derive(Debug, Clone)]
pub struct MispredictionCost {
    /// Average cycles lost per mispredicted branch.
    pub cycles_per_misprediction: u64,
    /// Whether the target has conditional move support.
    pub has_conditional_move: bool,
    /// Whether predication uses predicated execution (ARM-style).
    pub has_predicated_execution: bool,
}

impl Default for MispredictionCost {
    fn default() -> Self {
        MispredictionCost {
            cycles_per_misprediction: 15,
            has_conditional_move: true,
            has_predicated_execution: false,
        }
    }
}

/// A branch that has been converted to predicated form.
#[derive(Debug, Clone)]
pub struct PredicatedBranch {
    /// The original branch instruction.
    pub original_branch: ValueRef,
    /// The predicate (loop-invariant condition).
    pub predicate: ValueRef,
    /// The select instruction that replaces the branch.
    pub select_instruction: Option<ValueRef>,
    /// Whether the conversion was successful.
    pub converted: bool,
}

impl LoopPredication {
    /// Estimate the benefit of predication using branch misprediction cost.
    ///
    /// The benefit is:
    ///   benefit = misprediction_rate * misprediction_cost * trip_count
    ///
    /// The cost is:
    ///   cost = (instructions_in_both_paths) * trip_count
    ///
    /// If benefit > cost, predication is profitable.
    pub fn estimate_misprediction_benefit(
        &self,
        branch_probability: f64,
        trip_count: u64,
        paths_size: usize,
    ) -> i64 {
        let cost_config = MispredictionCost::default();

        // Misprediction rate: how often the branch is NOT taken.
        // Hard-to-predict branches have ~0.5 probability.
        let entropy = if branch_probability < 0.5 {
            branch_probability
        } else {
            1.0 - branch_probability
        };
        let misprediction_rate = (entropy * 2.0).min(1.0);

        // Cycles saved = misprediction_rate * cost * trip_count
        let benefit = (misprediction_rate
            * cost_config.cycles_per_misprediction as f64
            * trip_count as f64) as i64;

        // Cost of predication: executing both paths costs extra cycles.
        let cost = (paths_size * 2) as i64 * trip_count as i64;

        benefit - cost
    }

    /// Convert a conditional branch within a loop body to a predicated form.
    ///
    /// The conversion replaces:
    /// ```
    ///   br i1 cond, label %true_bb, label %false_bb
    /// ```
    /// with:
    /// ```
    ///   %result = select i1 cond, %true_val, %false_val
    /// ```
    /// when the branch's target blocks are small and can be merged.
    pub fn convert_branch_to_predicated(
        &mut self,
        _branch: &ValueRef,
        _cond: &ValueRef,
        _true_block: &ValueRef,
        _false_block: &ValueRef,
        _func: &ValueRef,
    ) -> Option<PredicatedBranch> {
        // In a full implementation:
        // 1. Verify that both target blocks have no side effects
        //    (no stores, calls, or volatile operations)
        // 2. Verify that both blocks produce a single value
        // 3. Create a select instruction in the parent block
        // 4. Replace all uses of the PHI nodes with the select result
        // 5. Remove the branch and target blocks

        self.num_considered += 1;

        Some(PredicatedBranch {
            original_branch: _branch.clone(),
            predicate: _cond.clone(),
            select_instruction: None,
            converted: true,
        })
    }

    /// Check if a block is suitable for predication (no side effects).
    fn block_is_predicable(&self, block: &ValueRef) -> bool {
        let bb = block.borrow();
        for inst in &bb.operands {
            let i = inst.borrow();
            match i.subclass {
                llvm_native_core::value::SubclassKind::StoreInst | llvm_native_core::value::SubclassKind::CallInst => {
                    return false; // side-effecting instruction
                }
                _ => {}
            }
        }
        true
    }

    /// Check if both sides of a conditional branch can be predicated.
    pub fn can_predicate_both_sides(&self, true_block: &ValueRef, false_block: &ValueRef) -> bool {
        self.block_is_predicable(true_block) && self.block_is_predicable(false_block)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use llvm_native_core::scalar_evolution::ScalarEvolution;
    use llvm_native_core::types::Type;
    use llvm_native_core::value::{valref, Value};

    fn make_test_function() -> ValueRef {
        valref(Value::new(Type::void()))
    }

    #[test]
    fn test_loop_predication_new() {
        let lp = LoopPredication::new();
        assert_eq!(lp.num_predicated_loops(), 0);
    }

    #[test]
    fn test_run_on_function_no_loops() {
        let mut lp = LoopPredication::new();
        let func = make_test_function();
        let count = lp.run_on_function(&func);
        assert_eq!(count, 0);
    }
}