llvm-native-core-ext 0.1.0

//! LLVM MachineBlockPlacement — optimal basic block layout using branch
//! probabilities to minimize taken branches and improve I-cache locality.
//! Clean-room behavioral reconstruction.
//!
//! This pass reorders the basic blocks within a function to minimize the
//! number of taken conditional branches, reduce branch mispredictions,
//! and improve instruction cache utilization. It is equivalent to the
//! MachineBlockPlacement pass in LLVM's code generation pipeline but
//! operates at the IR level for planning purposes.
//!
//! Algorithm (based on Pettis-Hansen and LLVM's approach):
//!   1. Build a control-flow graph (CFG) with edge weights
//!   2. Compute branch probabilities using static heuristics or profile data
//!   3. Build a chain of blocks using a greedy algorithm: start with the
//!      entry block, then repeatedly append the most likely successor
//!      that hasn't been placed yet
//!   4. Handle fall-through optimization: if a block ends with an
//!      unconditional branch, try to place the target immediately after
//!   5. Merge small blocks where profitable to reduce branch overhead
//!   6. Align hot blocks to cache-line boundaries where beneficial

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

// ============================================================================
// Machine Block Placement Pass
// ============================================================================

/// MachineBlockPlacement — reorders basic blocks for optimal layout.
pub struct MachineBlockPlacement {
    /// Number of basic blocks reordered.
    pub reordered: usize,
}

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

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

    /// Run block placement on a function. Returns the number of basic
    /// blocks whose position was changed.
    pub fn run_on_function(&mut self, func: &ValueRef) -> usize {
        self.reordered = 0;

        // Step 1: Compute branch probabilities
        let _probabilities = self.compute_branch_probabilities(func);

        // Step 2: Compute block frequencies
        let _frequencies = self.compute_block_frequencies(func);

        // Step 3: Compute the optimal layout
        let layout = self.compute_optimal_layout(func);

        // Step 4: Reorder the blocks
        self.reorder_blocks(func, &layout);

        // Step 5: Merge blocks where profitable
        self.merge_blocks(func);

        // Step 6: Align hot blocks
        self.align_blocks(func);

        self.reordered
    }

    // ========================================================================
    // Branch probability computation
    // ========================================================================

    /// Compute branch probabilities for all edges in the CFG.
    /// Returns a map from (source_block_vid, target_block_vid) to
    /// probability (0.0 to 1.0).
    fn compute_branch_probabilities(&self, func: &ValueRef) -> HashMap<(u64, u64), f64> {
        let mut probs: HashMap<(u64, u64), f64> = HashMap::new();
        let f = func.borrow();

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

            let succ_count = bb.successors.len();
            if succ_count == 0 {
                continue;
            }

            if succ_count == 1 {
                // Unconditional branch: 100% probability
                let target_vid = bb.successors[0].borrow().vid;
                probs.insert((bb.vid, target_vid), 1.0);
            } else if succ_count == 2 {
                // Conditional branch with two successors
                let succ0_vid = bb.successors[0].borrow().vid;
                let succ1_vid = bb.successors[1].borrow().vid;

                // Heuristic: check if one edge is a back edge
                let succ0_is_backedge = succ0_vid < bb.vid;
                let succ1_is_backedge = succ1_vid < bb.vid;

                if succ0_is_backedge && !succ1_is_backedge {
                    probs.insert((bb.vid, succ0_vid), 0.90);
                    probs.insert((bb.vid, succ1_vid), 0.10);
                } else if succ1_is_backedge && !succ0_is_backedge {
                    probs.insert((bb.vid, succ0_vid), 0.10);
                    probs.insert((bb.vid, succ1_vid), 0.90);
                } else {
                    probs.insert((bb.vid, succ0_vid), 0.70);
                    probs.insert((bb.vid, succ1_vid), 0.30);
                }
            } else {
                // Multiple successors (switch): distribute evenly
                let weight = 1.0 / succ_count as f64;
                for succ in &bb.successors {
                    probs.insert((bb.vid, succ.borrow().vid), weight);
                }
            }
        }

        probs
    }

    // ========================================================================
    // Block frequency computation
    // ========================================================================

    /// Compute the relative execution frequency of each basic block.
    /// Returns a vector of frequencies indexed by block position.
    fn compute_block_frequencies(&self, func: &ValueRef) -> Vec<f64> {
        let f = func.borrow();
        let mut block_list: Vec<ValueRef> = Vec::new();
        let mut block_index: HashMap<u64, usize> = HashMap::new();

        for op in &f.operands {
            if op.borrow().subclass == SubclassKind::BasicBlock {
                block_index.insert(op.borrow().vid, block_list.len());
                block_list.push(op.clone());
            }
        }

        let n = block_list.len();
        if n == 0 {
            return Vec::new();
        }

        let mut freqs = vec![0.0f64; n];
        freqs[0] = 1.0;

        let probs = self.compute_branch_probabilities(func);
        let mut worklist: VecDeque<usize> = VecDeque::new();
        worklist.push_back(0);
        let mut in_queue = vec![false; n];
        in_queue[0] = true;

        let max_iterations = n * 10;
        let mut iterations = 0;

        while let Some(bb_idx) = worklist.pop_front() {
            in_queue[bb_idx] = false;
            iterations += 1;
            if iterations > max_iterations {
                break;
            }

            let bb = &block_list[bb_idx];
            let bb_data = bb.borrow();
            let bb_freq = freqs[bb_idx];

            let succ_count = bb_data.successors.len();
            if succ_count == 0 {
                continue;
            }

            for succ in &bb_data.successors {
                let succ_vid = succ.borrow().vid;
                let edge_prob = probs
                    .get(&(bb_data.vid, succ_vid))
                    .copied()
                    .unwrap_or(1.0 / succ_count as f64);

                if let Some(&succ_idx) = block_index.get(&succ_vid) {
                    let incoming = bb_freq * edge_prob;
                    freqs[succ_idx] += incoming;

                    if !in_queue[succ_idx] {
                        worklist.push_back(succ_idx);
                        in_queue[succ_idx] = true;
                    }
                }
            }
        }

        // Normalize frequencies
        let sum: f64 = freqs.iter().sum();
        if sum > 0.0 {
            for f in &mut freqs {
                *f /= sum;
            }
        }

        freqs
    }

    // ========================================================================
    // Optimal layout computation
    // ========================================================================

    /// Compute the optimal block layout using a greedy chain-building
    /// approach. Returns a vector of block indices.
    fn compute_optimal_layout(&self, func: &ValueRef) -> Vec<usize> {
        let f = func.borrow();
        let mut block_list: Vec<ValueRef> = Vec::new();
        let mut block_index: HashMap<u64, usize> = HashMap::new();

        for (idx, op) in f.operands.iter().enumerate() {
            if op.borrow().subclass == SubclassKind::BasicBlock {
                block_index.insert(op.borrow().vid, idx);
                block_list.push(op.clone());
            }
        }

        let n = block_list.len();
        if n == 0 {
            return Vec::new();
        }

        let probs = self.compute_branch_probabilities(func);
        let freqs = self.compute_block_frequencies(func);

        let mut placed: HashSet<usize> = HashSet::new();
        let mut layout: Vec<usize> = Vec::new();

        let mut frontier: BinaryHeap<BlockPriority> = BinaryHeap::new();
        frontier.push(BlockPriority {
            freq: freqs[0],
            idx: 0,
        });

        while let Some(BlockPriority { idx, .. }) = frontier.pop() {
            if placed.contains(&idx) {
                continue;
            }

            layout.push(idx);
            placed.insert(idx);

            let bb = &block_list[idx];
            let bb_data = bb.borrow();
            for succ in &bb_data.successors {
                let succ_vid = succ.borrow().vid;
                if let Some(&succ_idx) = block_index.get(&succ_vid) {
                    if !placed.contains(&succ_idx) {
                        let edge_prob = probs.get(&(bb_data.vid, succ_vid)).copied().unwrap_or(0.5);
                        let priority = freqs[succ_idx] * edge_prob;
                        frontier.push(BlockPriority {
                            freq: priority,
                            idx: succ_idx,
                        });
                    }
                }
            }
        }

        // Any blocks not yet placed go at the end
        for idx in 0..n {
            if !placed.contains(&idx) {
                layout.push(idx);
            }
        }

        layout
    }

    // ========================================================================
    // Block reordering
    // ========================================================================

    /// Reorder the basic blocks in the function according to the
    /// computed layout.
    fn reorder_blocks(&mut self, func: &ValueRef, layout: &[usize]) {
        if layout.is_empty() {
            return;
        }

        let mut f = func.borrow_mut();

        let mut blocks: Vec<ValueRef> = Vec::new();
        let mut others: Vec<ValueRef> = Vec::new();

        for op in f.operands.drain(..) {
            if op.borrow().subclass == SubclassKind::BasicBlock {
                blocks.push(op);
            } else {
                others.push(op);
            }
        }

        let mut reordered: Vec<ValueRef> = Vec::new();
        for &idx in layout {
            if idx < blocks.len() {
                reordered.push(blocks[idx].clone());
            }
        }

        // Append any blocks not in layout
        for (i, block) in blocks.iter().enumerate() {
            if !layout.contains(&i) {
                reordered.push(block.clone());
            }
        }

        let mut changed = 0;
        for (i, block) in reordered.iter().enumerate() {
            if i < blocks.len() && block.borrow().vid != blocks[i].borrow().vid {
                changed += 1;
            }
        }
        self.reordered = changed;

        f.operands = others;
        f.operands.extend(reordered);
    }

    // ========================================================================
    // Block merging
    // ========================================================================

    /// Merge consecutive blocks where the predecessor ends with an
    /// unconditional branch to the successor and merging is profitable.
    fn merge_blocks(&mut self, func: &ValueRef) {
        let mut f = func.borrow_mut();
        let n = f.operands.len();

        if n < 2 {
            return;
        }

        let mut to_remove: HashSet<usize> = HashSet::new();

        for i in 0..n {
            if to_remove.contains(&i) {
                continue;
            }

            let bb = &f.operands[i];
            let bb_data = bb.borrow();

            let is_uncond_br = {
                let last_inst = bb_data.operands.last();
                last_inst.map_or(false, |inst| {
                    let ii = inst.borrow();
                    ii.opcode == Some(llvm_native_core::opcode::Opcode::Br)
                        && ii.operands.len() == 1
                        && ii.name.to_lowercase().contains("br")
                        && !ii.name.to_lowercase().contains("cond")
                })
            };

            if is_uncond_br && bb_data.successors.len() == 1 {
                let target_vid = bb_data.successors[0].borrow().vid;

                if i + 1 < n {
                    let next_bb = &f.operands[i + 1];
                    if next_bb.borrow().vid == target_vid {
                        let target_data = next_bb.borrow();
                        let target_insts: Vec<ValueRef> = target_data
                            .operands
                            .iter()
                            .filter(|inst| !inst.borrow().is_terminator())
                            .cloned()
                            .collect();

                        drop(bb_data);
                        drop(target_data);

                        {
                            let mut bb_mut = bb.borrow_mut();
                            let insert_pos = bb_mut.operands.len().saturating_sub(1);
                            for inst in target_insts.into_iter().rev() {
                                bb_mut.operands.insert(insert_pos, inst);
                            }
                        }

                        to_remove.insert(i + 1);
                    }
                }
            }
        }

        if !to_remove.is_empty() {
            let mut new_operands: Vec<ValueRef> = Vec::new();
            for (i, op) in f.operands.iter().enumerate() {
                if !to_remove.contains(&i) {
                    new_operands.push(op.clone());
                }
            }
            f.operands = new_operands;
        }
    }

    // ========================================================================
    // Block alignment
    // ========================================================================

    /// Align hot basic blocks to cache-line boundaries.
    fn align_blocks(&mut self, func: &ValueRef) {
        let freqs = self.compute_block_frequencies(func);
        let f = func.borrow();

        let total: f64 = freqs.iter().sum();
        let threshold = if total > 0.0 { 0.1 * total } else { 0.0 };

        for (idx, _op) in f.operands.iter().enumerate() {
            if idx < freqs.len() && freqs[idx] > threshold {
                // Hot block identified — would be aligned in codegen
            }
        }
    }
}

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

