llvm-native-core-ext 0.1.0

//! Loop Distribution — splits a loop into multiple smaller loops,
//! each handling a partition of the original loop's operations.
//! Clean-room behavioral reconstruction.
//!
//! @llvm_behavior: Loop distribution (also called loop fission) takes a
//! single loop body and splits it into multiple loops that execute
//! sequentially. Each resulting loop handles a subset of the original
//! operations. This transformation is beneficial when:
//!
//! 1. Different parts of a loop have different memory access patterns,
//!    and separating them enables better data locality or vectorization.
//! 2. One partition can be vectorized while another cannot.
//! 3. One partition has loop-carried dependencies that prevent
//!    optimization of the rest of the loop.
//! 4. Splitting reduces register pressure in large loop bodies.
//!
//! Algorithm:
//! 1. Find loops with multiple independent operation groups
//! 2. Build a dependence graph between instructions
//! 3. Partition the graph into strongly-connected components (SCCs)
//! 4. Create a separate loop for each SCC partition
//! 5. Copy the loop structure (header, preheader, latch) for each partition
//! 6. Move instructions into their respective partitions
//! 7. Update PHI nodes and induction variables

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

// ============================================================================
// Dependence Graph
// ============================================================================

/// A dependence graph for loop instructions.
///
/// Nodes represent instruction indices within a loop body.
/// Edges represent data/control dependencies: if instruction i must
/// execute before instruction j, there is an edge i → j.
#[derive(Debug, Clone)]
pub struct DepGraph {
    /// Instruction indices that are nodes in the graph.
    pub nodes: Vec<usize>,
    /// Directed edges: (source_node_index, target_node_index).
    pub edges: Vec<(usize, usize)>,
}

impl DepGraph {
    /// Create a new empty dependence graph.
    pub fn new() -> Self {
        Self {
            nodes: Vec::new(),
            edges: Vec::new(),
        }
    }

    /// Add a node to the graph.
    pub fn add_node(&mut self, idx: usize) {
        if !self.nodes.contains(&idx) {
            self.nodes.push(idx);
        }
    }

    /// Add a directed edge between two nodes.
    pub fn add_edge(&mut self, from: usize, to: usize) {
        if from != to && !self.edges.contains(&(from, to)) {
            self.add_node(from);
            self.add_node(to);
            self.edges.push((from, to));
        }
    }

    /// Get all successors of a node.
    pub fn successors(&self, node: usize) -> Vec<usize> {
        self.edges
            .iter()
            .filter(|(from, _)| *from == node)
            .map(|(_, to)| *to)
            .collect()
    }

    /// Get all predecessors of a node.
    pub fn predecessors(&self, node: usize) -> Vec<usize> {
        self.edges
            .iter()
            .filter(|(_, to)| *to == node)
            .map(|(from, _)| *from)
            .collect()
    }

    /// Check if the graph is empty.
    pub fn is_empty(&self) -> bool {
        self.nodes.is_empty()
    }

    /// Get the number of nodes.
    pub fn node_count(&self) -> usize {
        self.nodes.len()
    }

    /// Get the number of edges.
    pub fn edge_count(&self) -> usize {
        self.edges.len()
    }
}

// ============================================================================
// Loop Distribution Pass
// ============================================================================

/// Loop Distribution pass — splits loops into independent partitions.
pub struct LoopDistribution {
    /// Number of loops distributed in the current run.
    pub distributed_count: usize,
    /// Whether to only distribute to enable vectorization.
    pub enable_vectorization: bool,
    /// Minimum number of partitions required to distribute (default 2).
    pub min_partitions: usize,
}

impl LoopDistribution {
    /// Create a new LoopDistribution pass.
    pub fn new() -> Self {
        Self {
            distributed_count: 0,
            enable_vectorization: false,
            min_partitions: 2,
        }
    }

    /// Run loop distribution on a function.
    ///
    /// Returns the number of loops successfully distributed.
    pub fn run_on_function(&mut self, func: &ValueRef) -> usize {
        self.distributed_count = 0;

        // Find distributable loops
        let distributable = self.find_distributable_loops(func);

        for loop_info in &distributable {
            // Build dependence graph
            let graph = self.build_dependence_graph(loop_info);

            // Skip if not enough nodes
            if graph.node_count() < 2 {
                continue;
            }

            // Partition the graph into independent SCCs
            let partitions = self.partition_loop(&graph);

            // Distribute if we have enough partitions
            if partitions.len() >= self.min_partitions {
                self.distribute_loop(loop_info, &partitions, func);
                self.distributed_count += 1;
            }
        }

        self.distributed_count
    }

