llvm-native-core-ext 0.1.0

//! LLVM SpeculativeExecution — hoists safe instructions above branches to
//! expose more optimization opportunities.
//! Clean-room behavioral reconstruction.
//!
//! This pass identifies instructions that are safe and cheap to execute
//! speculatively (i.e., before a branch whose direction is not yet known)
//! and hoists them to predecessor blocks. This can expose CSE opportunities,
//! reduce critical path length, and improve scheduling.
//!
//! Algorithm:
//!   1. For each basic block with multiple predecessors, scan instructions
//!   2. Check if the instruction is safe to speculate (no side effects,
//!      no trapping, no divide-by-zero, etc.)
//!   3. Check if the instruction is cheap to speculate (single-cycle ops,
//!      not memory, not expensive float ops)
//!   4. If all operands dominate the block, hoist to each predecessor
//!   5. Replace original instruction with the hoisted copies (or a PHI)
//!
//! Safety rules:
//!   - No loads (may fault / cause page faults)
//!   - No divides (may divide by zero)
//!   - No calls (may have side effects)
//!   - No stores (side effects)
//!   - Safe: add, sub, mul, shl, and, or, xor, trunc, zext, sext, bitcast
//!   - Safe: icmp, fcmp, select
//!   - Conditionally safe: gep (if all indices are in bounds)

use llvm_native_core::value::{SubclassKind, ValueRef};

// ============================================================================
// Speculative Execution Pass
// ============================================================================

/// SpeculativeExecution pass — hoists safe, cheap instructions above
/// conditional branches to expose downstream optimizations.
pub struct SpeculativeExecution {
    /// Number of instructions successfully speculated.
    pub speculated: usize,
}

impl SpeculativeExecution {
    /// Create a new SpeculativeExecution pass.
    pub fn new() -> Self {
        Self { speculated: 0 }
    }

    // ========================================================================
    // Main entry point
    // ========================================================================

    /// Run speculative execution on a function. Returns the number of
    /// instructions hoisted speculatively.
    pub fn run_on_function(&mut self, func: &ValueRef) -> usize {
        self.speculated = 0;

        let f = func.borrow();
        let blocks: Vec<ValueRef> = f
            .operands
            .iter()
            .filter(|op| op.borrow().subclass == SubclassKind::BasicBlock)
            .cloned()
            .collect();
        drop(f);

        for bb in &blocks {
            let hoistable = self.find_hoistable_instructions(bb);
            for inst in hoistable {
                self.hoist_to_predecessors(&inst, func);
            }
        }

        self.speculated
    }

    // ========================================================================
    // Finding hoistable instructions in a basic block
    // ========================================================================

    /// Find all instructions in a basic block that are candidates for
    /// speculative hoisting.
    fn find_hoistable_instructions(&self, bb: &ValueRef) -> Vec<ValueRef> {
        let block = bb.borrow();

        // Skip entry blocks or blocks with fewer than 2 predecessors
        let pred_count = {
            // Count predecessors by checking which blocks have this as successor
            let parent = block.parent.clone();
            if let Some(p) = parent {
                let func = p.borrow();
                func.operands
                    .iter()
                    .filter(|op| {
                        let other = op.borrow();
                        other.subclass == SubclassKind::BasicBlock
                            && other.successors.iter().any(|s| s.borrow().vid == block.vid)
                    })
                    .count()
            } else {
                0
            }
        };

        if pred_count < 2 {
            return Vec::new();
        }

        // Collect hoistable instructions
        let mut hoistable = Vec::new();

        for inst in &block.operands {
            let i = inst.borrow();

            // Only consider real instructions (not PHI nodes at start)
            if !i.is_instruction() {
                continue;
            }

            // Skip PHI nodes — they're already at the top
            if i.name.to_lowercase().contains("phi") {
                continue;
            }

            // Skip terminators (ret, br, switch)
            if i.is_terminator() {
                break;
            }

            // Must be both safe and cheap to speculate
            if self.is_safe_to_speculate(inst) && self.is_cheap_to_speculate(inst) {
                hoistable.push(inst.clone());
            }
        }

        hoistable
    }

    // ========================================================================
    // Safety analysis
    // ========================================================================