// ============================================================================
// Helper types
// ============================================================================

/// Priority queue entry for greedy block layout.
#[derive(PartialEq)]
struct BlockPriority {
    freq: f64,
    idx: usize,
}

impl Eq for BlockPriority {}

impl PartialOrd for BlockPriority {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        self.freq.partial_cmp(&other.freq)
    }
}

impl Ord for BlockPriority {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.partial_cmp(other).unwrap_or(std::cmp::Ordering::Equal)
    }
}

// ============================================================================
// Branch Probability Information
// ============================================================================

/// BranchProbability represents the likelihood of taking an edge.
/// Values are in the range [0, PROB_MAX].
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct BranchProbability {
    /// Numerator of the probability fraction.
    pub num: u32,
    /// Denominator (maximum value).
    pub den: u32,
}

impl BranchProbability {
    /// Maximum probability value.
    pub const PROB_MAX: u32 = 0x8000_0000;

    /// Create a probability with a given numerator.
    pub fn new(num: u32) -> Self {
        Self {
            num,
            den: Self::PROB_MAX,
        }
    }

    /// Create a probability from a fraction.
    pub fn from_fraction(num: u32, den: u32) -> Self {
        if den == 0 {
            Self { num: 0, den: 1 }
        } else {
            Self {
                num,
                den: den.max(1),
            }
        }
    }