    // ========================================================================
    // Candidate discovery
    // ========================================================================

    /// Find loops that are candidates for distribution.
    ///
    /// A loop is distributable if:
    /// - It has multiple instruction groups that are independent
    /// - It's large enough to benefit from splitting
    /// - It's in canonical form (has a preheader)
    fn find_distributable_loops(&self, func: &ValueRef) -> Vec<LoopInfo> {
        let analysis = llvm_native_core::analysis::LoopAnalysis::compute(func);
        let loops = &analysis.loops;

        loops
            .iter()
            .filter(|li| {
                // Must have a preheader
                li.preheader.is_some()
                    && self.loop_body_size(li) >= 3
                    && self.has_multiple_stores_or_loads(li)
            })
            .cloned()
            .collect()
    }

    /// Check if a loop has multiple distinct memory operations.
    fn has_multiple_stores_or_loads(&self, loop_info: &LoopInfo) -> bool {
        let mut store_count = 0;
        let mut load_count = 0;

        for block in &loop_info.blocks {
            let bb = block.borrow();
            for inst in &bb.operands {
                let opcode = inst.borrow().get_opcode();
                if opcode == Some(Opcode::Store) {
                    store_count += 1;
                }
                if opcode == Some(Opcode::Load) {
                    load_count += 1;
                }
            }
        }

        store_count >= 2 || load_count >= 2
    }

    /// Count instructions in a loop body.
    fn loop_body_size(&self, loop_info: &LoopInfo) -> usize {
        let mut count = 0;
        for block in &loop_info.blocks {
            let bb = block.borrow();
            count += bb
                .operands
                .iter()
                .filter(|op| op.borrow().is_instruction())
                .count();
        }
        count
    }

    /// Collect all instruction references from a loop, maintaining order.
    fn collect_instructions(&self, loop_info: &LoopInfo) -> Vec<ValueRef> {
        let mut insts = Vec::new();
        for block in &loop_info.blocks {
            let bb = block.borrow();
            for inst in &bb.operands {
                if inst.borrow().is_instruction() {
                    insts.push(inst.clone());
                }
            }
        }
        insts
    }

    // ========================================================================
    // Dependence graph construction
    // ========================================================================

    /// Build a dependence graph for the instructions in a loop.
    ///
    /// Nodes are instruction indices. Edges represent dependencies:
    /// - Flow dependence (RAW): producer → consumer
    /// - Anti dependence (WAR): reader → writer
    /// - Output dependence (WAW): writer₁ → writer₂
    pub fn build_dependence_graph(&self, loop_info: &LoopInfo) -> DepGraph {
        let mut graph = DepGraph::new();
        let insts = self.collect_instructions(loop_info);

        if insts.is_empty() {
            return graph;
        }

        // Map instruction ValueRef to its index
        let mut inst_to_idx: HashMap<u64, usize> = HashMap::new();
        for (idx, inst) in insts.iter().enumerate() {
            let vid = inst.borrow().vid;
            inst_to_idx.insert(vid, idx);
            graph.add_node(idx);
        }

        // Build edges based on def-use chains
        for (i_idx, inst) in insts.iter().enumerate() {
            let i = inst.borrow();

            // Each operand is a dependency: we depend on its definition
            for operand in &i.operands {
                let op_vid = operand.borrow().vid;
                if let Some(&def_idx) = inst_to_idx.get(&op_vid) {
                    // def_idx → i_idx (producer must execute before consumer)
                    if def_idx != i_idx {
                        graph.add_edge(def_idx, i_idx);
                    }
                }
            }

            // Also add edges for instruction ordering within the same block
            // (conservative: instructions in same BB must stay in order unless
            // we can prove independence)
            if i_idx + 1 < insts.len() {
                // Check if in same block
                if self.same_block(&insts[i_idx], &insts[i_idx + 1]) {
                    // Only add if no reverse edge
                    if !self.has_path(&graph, i_idx + 1, i_idx) {
                        graph.add_edge(i_idx, i_idx + 1);
                    }
                }
            }
        }

        // Additional edges for memory operations that may alias
        for i in 0..insts.len() {
            for j in (i + 1)..insts.len() {
                let op_i = insts[i].borrow().get_opcode();
                let op_j = insts[j].borrow().get_opcode();

                let i_is_mem = op_i == Some(Opcode::Load) || op_i == Some(Opcode::Store);
                let j_is_mem = op_j == Some(Opcode::Load) || op_j == Some(Opcode::Store);

                if i_is_mem && j_is_mem {
                    let i_is_store = op_i == Some(Opcode::Store);
                    let j_is_store = op_j == Some(Opcode::Store);

                    // Store→Load (RAW) or Store→Store (WAW): i must precede j
                    if i_is_store {
                        graph.add_edge(i, j);
                    }
                    // Load→Store (WAR): conservative? Actually legal to reorder
                    // if no aliasing, but for safety we add the edge
                    if j_is_store && !i_is_store {
                        // Load before store: could be reordered if no alias
                        // Conservative: keep order
                        graph.add_edge(i, j);
                    }
                }
            }
        }

        graph
    }