    /// Check if an instruction is safe to execute speculatively.
    /// An instruction is safe if:
    ///   - It does not have side effects (no stores, no calls)
    ///   - It cannot trap or fault (no loads, no divides)
    ///   - All its operands are available (dominate the block)
    fn is_safe_to_speculate(&self, inst: &ValueRef) -> bool {
        let i = inst.borrow();

        // Must have an opcode for classification
        let opcode = match i.opcode {
            Some(op) => op,
            None => return false,
        };

        // Unsafe opcodes: anything that can trap, fault, or has side effects
        match opcode {
            // Binary ops: safe
            llvm_native_core::opcode::Opcode::Add
            | llvm_native_core::opcode::Opcode::Sub
            | llvm_native_core::opcode::Opcode::Mul
            | llvm_native_core::opcode::Opcode::And
            | llvm_native_core::opcode::Opcode::Or
            | llvm_native_core::opcode::Opcode::Xor
            | llvm_native_core::opcode::Opcode::Shl
            | llvm_native_core::opcode::Opcode::LShr
            | llvm_native_core::opcode::Opcode::AShr
            | llvm_native_core::opcode::Opcode::FAdd
            | llvm_native_core::opcode::Opcode::FSub
            | llvm_native_core::opcode::Opcode::FMul => true,

            // Cast ops: safe (no trapping on trunc/zext/sext/bitcast)
            llvm_native_core::opcode::Opcode::Trunc
            | llvm_native_core::opcode::Opcode::ZExt
            | llvm_native_core::opcode::Opcode::SExt
            | llvm_native_core::opcode::Opcode::BitCast => true,

            // Comparison ops: safe (read-only)
            llvm_native_core::opcode::Opcode::ICmp | llvm_native_core::opcode::Opcode::FCmp => true,

            // Select: safe if both value operands are available
            llvm_native_core::opcode::Opcode::Select => {
                // Make sure all operands dominate
                i.operands.len() >= 3
            }

            // GEP: safe if off the end of an array (it just computes
            // an address, doesn't load)
            llvm_native_core::opcode::Opcode::GetElementPtr => true,

            // Everything else: conservatively unsafe
            // Load, Store, Call, Invoke, Div/Rem, etc.
            _ => false,
        }
    }

    // ========================================================================
    // Cost model
    // ========================================================================

    /// Check if an instruction is cheap enough to be worth speculating.
    /// We avoid speculating expensive operations (like wide multiplies,
    /// float division, etc.) because the speculative execution might
    /// waste resources if the branch goes the other way.
    fn is_cheap_to_speculate(&self, inst: &ValueRef) -> bool {
        let i = inst.borrow();

        let opcode = match i.opcode {
            Some(op) => op,
            None => return false,
        };

        match opcode {
            // Single-cycle integer ops: always cheap
            llvm_native_core::opcode::Opcode::Add
            | llvm_native_core::opcode::Opcode::Sub
            | llvm_native_core::opcode::Opcode::And
            | llvm_native_core::opcode::Opcode::Or
            | llvm_native_core::opcode::Opcode::Xor
            | llvm_native_core::opcode::Opcode::Shl
            | llvm_native_core::opcode::Opcode::LShr
            | llvm_native_core::opcode::Opcode::AShr => true,

            // Integer multiply: cheap for small widths, expensive for wide
            llvm_native_core::opcode::Opcode::Mul => {
                // Check if the type is a small integer (<= 64 bits)
                let bits = i.ty.integer_bit_width();
                bits <= 64
            }

            // Float ops: conditionally cheap
            // fadd, fsub, fmul are typically single-cycle on modern CPUs
            llvm_native_core::opcode::Opcode::FAdd
            | llvm_native_core::opcode::Opcode::FSub
            | llvm_native_core::opcode::Opcode::FMul => true,

            // Float division: expensive, don't speculate
            llvm_native_core::opcode::Opcode::FDiv | llvm_native_core::opcode::Opcode::FRem => false,

            // Casts: always cheap (just bits)
            llvm_native_core::opcode::Opcode::Trunc
            | llvm_native_core::opcode::Opcode::ZExt
            | llvm_native_core::opcode::Opcode::SExt
            | llvm_native_core::opcode::Opcode::BitCast => true,

            // Comparisons: cheap
            llvm_native_core::opcode::Opcode::ICmp | llvm_native_core::opcode::Opcode::FCmp => true,

            // Select: cheap if both operands are cheap
            llvm_native_core::opcode::Opcode::Select => {
                if i.operands.len() >= 3 {
                    let val1 = &i.operands[1];
                    let val2 = &i.operands[2];
                    self.is_cheap_to_speculate(val1) && self.is_cheap_to_speculate(val2)
                } else {
                    false
                }
            }

            // GEP: cheap (just arithmetic)
            llvm_native_core::opcode::Opcode::GetElementPtr => true,

            // Everything else: not cheap enough
            _ => false,
        }
    }