    /// The "very likely" probability (80%).
    pub fn very_likely() -> Self {
        Self::new((Self::PROB_MAX as f64 * 0.8) as u32)
    }

    /// The "likely" probability (60%).
    pub fn likely() -> Self {
        Self::new((Self::PROB_MAX as f64 * 0.6) as u32)
    }

    /// The "unlikely" probability (20%).
    pub fn unlikely() -> Self {
        Self::new((Self::PROB_MAX as f64 * 0.2) as u32)
    }

    /// The "unknown" probability (50%).
    pub fn unknown() -> Self {
        Self::new(Self::PROB_MAX / 2)
    }

    /// Get as a float in [0.0, 1.0].
    pub fn as_f64(self) -> f64 {
        if self.den == 0 {
            0.0
        } else {
            self.num as f64 / self.den as f64
        }
    }

    /// Check if this branch is taken more often than not.
    pub fn is_taken(self) -> bool {
        self.num > self.den / 2
    }

    /// Check if this is a highly predictable branch.
    pub fn is_biased(self) -> bool {
        self.num > (self.den as f64 * 0.9) as u32 || self.num < (self.den as f64 * 0.1) as u32
    }

    /// Scale by a factor.
    pub fn scale(self, factor: f64) -> Self {
        Self {
            num: (self.num as f64 * factor) as u32,
            den: self.den,
        }
    }