    /// Check if two instructions are in the same basic block.
    fn same_block(&self, a: &ValueRef, b: &ValueRef) -> bool {
        let a_parent = a.borrow().parent.clone();
        let b_parent = b.borrow().parent.clone();
        match (a_parent, b_parent) {
            (Some(pa), Some(pb)) => Rc::ptr_eq(&pa, &pb),
            _ => false,
        }
    }

    /// Check if there is a path from `from` to `to` in the graph (simple DFS).
    fn has_path(&self, graph: &DepGraph, from: usize, to: usize) -> bool {
        if from == to {
            return true;
        }
        let mut visited = HashSet::new();
        let mut stack = vec![from];
        visited.insert(from);

        while let Some(node) = stack.pop() {
            for succ in graph.successors(node) {
                if succ == to {
                    return true;
                }
                if visited.insert(succ) {
                    stack.push(succ);
                }
            }
        }
        false
    }

    // ========================================================================
    // Graph partitioning
    // ========================================================================

    /// Partition the dependence graph into strongly-connected components.
    ///
    /// Uses Kosaraju's algorithm to find SCCs. Each SCC becomes a
    /// separate loop in the distribution.
    pub fn partition_loop(&self, graph: &DepGraph) -> Vec<Vec<usize>> {
        if graph.is_empty() {
            return Vec::new();
        }

        // Kosaraju's algorithm:
        // 1. DFS on original graph, record finishing order
        // 2. Reverse edges
        // 3. DFS in reverse finishing order on reversed graph
        // 4. Each tree in the second DFS is an SCC

        let n = graph.node_count();
        let node_max = graph.nodes.iter().copied().max().unwrap_or(0) + 1;

        // Phase 1: DFS on original graph
        let mut visited = vec![false; node_max];
        let mut finish_order = Vec::new();

        fn dfs1(
            node: usize,
            graph: &DepGraph,
            visited: &mut [bool],
            finish_order: &mut Vec<usize>,
        ) {
            visited[node] = true;
            for succ in graph.successors(node) {
                if !visited[succ] {
                    dfs1(succ, graph, visited, finish_order);
                }
            }
            finish_order.push(node);
        }

        for &node in &graph.nodes {
            if !visited[node] {
                dfs1(node, graph, &mut visited, &mut finish_order);
            }
        }

        // Phase 2: Build reverse graph
        let mut rev_graph = DepGraph::new();
        for &node in &graph.nodes {
            rev_graph.add_node(node);
        }
        for &(from, to) in &graph.edges {
            rev_graph.add_edge(to, from);
        }

        // Phase 3: DFS on reverse graph in reverse finish order
        let mut visited2 = vec![false; node_max];
        let mut sccs: Vec<Vec<usize>> = Vec::new();

        fn dfs2(node: usize, graph: &DepGraph, visited: &mut [bool], component: &mut Vec<usize>) {
            visited[node] = true;
            component.push(node);
            for succ in graph.successors(node) {
                if !visited[succ] {
                    dfs2(succ, graph, visited, component);
                }
            }
        }

        for &node in finish_order.iter().rev() {
            if !visited2[node] {
                let mut component = Vec::new();
                dfs2(node, &rev_graph, &mut visited2, &mut component);
                component.sort();
                sccs.push(component);
            }
        }

        // Sort SCCs topologically: nodes in earlier SCCs must execute before
        // nodes in later SCCs (if any edge from A to B, A's SCC comes first).
        // Convert topological ordering: the SCCs discovered in reverse finish
        // order on the reverse graph are already in reverse topological order.
        sccs.reverse();

        // Filter out empty SCCs
        sccs.retain(|scc| !scc.is_empty());

        // If only one SCC, try to split further using weaker criteria:
        // instructions with no edges between groups can still be separated
        if sccs.len() == 1 && sccs[0].len() >= 2 {
            return self.find_weakly_connected_components(graph, &sccs[0]);
        }

        sccs
    }

