analyssa 0.4.1

//! Worklist-based data flow solver.
//!
//! This module provides the iterative solver that computes fixpoints for
//! data flow analyses. It uses a worklist algorithm with traversal order
//! optimized for the analysis direction.
//!
//! # Algorithm
//!
//! The solver iterates until a fixpoint is reached:
//!
//! **Initialization**:
//! 1. All blocks start with the analysis's `initial()` value
//! 2. Set `boundary()` value at entry (forward) or exits (backward)
//! 3. Add all blocks to the worklist in traversal order (reverse postorder
//!    for forward, postorder for backward)
//!
//! **Iteration**:
//! 1. Pop a block from the worklist (deduplicated via `in_worklist` flag)
//! 2. Compute the input/output by meeting lattice values from all
//!    predecessors (forward) or successors (backward)
//! 3. Preserve boundary values for entry/exit blocks
//! 4. Apply the transfer function to compute the output/input
//! 5. If the result changed, add affected adjacent blocks
//!    (successors for forward, predecessors for backward) to the worklist
//!
//! **Finalization**:
//! 1. Call the analysis's `finalize()` hook for any post-processing
//! 2. Return `AnalysisResults` with per-block in/out states
//!
//! # Complexity
//!
//! On reducible CFGs, converges in O(d) iterations where d is the loop
//! nesting depth (typically small). Total work: O(n * d * h) where n is
//! the block count and h is the lattice height (maximum chain length from
//! top to bottom). Each iteration processes each block at most once via
//! the worklist deduplication flag.

use std::{cmp::Reverse, collections::BinaryHeap, marker::PhantomData};

use crate::{
    analysis::dataflow::{
        framework::{AnalysisResults, DataFlowAnalysis, DataFlowCfg, Direction},
        lattice::MeetSemiLattice,
    },
    bitset::BitSet,
    graph::NodeId,
    ir::function::SsaFunction,
    target::Target,
};

/// Worklist-based data flow solver.
///
/// This solver computes fixpoints for data flow analyses using an iterative
/// worklist algorithm. It supports both forward and backward analyses.
///
/// # Usage
///
/// ```rust
/// use analyssa::{
///     analysis::{
///         dataflow::{DataFlowSolver, ReachingDefinitions},
///         SsaCfg,
///     },
///     ir::SsaVarId,
///     testing,
/// };
///
/// let ssa = testing::diamond_phi_fixture();
/// let graph = SsaCfg::from_ssa(&ssa);
///
/// let analysis = ReachingDefinitions::new(&ssa);
/// let solver = DataFlowSolver::new(analysis);
/// let results = solver.solve(&ssa, &graph);
///
/// // Access results: the branch condition defined in block 0 reaches block 1.
/// let condition = SsaVarId::from_index(0);
/// let in_state = results.in_state(1).unwrap();
/// assert!(in_state.definitions().any(|var| var == condition));
/// ```
pub struct DataFlowSolver<T: Target, A: DataFlowAnalysis<T>> {
    /// The analysis being solved.
    analysis: A,
    /// Input state for each block.
    in_states: Vec<A::Lattice>,
    /// Output state for each block.
    out_states: Vec<A::Lattice>,
    /// Worklist of blocks to process, prioritized by traversal order so blocks
    /// are visited in (reverse) postorder. `Reverse((priority, block))` turns
    /// the max-heap into a min-heap on the priority, which reduces fixpoint
    /// iterations on loops compared with a plain FIFO.
    worklist: BinaryHeap<Reverse<(usize, usize)>>,
    /// Whether each block is currently in the worklist (for deduplication).
    in_worklist: Vec<bool>,
    /// Times each block's transfer function has run, handed to
    /// [`DataFlowAnalysis::widen`] so an unbounded-height lattice can stay
    /// exact for the first few passes and widen after.
    visits: Vec<usize>,
    /// Per-block traversal priority (position in RPO for forward analyses,
    /// postorder for backward); smaller is processed first.
    order_priority: Vec<usize>,
    /// Exit blocks, precomputed once per solve for O(1) membership in the
    /// backward solver hot loop (avoids re-allocating `cfg.exits()` per visit).
    exit_blocks: BitSet,
    /// Number of iterations performed.
    iterations: usize,
    _phantom: PhantomData<T>,
}

impl<T: Target, A: DataFlowAnalysis<T>> DataFlowSolver<T, A> {
    /// Creates a new solver for the given analysis.
    #[must_use]
    pub fn new(analysis: A) -> Self {
        Self {
            analysis,
            in_states: Vec::new(),
            out_states: Vec::new(),
            worklist: BinaryHeap::new(),
            in_worklist: Vec::new(),
            visits: Vec::new(),
            order_priority: Vec::new(),
            exit_blocks: BitSet::new(0),
            iterations: 0,
            _phantom: PhantomData,
        }
    }