    /// Get the complement probability (1 - this).
    pub fn complement(self) -> Self {
        Self {
            num: self.den.saturating_sub(self.num),
            den: self.den,
        }
    }
}

impl Default for BranchProbability {
    fn default() -> Self {
        Self::unknown()
    }
}

/// Edge weight combining branch probability with frequency.
#[derive(Debug, Clone, Copy)]
pub struct EdgeWeight {
    /// Branch probability for this edge.
    pub prob: BranchProbability,
    /// Execution frequency of the source block.
    pub freq: f64,
}

impl EdgeWeight {
    /// Create a new edge weight.
    pub fn new(prob: BranchProbability, freq: f64) -> Self {
        Self { prob, freq }
    }

    /// Get the expected number of times this edge is taken.
    pub fn expected_count(&self) -> f64 {
        self.freq * self.prob.as_f64()
    }
}

// ============================================================================
// BlockChain — Ordered sequence of basic blocks for layout
// ============================================================================

/// A BlockChain is an ordered sequence of basic blocks that are placed
/// consecutively in memory for optimal fallthrough behavior.
#[derive(Debug, Clone)]
pub struct BlockChain {
    /// Blocks in this chain, in layout order.
    pub blocks: Vec<usize>,
    /// Block indices that are part of this chain.
    pub member_set: HashSet<usize>,
    /// Total execution frequency of this chain.
    pub total_freq: f64,
    /// Whether this chain contains the entry block.
    pub is_entry: bool,
    /// Whether this chain is a loop.
    pub is_loop: bool,
}

impl BlockChain {
    /// Create a new chain starting with a single block.
    pub fn new(block_idx: usize, freq: f64) -> Self {
        let mut member_set = HashSet::new();
        member_set.insert(block_idx);
        Self {
            blocks: vec![block_idx],
            member_set,
            total_freq: freq,
            is_entry: block_idx == 0,
            is_loop: false,
        }
    }

    /// Check if a block is a member of this chain.
    pub fn contains(&self, block_idx: usize) -> bool {
        self.member_set.contains(&block_idx)
    }

    /// Get the first block in the chain.
    pub fn head(&self) -> Option<usize> {
        self.blocks.first().copied()
    }

    /// Get the last block in the chain.
    pub fn tail(&self) -> Option<usize> {
        self.blocks.last().copied()
    }

    /// Add a block to the end of the chain.
    pub fn append(&mut self, block_idx: usize, freq: f64) {
        if !self.member_set.contains(&block_idx) {
            self.blocks.push(block_idx);
            self.member_set.insert(block_idx);
            self.total_freq += freq;
        }
    }

    /// Prepend a block to the beginning of the chain.
    pub fn prepend(&mut self, block_idx: usize, freq: f64) {
        if !self.member_set.contains(&block_idx) {
            self.blocks.insert(0, block_idx);
            self.member_set.insert(block_idx);
            self.total_freq += freq;
            if block_idx == 0 {
                self.is_entry = true;
            }
        }
    }

    /// Merge another chain at the end of this chain.
    pub fn merge_chain(&mut self, other: &BlockChain) {
        for &block in &other.blocks {
            if !self.member_set.contains(&block) {
                self.blocks.push(block);
                self.member_set.insert(block);
            }
        }
        self.total_freq += other.total_freq;
    }

    /// Check if this chain is empty.
    pub fn is_empty(&self) -> bool {
        self.blocks.is_empty()
    }

    /// Get the number of blocks in this chain.
    pub fn len(&self) -> usize {
        self.blocks.len()
    }
}

/// BlockChainBuilder constructs chains using a greedy trace-formation algorithm.
pub struct BlockChainBuilder {
    /// All chains formed.
    pub chains: Vec<BlockChain>,
    /// Mapping from block index to its chain index.
    pub block_to_chain: HashMap<usize, usize>,
    /// Edge weights between blocks.
    pub edges: HashMap<(usize, usize), EdgeWeight>,
    /// Block frequencies.
    pub frequencies: Vec<f64>,
}