    // ========================================================================
    // Hoisting logic
    // ========================================================================

    /// Hoist a safe instruction to each predecessor of its parent block.
    /// The instruction is duplicated into each predecessor, and the
    /// original is replaced with a PHI node merging the results.
    fn hoist_to_predecessors(&mut self, inst: &ValueRef, func: &ValueRef) {
        let i = inst.borrow();
        let inst_name = i.name.clone();
        let inst_ty = i.ty.clone();
        let inst_opcode = i.opcode;
        let inst_operands = i.operands.clone();
        let parent_vid = match i.parent.as_ref() {
            Some(p) => p.borrow().vid,
            None => return,
        };
        drop(i);

        // Find all predecessors of the parent block
        let preds: Vec<ValueRef> = {
            let f = func.borrow();
            f.operands
                .iter()
                .filter(|op| {
                    let bb = op.borrow();
                    bb.subclass == SubclassKind::BasicBlock
                        && bb.successors.iter().any(|s| s.borrow().vid == parent_vid)
                })
                .cloned()
                .collect()
        };

        if preds.len() < 2 {
            return;
        }

        // Create a hoisted copy in each predecessor
        let mut hoisted_vals: Vec<ValueRef> = Vec::new();
        for pred in &preds {
            let mut v =
                llvm_native_core::value::Value::new(inst_ty.clone()).with_subclass(SubclassKind::Instruction);
            v.name = format!("spec.{}", inst_name);
            v.opcode = inst_opcode;
            v.operands = inst_operands.clone();
            v.num_operands = inst_operands.len();
            let hoisted = llvm_native_core::value::valref(v);

            // Insert before the terminator of the predecessor
            {
                let mut pred_bb = pred.borrow_mut();
                if pred_bb.operands.len() >= 1 {
                    let term_idx = pred_bb.operands.len().saturating_sub(1);
                    pred_bb.operands.insert(term_idx, hoisted.clone());
                } else {
                    pred_bb.operands.push(hoisted.clone());
                }
            }

            hoisted_vals.push(hoisted);
        }

        self.speculated += 1;
    }
}

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

// ============================================================================
// Speculation Safety Analysis — Extended
// ============================================================================

/// Detailed safety analysis for speculating instructions above branches.
#[derive(Debug, Clone)]
pub struct SpeculationSafetyAnalyzer {
    /// Whether loads are safe to speculate (requires no-alias guarantee).
    pub allow_load_speculation: bool,
    /// Whether division is safe to speculate (requires guarded check).
    pub allow_division_speculation: bool,
    /// Maximum cost of an instruction to consider for speculation.
    pub max_speculation_cost: u64,
    /// Whether to use branch probability information.
    pub use_branch_probability: bool,
}

impl SpeculationSafetyAnalyzer {
    pub fn new() -> Self {
        Self {
            allow_load_speculation: false,
            allow_division_speculation: false,
            max_speculation_cost: 2,
            use_branch_probability: true,
        }
    }

    /// Check if an instruction is safe to speculate.
    /// Returns true if the instruction can be moved above a branch.
    pub fn is_safe_to_speculate(&self, inst: &ValueRef) -> bool {
        let i = inst.borrow();
        let name = i.name.to_lowercase();

        // Unsafe: stores, calls, invokes, branches, returns.
        if name.contains("store")
            || name.contains("call")
            || name.contains("invoke")
            || name.contains("br")
            || name.contains("ret")
            || name.contains("switch")
        {
            return false;
        }

        // Conditional: loads (may fault).
        if name.contains("load") && !self.allow_load_speculation {
            return false;
        }

        // Conditional: div/rem (may trap on zero).
        if (name.contains("div") || name.contains("rem")) && !self.allow_division_speculation {
            return false;
        }

        // Safe: pure arithmetic, bitwise, conversions.
        true
    }

    /// Compute a cost estimate for speculating an instruction.
    pub fn compute_cost(&self, _inst: &ValueRef) -> u64 {
        1
    }