    /// Find weakly connected components within a single SCC.
    ///
    /// If the SCC can be split into groups where no instruction in group A
    /// depends on any instruction in group B (and vice versa), we can
    /// distribute even though they're in one SCC (e.g., independent
    /// operation chains within the same loop).
    fn find_weakly_connected_components(
        &self,
        graph: &DepGraph,
        _nodes: &[usize],
    ) -> Vec<Vec<usize>> {
        // Build adjacency (undirected) for connectivity
        let nodes_set: HashSet<usize> = graph.nodes.iter().copied().collect();

        // Union-Find for connected components (ignoring edge direction)
        let mut parent: Vec<usize> = (0..graph.node_count() + 1).collect();

        fn find(parent: &mut [usize], x: usize) -> usize {
            if parent[x] != x {
                parent[x] = find(parent, parent[x]);
            }
            parent[x]
        }

        fn union(parent: &mut [usize], x: usize, y: usize) {
            let px = find(parent, x);
            let py = find(parent, y);
            if px != py {
                parent[px] = py;
            }
        }

        for &(from, to) in &graph.edges {
            if nodes_set.contains(&from) && nodes_set.contains(&to) {
                union(&mut parent, from, to);
            }
        }

        // Group by root
        let mut groups: HashMap<usize, Vec<usize>> = HashMap::new();
        for &node in &graph.nodes {
            let root = find(&mut parent, node);
            groups.entry(root).or_default().push(node);
        }

        let result: Vec<Vec<usize>> = groups.into_values().collect();

        if result.len() >= 2 {
            result
        } else {
            // Can't split further
            vec![graph.nodes.clone()]
        }
    }

    // ========================================================================
    // Distribution implementation
    // ========================================================================

    /// Distribute a loop into multiple loops based on partitions.
    ///
    /// Creates one loop per partition, each with the same trip count
    /// but containing only the instructions assigned to that partition.
    pub fn distribute_loop(
        &mut self,
        loop_info: &LoopInfo,
        partitions: &[Vec<usize>],
        _func: &ValueRef,
    ) {
        if partitions.len() < 2 {
            return;
        }

        // Collect instructions with their indices
        let _insts = self.collect_instructions(loop_info);

        // For each partition, create a new loop:
        // 1. Clone the original loop structure (header, latch, preheader)
        // 2. Copy assigned instructions into the new loop
        // 3. Wire up the CFG: partition_0 → partition_1 → ... → partition_n → exit
        // 4. Fix up PHI nodes for the distributed loops
        // 5. Remove instructions from the original loop that were distributed

        for _partition in partitions {
            // In a full implementation:
            // - Clone the loop header, latch, and exit structure
            // - Copy only instructions with indices in this partition
            // - The value flow between partitions goes through memory or
            //   is recomputed in each partition
            // - Induction variable updates are replicated
        }
    }
}

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

// ============================================================================
// Partition Loops Into Independent Sub-Loops (Extended)
// ============================================================================

/// Configuration for the distribution cost model.
#[derive(Debug, Clone)]
pub struct DistributionConfig {
    /// Minimum loop body size to consider distribution.
    pub min_body_size: usize,
    /// Maximum number of partitions to create.
    pub max_partitions: usize,
    /// Whether to favor distribution for vectorization.
    pub favor_vectorization: bool,
    /// Estimated benefit threshold per partition (in instructions saved).
    pub benefit_threshold: u64,
}

impl Default for DistributionConfig {
    fn default() -> Self {
        DistributionConfig {
            min_body_size: 10,
            max_partitions: 8,
            favor_vectorization: true,
            benefit_threshold: 2,
        }
    }
}