impl BlockChainBuilder {
    /// Create a new chain builder.
    pub fn new() -> Self {
        Self {
            chains: Vec::new(),
            block_to_chain: HashMap::new(),
            edges: HashMap::new(),
            frequencies: Vec::new(),
        }
    }

    /// Build chains from block frequencies and edge weights.
    pub fn build_chains(
        &mut self,
        freqs: &[f64],
        edge_weights: &HashMap<(usize, usize), EdgeWeight>,
        num_blocks: usize,
    ) {
        self.chains.clear();
        self.block_to_chain.clear();
        self.frequencies = freqs.to_vec();
        self.edges = edge_weights.clone();

        if num_blocks == 0 {
            return;
        }

        // Phase 1: Create singleton chains for each block
        for i in 0..num_blocks {
            let freq = freqs.get(i).copied().unwrap_or(0.0);
            let chain = BlockChain::new(i, freq);
            let chain_idx = self.chains.len();
            self.chains.push(chain);
            self.block_to_chain.insert(i, chain_idx);
        }

        // Phase 2: Merge chains greedily based on edge weights
        let mut changed = true;
        while changed {
            changed = false;

            // Find the best edge to merge across chains
            let mut best_edge: Option<((usize, usize), f64)> = None;

            for ((src, dst), weight) in &self.edges {
                if let (Some(&chain_src), Some(&chain_dst)) =
                    (self.block_to_chain.get(src), self.block_to_chain.get(dst))
                {
                    if chain_src == chain_dst {
                        continue; // Already in same chain
                    }

                    let src_chain = &self.chains[chain_src];
                    let dst_chain = &self.chains[chain_dst];

                    // Only merge if src is at the tail of its chain
                    // and dst is at the head of its chain
                    if src_chain.tail() != Some(*src) || dst_chain.head() != Some(*dst) {
                        continue;
                    }

                    let score = weight.expected_count();
                    if best_edge.map(|(_, s)| score > s).unwrap_or(true) {
                        best_edge = Some(((*src, *dst), score));
                    }
                }
            }

            if let Some(((src, dst), _score)) = best_edge {
                let chain_src = self.block_to_chain[&src];
                let chain_dst = self.block_to_chain[&dst];

                // Merge dst chain into src chain
                let dst_blocks = self.chains[chain_dst].blocks.clone();
                for &block in &dst_blocks {
                    self.chains[chain_src]
                        .append(block, self.frequencies.get(block).copied().unwrap_or(0.0));
                    self.block_to_chain.insert(block, chain_src);
                }
                self.chains[chain_dst].blocks.clear();
                self.chains[chain_dst].member_set.clear();
                changed = true;
            }
        }

        // Phase 3: Remove empty chains
        self.chains.retain(|c| !c.is_empty());
    }

    /// Get the final block layout order from built chains.
    pub fn get_layout_order(&self) -> Vec<usize> {
        // Sort chains: entry chain first, then by total frequency descending
        let mut chain_order: Vec<(usize, f64, bool)> = self
            .chains
            .iter()
            .enumerate()
            .filter(|(_, c)| !c.is_empty())
            .map(|(i, c)| (i, c.total_freq, c.is_entry))
            .collect();

        chain_order.sort_by(|a, b| {
            b.2.cmp(&a.2) // Entry chains first
                .then_with(|| {
                    b.1.partial_cmp(&a.1)
                        .unwrap_or(std::cmp::Ordering::Equal)
                })
        });

        let mut order = Vec::new();
        for (chain_idx, _, _) in &chain_order {
            order.extend_from_slice(&self.chains[*chain_idx].blocks);
        }
        order
    }

    /// Print chain information for debugging.
    pub fn print_chains(&self) {
        for (i, chain) in self.chains.iter().enumerate() {
            if chain.is_empty() {
                continue;
            }
            eprintln!(
                "Chain {}: {:?} (freq={:.2}, entry={})",
                i, chain.blocks, chain.total_freq, chain.is_entry
            );
        }
    }
}

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

// ============================================================================
// Fallthrough Optimization
// ============================================================================

/// FallthroughOptimizer ensures that conditional branches have their
/// most likely target as the fallthrough (next block in layout),
/// which improves branch prediction and reduces taken branches.
pub struct FallthroughOptimizer {
    /// Number of blocks where fallthrough was optimized.
    pub optimized: usize,
    /// Number of condition reversals performed.
    pub conditions_reversed: usize,
}

impl FallthroughOptimizer {
    /// Create a new FallthroughOptimizer.
    pub fn new() -> Self {
        Self {
            optimized: 0,
            conditions_reversed: 0,
        }
    }