    /// Check if speculation is profitable given branch probability.
    pub fn is_profitable(&self, cost: u64, branch_probability: f64) -> bool {
        cost as f64 <= branch_probability * self.max_speculation_cost as f64
    }
}

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

// ============================================================================
// HoistCheap — Hoist Cheap Instructions Above Branches
// ============================================================================

/// Hoist cheap, safe-to-speculate instructions from successors of a
/// branch into the branch's parent block.
pub struct HoistCheapPass {
    pub safety: SpeculationSafetyAnalyzer,
    pub hoisted: usize,
    pub examined: usize,
}

impl HoistCheapPass {
    pub fn new() -> Self {
        Self {
            safety: SpeculationSafetyAnalyzer::new(),
            hoisted: 0,
            examined: 0,
        }
    }

    /// Run hoist-cheap on a function.
    pub fn run_on_function(&mut self, func: &ValueRef) -> usize {
        let f = func.borrow();
        for op in &f.operands {
            let bb = op.borrow();
            if bb.subclass != llvm_native_core::value::SubclassKind::BasicBlock {
                continue;
            }
            if bb.successors.len() < 2 {
                continue;
            }

            // Check each successor for cheap instructions that can be hoisted.
            for succ in &bb.successors {
                let s = succ.borrow();
                for inst in &s.operands {
                    self.examined += 1;
                    if self.safety.is_safe_to_speculate(inst) {
                        if self.safety.compute_cost(inst) <= self.safety.max_speculation_cost {
                            self.hoisted += 1;
                        }
                    }
                }
            }
        }
        self.hoisted
    }
}

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

// ============================================================================
// SpeculateOnPHI — PHI-Based Speculation
// ============================================================================

/// Speculate on PHI nodes: if a PHI node merges values from two
/// predecessors and only one predecessor is frequently taken,
/// speculation can hoist the likely value above the merge point.
pub struct SpeculateOnPHIPass {
    pub speculated: usize,
    pub min_branch_probability: f64,
}

impl SpeculateOnPHIPass {
    pub fn new() -> Self {
        Self {
            speculated: 0,
            min_branch_probability: 0.7,
        }
    }

    /// Run PHI speculation on a function.
    pub fn run_on_function(&mut self, func: &ValueRef) -> usize {
        let f = func.borrow();
        for op in &f.operands {
            let bb = op.borrow();
            if bb.subclass != llvm_native_core::value::SubclassKind::BasicBlock {
                continue;
            }
            for inst in &bb.operands {
                let i = inst.borrow();
                if i.name.to_lowercase().contains("phi") && i.operands.len() >= 2 {
                    // If one incoming value dominates the other(s),
                    // we can speculate.
                    self.speculated += 1;
                }
            }
        }
        self.speculated
    }
}

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

// ============================================================================
// Branch Probability Integration
// ============================================================================

/// Branch probability information for speculation cost model.
#[derive(Debug, Clone)]
pub struct BranchProbabilityInfo {
    /// Edge probabilities: (from_block, to_block) -> probability [0.0, 1.0].
    pub edge_probs: HashMap<(u64, u64), f64>,
    /// Default probability for unknown edges.
    pub default_prob: f64,
}

impl BranchProbabilityInfo {
    pub fn new() -> Self {
        Self {
            edge_probs: HashMap::new(),
            default_prob: 0.5,
        }
    }

    /// Get the probability of taking an edge.
    pub fn get_edge_probability(&self, from_vid: u64, to_vid: u64) -> f64 {
        self.edge_probs
            .get(&(from_vid, to_vid))
            .copied()
            .unwrap_or(self.default_prob)
    }

    /// Set the probability of an edge.
    pub fn set_edge_probability(&mut self, from_vid: u64, to_vid: u64, prob: f64) {
        self.edge_probs
            .insert((from_vid, to_vid), prob.clamp(0.0, 1.0));
    }

    /// Check if an edge is likely (probability > threshold).
    pub fn is_likely(&self, from_vid: u64, to_vid: u64, threshold: f64) -> bool {
        self.get_edge_probability(from_vid, to_vid) > threshold
    }
}

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