    /// Solves the data flow analysis to a fixpoint.
    ///
    /// Returns the analysis results containing input and output states
    /// for each basic block.
    pub fn solve<C: DataFlowCfg>(
        mut self,
        ssa: &SsaFunction<T>,
        cfg: &C,
    ) -> AnalysisResults<A::Lattice>
    where
        A::Lattice: Clone,
    {
        let num_blocks = ssa.block_count();
        if num_blocks == 0 {
            return AnalysisResults::new(Vec::new(), Vec::new());
        }

        // Initialize states
        self.initialize(ssa, cfg);

        // Main iteration loop
        self.iterate(ssa, cfg);

        // Finalize
        self.analysis
            .finalize(&self.in_states, &self.out_states, ssa);

        AnalysisResults::new(self.in_states, self.out_states)
    }

    /// Returns the number of iterations performed.
    #[must_use]
    pub const fn iterations(&self) -> usize {
        self.iterations
    }

    /// Initializes the solver state.
    fn initialize<C: DataFlowCfg>(&mut self, ssa: &SsaFunction<T>, cfg: &C)
    where
        A::Lattice: Clone,
    {
        let num_blocks = ssa.block_count();
        let initial = self.analysis.initial(ssa);
        let boundary = self.analysis.boundary(ssa);

        // Initialize all blocks with the initial value
        self.in_states = vec![initial.clone(); num_blocks];
        self.out_states = vec![initial; num_blocks];
        self.in_worklist = vec![false; num_blocks];
        self.visits = vec![0; num_blocks];

        // Set boundary conditions based on direction
        match A::DIRECTION {
            Direction::Forward => {
                // Entry block gets boundary value
                let entry = cfg.entry().index();
                if let Some(slot) = self.in_states.get_mut(entry) {
                    *slot = boundary;
                }
            }
            Direction::Backward => {
                // Exit blocks get boundary value. Cache them in a BitSet so the
                // backward fixpoint loop can test membership without rebuilding
                // (and reallocating) `cfg.exits()` on every block visit.
                let mut exit_blocks = BitSet::new(num_blocks);
                for exit in cfg.exits() {
                    let idx = exit.index();
                    exit_blocks.insert(idx);
                    if let Some(slot) = self.out_states.get_mut(idx) {
                        *slot = boundary.clone();
                    }
                }
                self.exit_blocks = exit_blocks;
            }
        }

        // Assign each block a traversal priority (position in RPO for forward,
        // postorder for backward) and seed the priority worklist with it.
        let order = match A::DIRECTION {
            Direction::Forward => cfg.reverse_postorder(),
            Direction::Backward => cfg.postorder(),
        };

        self.order_priority = vec![usize::MAX; num_blocks];
        for (pos, node) in order.iter().enumerate() {
            if let Some(slot) = self.order_priority.get_mut(node.index()) {
                *slot = pos;
            }
        }

        for node in &order {
            let idx = node.index();
            if let Some(slot) = self.in_worklist.get_mut(idx) {
                let prio = self.order_priority.get(idx).copied().unwrap_or(usize::MAX);
                self.worklist.push(Reverse((prio, idx)));
                *slot = true;
            }
        }
    }

    /// Main iteration loop.
    fn iterate<C: DataFlowCfg>(&mut self, ssa: &SsaFunction<T>, cfg: &C)
    where
        A::Lattice: Clone,
    {
        while let Some(Reverse((_, block_idx))) = self.worklist.pop() {
            if let Some(slot) = self.in_worklist.get_mut(block_idx) {
                *slot = false;
            }
            self.iterations = self.iterations.saturating_add(1);

            let changed = match A::DIRECTION {
                Direction::Forward => self.process_forward(block_idx, ssa, cfg),
                Direction::Backward => self.process_backward(block_idx, ssa, cfg),
            };

            if changed {
                // Add affected blocks to worklist
                self.add_affected_to_worklist(block_idx, cfg);
            }
        }
    }