    /// Optimize fallthrough for all blocks in layout order.
    pub fn run_on_function(
        &mut self,
        block_order: &[usize],
        freqs: &[f64],
        successors: &[Vec<usize>],
        branch_probs: &HashMap<(usize, usize), BranchProbability>,
    ) -> usize {
        self.optimized = 0;
        self.conditions_reversed = 0;

        for i in 0..block_order.len().saturating_sub(1) {
            let cur = block_order[i];
            let next = block_order[i + 1];

            let succs = successors.get(cur);
            if succs.is_none() || succs.unwrap().len() != 2 {
                continue;
            }

            let succs = succs.unwrap();
            let taken = succs[0];
            let not_taken = succs[1];

            // If the not-taken target is the next block, the fallthrough is
            // already optimal (branch is taken only when condition true).
            if not_taken == next {
                // Already optimal
                continue;
            }

            // If the taken target is the next block, we might want to reverse
            // the condition so the hot path falls through.
            if taken == next {
                let prob_taken = branch_probs.get(&(cur, taken)).copied().unwrap_or_default();

                // If the taken path is likely (>50%), keep as-is
                if prob_taken.is_taken() {
                    continue;
                }

                // Otherwise, reversing the condition would make the
                // more likely path fall through.
                self.conditions_reversed += 1;
                self.optimized += 1;
                continue;
            }

            // Neither successor is the next block — check if we should
            // reorder blocks to make the hot successor fall through.
            let prob_taken = branch_probs.get(&(cur, taken)).copied().unwrap_or_default();
            let prob_not_taken = branch_probs
                .get(&(cur, not_taken))
                .copied()
                .unwrap_or_default();

            if prob_taken.is_taken() && prob_taken.num > prob_not_taken.num {
                // The taken edge is more likely — keep as-is;
                // the block layout may be adjusted elsewhere.
            } else if prob_not_taken.is_taken() {
                // The not-taken edge is more likely — reverse to make it
                // fall through.
                self.conditions_reversed += 1;
                self.optimized += 1;
            }
        }

        self.optimized
    }
}

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

// ============================================================================
// Loop Alignment
// ============================================================================

/// LoopAlignment ensures loop headers are aligned to cache-line boundaries
/// for better I-cache utilization and branch prediction.
#[derive(Debug, Clone)]
pub struct LoopAlignment {
    /// Loop header block index.
    pub header: usize,
    /// Whether alignment is recommended.
    pub should_align: bool,
    /// Alignment boundary (e.g., 16, 32, 64 bytes).
    pub alignment: u32,
    /// Loop nesting depth.
    pub depth: u32,
    /// Estimated loop trip count.
    pub trip_count: f64,
}

impl LoopAlignment {
    /// Create a new loop alignment hint.
    pub fn new(header: usize) -> Self {
        Self {
            header,
            should_align: false,
            alignment: 16,
            depth: 1,
            trip_count: 0.0,
        }
    }

    /// Determine if alignment is beneficial for this loop.
    pub fn compute_alignment(&mut self, loop_freq: f64, total_freq: f64) {
        let ratio = if total_freq > 0.0 {
            loop_freq / total_freq
        } else {
            0.0
        };

        // Align loops that execute more than 1% of the time
        // and use larger alignment for hotter loops.
        if ratio > 0.01 {
            self.should_align = true;
            if ratio > 0.10 {
                self.alignment = 32;
            }
            if ratio > 0.25 {
                self.alignment = 64;
            }
        }

        self.trip_count = loop_freq;
    }
}

/// LoopAlignmentPass computes alignment recommendations for loops.
pub struct LoopAlignmentPass {
    /// Alignment recommendations per block.
    pub alignments: HashMap<usize, LoopAlignment>,
    /// Total alignment padding added.
    pub total_padding: u32,
}

impl LoopAlignmentPass {
    /// Create a new loop alignment pass.
    pub fn new() -> Self {
        Self {
            alignments: HashMap::new(),
            total_padding: 0,
        }
    }

    /// Compute alignment for all loops.
    pub fn run(&mut self, loop_headers: &[usize], loop_freqs: &[f64], total_freq: f64) {
        self.alignments.clear();
        self.total_padding = 0;

        for (depth, (&header, &freq)) in loop_headers.iter().zip(loop_freqs.iter()).enumerate() {
            let mut align = LoopAlignment::new(header);
            align.depth = (depth + 1) as u32;
            align.compute_alignment(freq, total_freq);

            if align.should_align {
                self.total_padding += align.alignment;
            }

            self.alignments.insert(header, align);
        }
    }

    /// Check if a block should be aligned.
    pub fn should_align(&self, block_idx: usize) -> bool {
        self.alignments
            .get(&block_idx)
            .map(|a| a.should_align)
            .unwrap_or(false)
    }

    /// Get the recommended alignment for a block.
    pub fn get_alignment(&self, block_idx: usize) -> u32 {
        self.alignments
            .get(&block_idx)
            .map(|a| a.alignment)
            .unwrap_or(1)
    }
}

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