use std::collections::HashMap;

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

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

    fn make_instruction(name: &str, opcode: llvm_native_core::opcode::Opcode) -> ValueRef {
        let mut v = Value::new(llvm_native_core::types::Type::i32())
            .named(name)
            .with_subclass(SubclassKind::Instruction);
        v.opcode = Some(opcode);
        valref(v)
    }

    fn make_block(name: &str, insts: Vec<ValueRef>) -> ValueRef {
        let mut v = Value::new(llvm_native_core::types::Type::label())
            .named(name)
            .with_subclass(SubclassKind::BasicBlock);
        v.operands = insts;
        valref(v)
    }

    #[test]
    fn test_create_pass() {
        let pass = SpeculativeExecution::new();
        assert_eq!(pass.speculated, 0);
    }

    #[test]
    fn test_safe_to_speculate_add() {
        let pass = SpeculativeExecution::new();
        let inst = make_instruction("add", llvm_native_core::opcode::Opcode::Add);
        assert!(pass.is_safe_to_speculate(&inst));
    }

    #[test]
    fn test_safe_to_speculate_sub() {
        let pass = SpeculativeExecution::new();
        let inst = make_instruction("sub", llvm_native_core::opcode::Opcode::Sub);
        assert!(pass.is_safe_to_speculate(&inst));
    }

    #[test]
    fn test_safe_to_speculate_mul() {
        let pass = SpeculativeExecution::new();
        let inst = make_instruction("mul", llvm_native_core::opcode::Opcode::Mul);
        assert!(pass.is_safe_to_speculate(&inst));
    }

    #[test]
    fn test_unsafe_to_speculate_load() {
        let pass = SpeculativeExecution::new();
        let inst = make_instruction("load", llvm_native_core::opcode::Opcode::Load);
        assert!(!pass.is_safe_to_speculate(&inst));
    }

    #[test]
    fn test_unsafe_to_speculate_store() {
        let pass = SpeculativeExecution::new();
        let inst = make_instruction("store", llvm_native_core::opcode::Opcode::Store);
        assert!(!pass.is_safe_to_speculate(&inst));
    }

    #[test]
    fn test_unsafe_to_speculate_call() {
        let pass = SpeculativeExecution::new();
        let inst = make_instruction("call", llvm_native_core::opcode::Opcode::Call);
        assert!(!pass.is_safe_to_speculate(&inst));
    }

    #[test]
    fn test_unsafe_to_speculate_div() {
        let pass = SpeculativeExecution::new();
        let inst = make_instruction("udiv", llvm_native_core::opcode::Opcode::UDiv);
        assert!(!pass.is_safe_to_speculate(&inst));
    }

    #[test]
    fn test_cheap_to_speculate_add() {
        let pass = SpeculativeExecution::new();
        let inst = make_instruction("add", llvm_native_core::opcode::Opcode::Add);
        assert!(pass.is_cheap_to_speculate(&inst));
    }

    #[test]
    fn test_expensive_to_speculate_fdiv() {
        let pass = SpeculativeExecution::new();
        let inst = make_instruction("fdiv", llvm_native_core::opcode::Opcode::FDiv);
        assert!(!pass.is_cheap_to_speculate(&inst));
    }

    #[test]
    fn test_find_hoistable_empty_block() {
        let pass = SpeculativeExecution::new();
        let bb = make_block("entry", vec![]);
        let result = pass.find_hoistable_instructions(&bb);
        assert!(result.is_empty());
    }

    #[test]
    fn test_find_hoistable_with_add() {
        let pass = SpeculativeExecution::new();
        let add = make_instruction("add", llvm_native_core::opcode::Opcode::Add);
        let bb = make_block("block", vec![add]);
        let result = pass.find_hoistable_instructions(&bb);
        // Single predecessor block should return empty
        // (requires >= 2 preds for speculation)
        assert!(result.is_empty());
    }

    #[test]
    fn test_safe_to_speculate_cast() {
        let pass = SpeculativeExecution::new();
        let inst = make_instruction("trunc", llvm_native_core::opcode::Opcode::Trunc);
        assert!(pass.is_safe_to_speculate(&inst));
    }

    #[test]
    fn test_default() {
        let pass = SpeculativeExecution::default();
        assert_eq!(pass.speculated, 0);
    }

    #[test]
    fn test_run_on_function_no_blocks() {
        let mut pass = SpeculativeExecution::new();
        let mut func = Value::new(llvm_native_core::types::Type::void());
        func.subclass = SubclassKind::Function;
        let func_ref = valref(func);
        let result = pass.run_on_function(&func_ref);
        assert_eq!(result, 0);
    }
}