    /// Processes a block in forward direction.
    ///
    /// Returns `true` if the output state changed.
    fn process_forward<C: DataFlowCfg>(
        &mut self,
        block_idx: usize,
        ssa: &SsaFunction<T>,
        cfg: &C,
    ) -> bool
    where
        A::Lattice: Clone,
    {
        // Compute input by meeting all predecessor outputs
        let node = NodeId::new(block_idx);
        let Some(current_in) = self.in_states.get(block_idx).cloned() else {
            return false;
        };
        let mut input = if cfg.predecessors(node).next().is_none() {
            // Entry block or unreachable - keep current in_state
            current_in.clone()
        } else {
            // Meet all predecessor outputs. The accumulator is seeded from the
            // first predecessor and the rest are folded in place, so the whole
            // fold costs one clone rather than one per predecessor.
            let mut result: Option<A::Lattice> = None;
            for pred in cfg.predecessors(node) {
                let Some(pred_out) = self.out_states.get(pred.index()) else {
                    continue;
                };
                match result.as_mut() {
                    None => result = Some(pred_out.clone()),
                    Some(acc) => acc.meet_into(pred_out),
                }
            }
            result.unwrap_or_else(|| current_in.clone())
        };

        // Special case: entry block keeps its boundary value
        if node == cfg.entry() {
            input = current_in.clone();
        }

        if let Some(slot) = self.in_states.get_mut(block_idx) {
            *slot = input.clone();
        }

        // Apply transfer function
        let Some(block) = ssa.block(block_idx) else {
            return false;
        };
        let output = self.analysis.transfer(block_idx, block, &input, ssa);

        // Check if output changed
        let visit = self.bump_visit(block_idx);
        let Some(out_slot) = self.out_states.get_mut(block_idx) else {
            return false;
        };
        let output = self.analysis.widen(block_idx, out_slot, output, visit);
        let changed = output != *out_slot;
        *out_slot = output;

        changed
    }

    /// Records another transfer of `block_idx` and returns the running count.
    fn bump_visit(&mut self, block_idx: usize) -> usize {
        let Some(slot) = self.visits.get_mut(block_idx) else {
            return 0;
        };
        *slot = slot.saturating_add(1);
        *slot
    }

    /// Processes a block in backward direction.
    ///
    /// Returns `true` if the input state changed.
    fn process_backward<C: DataFlowCfg>(
        &mut self,
        block_idx: usize,
        ssa: &SsaFunction<T>,
        cfg: &C,
    ) -> bool
    where
        A::Lattice: Clone,
    {
        // Compute output by meeting all successor inputs
        let node = NodeId::new(block_idx);
        let Some(current_out) = self.out_states.get(block_idx).cloned() else {
            return false;
        };
        let mut output = if cfg.successors(node).next().is_none() {
            // Exit block or dead end - keep current out_state
            current_out.clone()
        } else {
            // Meet all successor inputs
            let mut result: Option<A::Lattice> = None;
            for succ in cfg.successors(node) {
                let Some(succ_in) = self.in_states.get(succ.index()) else {
                    continue;
                };
                result = Some(match result {
                    None => succ_in.clone(),
                    Some(acc) => acc.meet(succ_in),
                });
            }
            result.unwrap_or_else(|| current_out.clone())
        };

        // Special case: exit blocks keep their boundary value
        if self.exit_blocks.contains(node.index()) {
            output = current_out.clone();
        }

        if let Some(slot) = self.out_states.get_mut(block_idx) {
            *slot = output.clone();
        }

        // Apply transfer function (backward: input = transfer(output))
        let Some(block) = ssa.block(block_idx) else {
            return false;
        };
        let input = self.analysis.transfer(block_idx, block, &output, ssa);

        // Check if input changed
        let visit = self.bump_visit(block_idx);
        let Some(in_slot) = self.in_states.get_mut(block_idx) else {
            return false;
        };
        let input = self.analysis.widen(block_idx, in_slot, input, visit);
        let changed = input != *in_slot;
        *in_slot = input;

        changed
    }

    /// Adds affected blocks to the worklist after a change.
    fn add_affected_to_worklist<C: DataFlowCfg>(&mut self, block_idx: usize, cfg: &C) {
        let node = NodeId::new(block_idx);

        let priority = &self.order_priority;
        let enqueue =
            |idx: usize, list: &mut Vec<bool>, work: &mut BinaryHeap<Reverse<(usize, usize)>>| {
                if let Some(slot) = list.get_mut(idx)
                    && !*slot
                {
                    let prio = priority.get(idx).copied().unwrap_or(usize::MAX);
                    work.push(Reverse((prio, idx)));
                    *slot = true;
                }
            };

        match A::DIRECTION {
            Direction::Forward => {
                // Forward: successors are affected
                for succ in cfg.successors(node) {
                    enqueue(succ.index(), &mut self.in_worklist, &mut self.worklist);
                }
            }
            Direction::Backward => {
                // Backward: predecessors are affected
                for pred in cfg.predecessors(node) {
                    enqueue(pred.index(), &mut self.in_worklist, &mut self.worklist);
                }
            }
        }
    }
}