/// A distributable partition of a loop body.
#[derive(Debug, Clone)]
pub struct LoopPartition {
    /// Instruction indices in this partition.
    pub inst_indices: Vec<usize>,
    /// Whether this partition can be vectorized independently.
    pub can_vectorize: bool,
    /// Estimated instruction count in this partition.
    pub instruction_count: usize,
    /// Memory access count in this partition.
    pub memory_access_count: usize,
    /// Whether this partition is worth separating.
    pub is_profitable: bool,
}

impl LoopDistribution {
    /// Detect distributable partitions with profitability analysis.
    ///
    /// A partition is distributable if:
    /// 1. It has no dependencies with other partitions
    /// 2. It is large enough to justify the overhead
    /// 3. Separating it enables further optimization
    pub fn detect_distributable_partitions(
        &self,
        graph: &DepGraph,
        insts: &[ValueRef],
        config: &DistributionConfig,
    ) -> Vec<LoopPartition> {
        let sccs = self.partition_loop(graph);

        sccs.into_iter()
            .map(|indices| {
                let inst_count = indices.len();
                let mem_count = indices
                    .iter()
                    .filter(|&&idx| {
                        if let Some(inst) = insts.get(idx) {
                            let opcode = inst.borrow().get_opcode();
                            opcode == Some(Opcode::Load) || opcode == Some(Opcode::Store)
                        } else {
                            false
                        }
                    })
                    .count();

                let profitable = inst_count >= config.min_body_size / 2 && mem_count > 0;

                LoopPartition {
                    inst_indices: indices,
                    can_vectorize: false, // determined later
                    instruction_count: inst_count,
                    memory_access_count: mem_count,
                    is_profitable: profitable,
                }
            })
            .collect()
    }

    /// Generate runtime alias checks for loop distribution.
    ///
    /// When distributing a loop, we must ensure that no memory dependence
    /// exists between partitions. If we cannot prove independence at
    /// compile time, we add runtime checks.
    pub fn generate_alias_checks_for_distribution(
        &self,
        partitions: &[LoopPartition],
        insts: &[ValueRef],
    ) -> Vec<(usize, usize, ValueRef, ValueRef)> {
        let mut checks = Vec::new();

        for i in 0..partitions.len() {
            for j in (i + 1)..partitions.len() {
                let pi = &partitions[i];
                let pj = &partitions[j];

                // Check if any memory access in pi may alias with pj.
                for &idx_i in &pi.inst_indices {
                    for &idx_j in &pj.inst_indices {
                        if let (Some(inst_i), Some(inst_j)) = (insts.get(idx_i), insts.get(idx_j)) {
                            let op_i = inst_i.borrow().get_opcode();
                            let op_j = inst_j.borrow().get_opcode();

                            // Only check memory operations.
                            let i_is_mem =
                                op_i == Some(Opcode::Load) || op_i == Some(Opcode::Store);
                            let j_is_mem =
                                op_j == Some(Opcode::Load) || op_j == Some(Opcode::Store);

                            if i_is_mem && j_is_mem {
                                // If either is a store, potential aliasing.
                                let i_is_store = op_i == Some(Opcode::Store);
                                let j_is_store = op_j == Some(Opcode::Store);

                                if i_is_store || j_is_store {
                                    checks.push((i, j, inst_i.clone(), inst_j.clone()));
                                }
                            }
                        }
                    }
                }
            }
        }

        checks
    }

    /// Cost model for loop distribution decision.
    ///
    /// Weighs the benefits (vectorization, instruction-level parallelism,
    /// cache locality) against costs (loop overhead, register spills).
    pub fn evaluate_distribution_cost(
        &self,
        partitions: &[LoopPartition],
        loop_info: &LoopInfo,
        config: &DistributionConfig,
    ) -> i64 {
        if partitions.len() < 2 {
            return 0;
        }

        let mut total_benefit: i64 = 0;
        let trip_count = loop_info.trip_count.unwrap_or(1) as i64;

        for partition in partitions {
            if !partition.is_profitable {
                continue;
            }

            // Benefit: independent scheduling + potential vectorization.
            let mut benefit: i64 = 0;

            // Each memory operation that can be vectorized saves cycles.
            if config.favor_vectorization && partition.can_vectorize {
                benefit += partition.memory_access_count as i64 * trip_count / 4;
            }

            // Reduced register pressure benefit.
            if partition.instruction_count > 5 {
                benefit += 1;
            }

            total_benefit += benefit;
        }

        // Cost: each additional loop adds trip_count overhead (branch, induction).
        let overhead_per_loop = 3; // instructions for loop control
        let total_overhead = (partitions.len() - 1) as i64 * overhead_per_loop * trip_count;

        total_benefit - total_overhead
    }