// ============================================================================
// Top-Down Block Placement
// ============================================================================

/// TopDownPlacer builds the final block layout using a top-down greedy
/// approach that prioritizes the most frequent edges.
pub struct TopDownPlacer {
    /// The final block order.
    pub order: Vec<usize>,
    /// Blocks that have been placed.
    pub placed: HashSet<usize>,
}

impl TopDownPlacer {
    /// Create a new top-down placer.
    pub fn new() -> Self {
        Self {
            order: Vec::new(),
            placed: HashSet::new(),
        }
    }

    /// Place blocks using top-down greedy algorithm.
    pub fn place(
        &mut self,
        num_blocks: usize,
        successors: &[Vec<usize>],
        edge_weights: &HashMap<(usize, usize), EdgeWeight>,
    ) -> Vec<usize> {
        self.order.clear();
        self.placed.clear();

        if num_blocks == 0 {
            return Vec::new();
        }

        // Start with entry block (index 0)
        let mut worklist: VecDeque<usize> = VecDeque::new();
        worklist.push_back(0);

        while let Some(block) = worklist.pop_front() {
            if self.placed.contains(&block) {
                continue;
            }

            self.order.push(block);
            self.placed.insert(block);

            // Find the best successor that hasn't been placed
            if let Some(succs) = successors.get(block) {
                // Sort successors by edge weight descending
                let mut ranked: Vec<(usize, f64)> = succs
                    .iter()
                    .filter_map(|&s| {
                        if self.placed.contains(&s) {
                            None
                        } else {
                            let weight = edge_weights
                                .get(&(block, s))
                                .map(|w| w.expected_count())
                                .unwrap_or(0.0);
                            Some((s, weight))
                        }
                    })
                    .collect();

                ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

                for (succ, _weight) in &ranked {
                    if !self.placed.contains(succ) {
                        worklist.push_front(*succ);
                        break; // Only push the best one
                    }
                }

                // Push remaining successors to back
                for (succ, _weight) in ranked.iter().skip(1) {
                    if !self.placed.contains(succ) {
                        worklist.push_back(*succ);
                    }
                }
            }
        }

        // Place any remaining unplaced blocks at the end
        for i in 0..num_blocks {
            if !self.placed.contains(&i) {
                self.order.push(i);
                self.placed.insert(i);
            }
        }

        self.order.clone()
    }
}

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

// ============================================================================
// Cold Block Clustering
// ============================================================================

/// ColdBlockCluster groups rarely-executed blocks together at the end
/// of the function to improve I-cache utilization for the hot path.
#[derive(Debug, Clone)]
pub struct ColdBlockCluster {
    /// Blocks identified as cold.
    pub cold_blocks: HashSet<usize>,
    /// Cold execution threshold (fraction of max frequency).
    pub cold_threshold: f64,
    /// Number of cold blocks clustered.
    pub clustered: usize,
}

impl ColdBlockCluster {
    /// Create a new cold block cluster pass.
    pub fn new() -> Self {
        Self {
            cold_blocks: HashSet::new(),
            cold_threshold: 0.05,
            clustered: 0,
        }
    }

    /// Analyze block frequencies and identify cold blocks.
    pub fn analyze(&mut self, freqs: &[f64]) {
        self.cold_blocks.clear();

        if freqs.is_empty() {
            return;
        }

        let max_freq = freqs.iter().cloned().fold(0.0f64, f64::max);

        if max_freq <= 0.0 {
            return;
        }

        let threshold = max_freq * self.cold_threshold;

        for (i, &freq) in freqs.iter().enumerate() {
            if freq < threshold && i != 0 {
                self.cold_blocks.insert(i);
            }
        }
    }

    /// Reorder blocks to cluster cold blocks at the end.
    pub fn apply(&self, order: &mut [usize], _successors: &[Vec<usize>]) -> usize {
        let mut hot: Vec<usize> = Vec::new();
        let mut cold: Vec<usize> = Vec::new();

        for &block in order.iter() {
            if self.cold_blocks.contains(&block) && block != 0 {
                cold.push(block);
            } else {
                hot.push(block);
            }
        }

        let clustered = cold.len();

        // Rebuild order: hot first, then cold
        hot.append(&mut cold);
        let new_order = hot;

        // Copy back
        let len = order.len().min(new_order.len());
        order[..len].copy_from_slice(&new_order[..len]);

        clustered
    }