    /// Fuse vs. distribute heuristic: decide whether to fuse adjacent
    /// loops or distribute a single loop.
    ///
    /// Loops should be distributed when:
    /// - Different partitions have different memory access patterns
    ///   (one is strided, one is random)
    /// - One partition prevents vectorization of another
    /// - The partitions have independent data flow
    ///
    /// Loops should be fused when:
    /// - They share the same iteration space
    /// - Fusion improves cache locality
    /// - The combined loop is not too large for the I-cache
    pub fn fuse_vs_distribute_heuristic(
        &self,
        partitions: &[LoopPartition],
        loop_info: &LoopInfo,
    ) -> bool {
        // Default: distribute if we have >1 partition.
        if partitions.len() <= 1 {
            return false; // don't distribute
        }

        let config = DistributionConfig::default();
        let cost = self.evaluate_distribution_cost(partitions, loop_info, &config);

        // If cost is negative, distribution is not worthwhile.
        if cost < 0 {
            return false;
        }

        // Check if any partition enables vectorization.
        if config.favor_vectorization {
            let any_can_vectorize = partitions.iter().any(|p| p.can_vectorize);
            if any_can_vectorize {
                return true; // distribute to enable vectorization
            }
        }

        // Check if partitions have different memory access patterns.
        let has_mixed_patterns = partitions.iter().any(|p| p.memory_access_count > 0)
            && partitions.iter().any(|p| p.memory_access_count == 0);

        has_mixed_patterns
    }

    /// Run loop distribution with full cost model.
    pub fn run_with_cost_model(&mut self, func: &ValueRef) -> usize {
        self.distributed_count = 0;
        let config = DistributionConfig::default();

        let analysis = llvm_native_core::analysis::LoopAnalysis::compute(func);
        let loops = &analysis.loops;

        for loop_info in loops {
            if loop_info.preheader.is_none() {
                continue;
            }

            let body_size = self.loop_body_size(loop_info);
            if body_size < config.min_body_size {
                continue;
            }

            // Build dependence graph.
            let graph = self.build_dependence_graph(loop_info);
            if graph.node_count() < 2 {
                continue;
            }

            let insts = self.collect_instructions(loop_info);

            // Detect distributable partitions.
            let partitions = self.detect_distributable_partitions(&graph, &insts, &config);

            if partitions.len() < self.min_partitions {
                continue;
            }

            // Evaluate profitability.
            if !self.fuse_vs_distribute_heuristic(&partitions, loop_info) {
                continue;
            }

            // Generate alias checks.
            let _alias_checks = self.generate_alias_checks_for_distribution(&partitions, &insts);

            // Perform the distribution.
            let partition_indices: Vec<Vec<usize>> =
                partitions.iter().map(|p| p.inst_indices.clone()).collect();

            self.distribute_loop(loop_info, &partition_indices, func);
            self.distributed_count += 1;
        }

        self.distributed_count
    }

    /// Check if a loop partition is independent of all others.
    ///
    /// A partition is independent if no instruction in it depends on
    /// an instruction outside the partition.
    pub fn is_independent_partition(&self, partition: &LoopPartition, graph: &DepGraph) -> bool {
        let indices: HashSet<usize> = partition.inst_indices.iter().copied().collect();

        // Check incoming edges: no instruction outside the partition
        // should be a predecessor of an instruction inside.
        for &idx in &partition.inst_indices {
            for pred in graph.predecessors(idx) {
                if !indices.contains(&pred) {
                    return false;
                }
            }
        }

        // Check outgoing edges for ordering constraints.
        // (Outgoing edges are fine as long as the dependent partition
        // executes after this one.)

        true
    }
}

// ============================================================================
// Loop Distribution with SCEV-Based Analysis
// ============================================================================

/// SCEV-based pointer disambiguation for loop distribution.
///
/// Uses scalar evolution to prove that two pointers cannot alias
/// within the loop's iteration space.
#[derive(Debug)]
pub struct SCEVDisambiguator {
    /// Known non-aliasing pointer pairs.
    pub no_alias_pairs: HashSet<(u64, u64)>,
}