    /// Check if a block is cold.
    pub fn is_cold(&self, block_idx: usize) -> bool {
        self.cold_blocks.contains(&block_idx)
    }
}

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

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

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

    fn make_block_named(_name: &str) -> ValueRef {
        let mut v = Value::new(llvm_native_core::types::Type::label())
            .named(_name)
            .with_subclass(SubclassKind::BasicBlock);
        valref(v)
    }

    fn set_successors(block: &ValueRef, succs: Vec<ValueRef>) {
        let mut bb = block.borrow_mut();
        bb.successors = succs;
    }

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

    #[test]
    fn test_compute_branch_probabilities_empty() {
        let pass = MachineBlockPlacement::new();
        let mut func = Value::new(llvm_native_core::types::Type::void());
        func.subclass = SubclassKind::Function;
        let func_ref = valref(func);
        let probs = pass.compute_branch_probabilities(&func_ref);
        assert!(probs.is_empty());
    }

    #[test]
    fn test_compute_branch_probabilities_single_block() {
        let pass = MachineBlockPlacement::new();
        let bb = make_block_named("entry");

        let mut func = Value::new(llvm_native_core::types::Type::void());
        func.subclass = SubclassKind::Function;
        func.operands = vec![bb];
        let func_ref = valref(func);

        let probs = pass.compute_branch_probabilities(&func_ref);
        assert!(probs.is_empty());
    }

    #[test]
    fn test_compute_branch_probabilities_uncond_br() {
        let pass = MachineBlockPlacement::new();
        let bb1 = make_block_named("entry");
        let bb2 = make_block_named("target");
        set_successors(&bb1, vec![bb2.clone()]);

        let mut func = Value::new(llvm_native_core::types::Type::void());
        func.subclass = SubclassKind::Function;
        func.operands = vec![bb1];
        let func_ref = valref(func);

        let probs = pass.compute_branch_probabilities(&func_ref);
        assert!(!probs.is_empty());
    }

    #[test]
    fn test_compute_block_frequencies_empty() {
        let pass = MachineBlockPlacement::new();
        let mut func = Value::new(llvm_native_core::types::Type::void());
        func.subclass = SubclassKind::Function;
        let func_ref = valref(func);
        let freqs = pass.compute_block_frequencies(&func_ref);
        assert!(freqs.is_empty());
    }

    #[test]
    fn test_compute_block_frequencies_single() {
        let pass = MachineBlockPlacement::new();
        let bb = make_block_named("entry");

        let mut func = Value::new(llvm_native_core::types::Type::void());
        func.subclass = SubclassKind::Function;
        func.operands = vec![bb];
        let func_ref = valref(func);

        let freqs = pass.compute_block_frequencies(&func_ref);
        assert_eq!(freqs.len(), 1);
        assert!(freqs[0] > 0.0);
    }

    #[test]
    fn test_compute_optimal_layout_empty() {
        let pass = MachineBlockPlacement::new();
        let mut func = Value::new(llvm_native_core::types::Type::void());
        func.subclass = SubclassKind::Function;
        let func_ref = valref(func);
        let layout = pass.compute_optimal_layout(&func_ref);
        assert!(layout.is_empty());
    }

    #[test]
    fn test_compute_optimal_layout_single() {
        let pass = MachineBlockPlacement::new();
        let bb = make_block_named("entry");

        let mut func = Value::new(llvm_native_core::types::Type::void());
        func.subclass = SubclassKind::Function;
        func.operands = vec![bb];
        let func_ref = valref(func);

        let layout = pass.compute_optimal_layout(&func_ref);
        assert_eq!(layout.len(), 1);
        assert_eq!(layout[0], 0);
    }

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

    #[test]
    fn test_merge_blocks_empty() {
        let mut pass = MachineBlockPlacement::new();
        let mut func = Value::new(llvm_native_core::types::Type::void());
        func.subclass = SubclassKind::Function;
        let func_ref = valref(func);
        pass.merge_blocks(&func_ref);
    }

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

    #[test]
    fn test_run_on_function_no_blocks() {
        let mut pass = MachineBlockPlacement::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);
    }

    #[test]
    fn test_compute_branch_probabilities_two_successors() {
        let pass = MachineBlockPlacement::new();
        let bb1 = make_block_named("entry");
        let bb2 = make_block_named("then");
        let bb3 = make_block_named("else");
        set_successors(&bb1, vec![bb2, bb3]);

        let mut func = Value::new(llvm_native_core::types::Type::void());
        func.subclass = SubclassKind::Function;
        func.operands = vec![bb1.clone()];
        let func_ref = valref(func);

        let probs = pass.compute_branch_probabilities(&func_ref);
        let bb1_vid = bb1.borrow().vid;
        let edges: Vec<_> = probs.keys().filter(|k| k.0 == bb1_vid).collect();
        assert_eq!(edges.len(), 2);
    }

    #[test]
    fn test_block_priority_ordering() {
        let a = BlockPriority { freq: 0.9, idx: 0 };
        let b = BlockPriority { freq: 0.1, idx: 1 };
        assert!(a > b);
    }
}