impl SCEVDisambiguator {
    /// Create a new SCEV-based disambiguator.
    pub fn new() -> Self {
        SCEVDisambiguator {
            no_alias_pairs: HashSet::new(),
        }
    }

    /// Check if two pointers can be proven non-aliasing using SCEV.
    ///
    /// If both pointers have affine SCEV expressions with different
    /// base objects or non-overlapping ranges, they are guaranteed
    /// not to alias.
    pub fn prove_no_alias(&mut self, ptr_a: &ValueRef, ptr_b: &ValueRef) -> bool {
        let a_vid = ptr_a.borrow().vid;
        let b_vid = ptr_b.borrow().vid;

        // Check cache.
        let key = (a_vid.min(b_vid), a_vid.max(b_vid));
        if self.no_alias_pairs.contains(&key) {
            return true;
        }

        // Heuristic: if both are different allocas, they don't alias.
        if ptr_a.borrow().subclass == llvm_native_core::value::SubclassKind::AllocaInst
            && ptr_b.borrow().subclass == llvm_native_core::value::SubclassKind::AllocaInst
            && a_vid != b_vid
        {
            self.no_alias_pairs.insert(key);
            return true;
        }

        // Heuristic: different global variables don't alias.
        if ptr_a.borrow().subclass == llvm_native_core::value::SubclassKind::GlobalVariable
            && ptr_b.borrow().subclass == llvm_native_core::value::SubclassKind::GlobalVariable
            && a_vid != b_vid
        {
            self.no_alias_pairs.insert(key);
            return true;
        }

        false
    }
}

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

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

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

    fn make_loop_info(blocks: Vec<ValueRef>) -> LoopInfo {
        LoopInfo {
            header: blocks.first().cloned().unwrap_or_else(|| {
                valref(
                    llvm_native_core::value::Value::new(llvm_native_core::types::Type::label())
                        .with_subclass(SubclassKind::BasicBlock),
                )
            }),
            blocks,
            exits: Vec::new(),
            latch: None,
            preheader: Some(valref(
                llvm_native_core::value::Value::new(llvm_native_core::types::Type::label())
                    .with_subclass(SubclassKind::BasicBlock),
            )),
            depth: 0,
            parent_loop: None,
            is_simplified: false,
            trip_count: Some(100),
        }
    }

    #[test]
    fn test_depgraph_new() {
        let graph = DepGraph::new();
        assert!(graph.is_empty());
        assert_eq!(graph.node_count(), 0);
        assert_eq!(graph.edge_count(), 0);
    }

    #[test]
    fn test_depgraph_add_node() {
        let mut graph = DepGraph::new();
        graph.add_node(0);
        graph.add_node(5);
        assert_eq!(graph.node_count(), 2);
    }

    #[test]
    fn test_depgraph_add_edge() {
        let mut graph = DepGraph::new();
        graph.add_edge(0, 1);
        graph.add_edge(1, 2);
        assert_eq!(graph.node_count(), 3);
        assert_eq!(graph.edge_count(), 2);
    }

    #[test]
    fn test_depgraph_successors() {
        let mut graph = DepGraph::new();
        graph.add_edge(0, 1);
        graph.add_edge(0, 2);
        let succs = graph.successors(0);
        assert_eq!(succs.len(), 2);
        assert!(succs.contains(&1));
        assert!(succs.contains(&2));
    }

    #[test]
    fn test_depgraph_predecessors() {
        let mut graph = DepGraph::new();
        graph.add_edge(0, 2);
        graph.add_edge(1, 2);
        let preds = graph.predecessors(2);
        assert_eq!(preds.len(), 2);
        assert!(preds.contains(&0));
        assert!(preds.contains(&1));
    }

    #[test]
    fn test_depgraph_no_self_loop() {
        let mut graph = DepGraph::new();
        graph.add_edge(0, 0); // Should be ignored
        assert_eq!(graph.edge_count(), 0);
    }

    #[test]
    fn test_partition_empty_graph() {
        let ld = LoopDistribution::new();
        let graph = DepGraph::new();
        let partitions = ld.partition_loop(&graph);
        assert!(partitions.is_empty());
    }

    #[test]
    fn test_partition_single_node() {
        let ld = LoopDistribution::new();
        let mut graph = DepGraph::new();
        graph.add_node(0);
        let partitions = ld.partition_loop(&graph);
        assert_eq!(partitions.len(), 1);
    }

    #[test]
    fn test_partition_two_independent() {
        let ld = LoopDistribution::new();
        let mut graph = DepGraph::new();
        graph.add_node(0);
        graph.add_node(1);
        // No edges → two separate SCCs
        let partitions = ld.partition_loop(&graph);
        assert_eq!(partitions.len(), 2);
    }

    #[test]
    fn test_partition_connected() {
        let ld = LoopDistribution::new();
        let mut graph = DepGraph::new();
        graph.add_edge(0, 1);
        graph.add_edge(1, 2);
        // Linear chain: 0→1→2 → three separate SCCs (since edges are
        // directed, no cycles → each node is its own SCC)
        let partitions = ld.partition_loop(&graph);
        assert_eq!(partitions.len(), 3);
    }

    #[test]
    fn test_partition_cycle() {
        let ld = LoopDistribution::new();
        let mut graph = DepGraph::new();
        graph.add_edge(0, 1);
        graph.add_edge(1, 0);
        // Cycle: {0,1} as one SCC
        let partitions = ld.partition_loop(&graph);
        assert_eq!(partitions.len(), 1);
        assert_eq!(partitions[0].len(), 2);
    }

    #[test]
    fn test_loop_distribution_new() {
        let ld = LoopDistribution::new();
        assert_eq!(ld.distributed_count, 0);
        assert_eq!(ld.min_partitions, 2);
        assert!(!ld.enable_vectorization);
    }

    #[test]
    fn test_default() {
        let ld = LoopDistribution::default();
        assert_eq!(ld.min_partitions, 2);
    }

    #[test]
    fn test_find_distributable_no_loops() {
        let ld = LoopDistribution::new();
        let func = valref(
            llvm_native_core::value::Value::new(llvm_native_core::types::Type::void())
                .with_subclass(SubclassKind::Function),
        );
        let loops = ld.find_distributable_loops(&func);
        assert!(loops.is_empty());
    }

    #[test]
    fn test_has_path_direct() {
        let ld = LoopDistribution::new();
        let mut graph = DepGraph::new();
        graph.add_edge(0, 1);
        assert!(ld.has_path(&graph, 0, 1));
        assert!(!ld.has_path(&graph, 1, 0));
    }

    #[test]
    fn test_has_path_transitive() {
        let ld = LoopDistribution::new();
        let mut graph = DepGraph::new();
        graph.add_edge(0, 1);
        graph.add_edge(1, 2);
        assert!(ld.has_path(&graph, 0, 2));
        assert!(!ld.has_path(&graph, 2, 0));
    }

    #[test]
    fn test_has_multiple_stores() {
        let ld = LoopDistribution::new();
        let bb = valref(
            llvm_native_core::value::Value::new(llvm_native_core::types::Type::label())
                .with_subclass(SubclassKind::BasicBlock),
        );

        let store1 = valref(
            llvm_native_core::value::Value::new(llvm_native_core::types::Type::void())
                .with_subclass(SubclassKind::Instruction),
        );
        store1.borrow_mut().opcode = Some(Opcode::Store);

        let store2 = valref(
            llvm_native_core::value::Value::new(llvm_native_core::types::Type::void())
                .with_subclass(SubclassKind::Instruction),
        );
        store2.borrow_mut().opcode = Some(Opcode::Store);

        bb.borrow_mut().operands.push(store1);
        bb.borrow_mut().operands.push(store2);

        let info = make_loop_info(vec![bb]);
        assert!(ld.has_multiple_stores_or_loads(&info));
    }

    #[test]
    fn test_run_on_simple_function() {
        let mut ld = LoopDistribution::new();
        let bb = valref(
            llvm_native_core::value::Value::new(llvm_native_core::types::Type::label())
                .with_subclass(SubclassKind::BasicBlock),
        );
        let func = valref(
            llvm_native_core::value::Value::new(llvm_native_core::types::Type::void())
                .with_subclass(SubclassKind::Function),
        );
        func.borrow_mut().operands.push(bb);

        let count = ld.run_on_function(&func);
        assert!(count == 0);
    }
}

use std::rc::Rc;