weirflow 0.1.0

use super::reaching_defs_step;
use crate::dispatch_input_cache::OwnedDispatchInputs;
use crate::soundness::{DataflowEvidence, Soundness};
use vyre_primitives::graph::program_graph::ProgramGraphShape;

pub fn reaching_closure_evidence_via<F>(
    dispatch: &F,
    node_count: u32,
    edge_offsets: &[u32],
    edge_targets: &[u32],
    edge_kind_mask: &[u32],
    seed_bits: &[u32],
    max_iterations: u32,
) -> Result<DataflowEvidence<Vec<u32>>, String>
where
    F: Fn(&vyre::ir::Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    reaching_closure_via(
        dispatch,
        node_count,
        edge_offsets,
        edge_targets,
        edge_kind_mask,
        seed_bits,
        max_iterations,
    )
    .map(|value| DataflowEvidence::new(value, Soundness::Exact))
}

pub fn reaching_closure_via<F>(
    dispatch: &F,
    node_count: u32,
    edge_offsets: &[u32],
    edge_targets: &[u32],
    edge_kind_mask: &[u32],
    seed_bits: &[u32],
    max_iterations: u32,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    let owned_inputs = OwnedDispatchInputs::new();
    reaching_closure_borrowed_via(
        &|program, inputs, grid| {
            owned_inputs.dispatch_refreshing_slots(program, inputs, grid, &[5], dispatch)
        },
        node_count,
        edge_offsets,
        edge_targets,
        edge_kind_mask,
        seed_bits,
        max_iterations,
    )
}

/// GPU reaching-definition closure using borrowed dispatch inputs.
///
/// Invariant graph buffers are packed once and reused across fixpoint
/// iterations; only the changing frontier is repacked each round.
pub fn reaching_closure_borrowed_via<F>(
    dispatch: &F,
    node_count: u32,
    edge_offsets: &[u32],
    edge_targets: &[u32],
    edge_kind_mask: &[u32],
    seed_bits: &[u32],
    max_iterations: u32,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    crate::dispatch_decode::require_positive_iterations("reaching_closure", max_iterations)?;
    crate::dispatch_decode::require_bitset_tail_clear("reaching_closure", seed_bits, node_count)?;
    let out = reaching_closure_borrowed_into_via(
        &|program, inputs, grid, outputs| {
            let result = dispatch(program, inputs, grid)?;
            crate::output_scratch::replace_outputs_preserving_slots(outputs, result);
            Ok(())
        },
        node_count,
        edge_offsets,
        edge_targets,
        edge_kind_mask,
        seed_bits,
        max_iterations,
    )?;
    crate::dispatch_decode::require_bitset_tail_clear("reaching_closure output", &out, node_count)?;
    Ok(out)
}

/// GPU reaching-definition closure using borrowed dispatch inputs and
/// caller-owned output storage.
///
/// This is the zero-output-vector-churn hot path for backends that expose an
/// `_into` dispatch API. Invariant graph buffers are packed once, frontier
/// bytes are reused every iteration, and backend outputs are decoded from the
/// same caller-owned output vector.
pub fn reaching_closure_borrowed_into_via<F>(
    dispatch: &F,
    node_count: u32,
    edge_offsets: &[u32],
    edge_targets: &[u32],
    edge_kind_mask: &[u32],
    seed_bits: &[u32],
    max_iterations: u32,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let mut scratch = crate::fixed_point_scratch::FixedPointScratch::default();
    reaching_closure_borrowed_into_with_scratch_via(
        dispatch,
        node_count,
        edge_offsets,
        edge_targets,
        edge_kind_mask,
        seed_bits,
        max_iterations,
        &mut scratch,
    )
}

/// GPU reaching-definition closure using caller-owned fixed-point scratch.
pub fn reaching_closure_borrowed_into_with_scratch_via<F>(
    dispatch: &F,
    node_count: u32,
    edge_offsets: &[u32],
    edge_targets: &[u32],
    edge_kind_mask: &[u32],
    seed_bits: &[u32],
    max_iterations: u32,
    scratch: &mut crate::fixed_point_scratch::FixedPointScratch,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let mut result = crate::staging_reserve::reserved_vec(
        crate::dispatch_decode::bitset_word_capacity("reaching_closure result", node_count)?,
        "fixed-point result word",
    )?;
    reaching_closure_borrowed_into_result_with_scratch_via(
        dispatch,
        node_count,
        edge_offsets,
        edge_targets,
        edge_kind_mask,
        seed_bits,
        max_iterations,
        scratch,
        &mut result,
    )?;
    Ok(result)
}

/// GPU reaching-definition closure using caller-owned fixed-point scratch and
/// caller-owned result storage.
pub fn reaching_closure_borrowed_into_result_with_scratch_via<F>(
    dispatch: &F,
    node_count: u32,
    edge_offsets: &[u32],
    edge_targets: &[u32],
    edge_kind_mask: &[u32],
    seed_bits: &[u32],
    max_iterations: u32,
    scratch: &mut crate::fixed_point_scratch::FixedPointScratch,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    crate::dispatch_decode::require_positive_iterations("reaching_closure", max_iterations)?;
    let words = crate::dispatch_decode::bitset_word_capacity("reaching_closure", node_count)?;
    crate::dispatch_decode::require_bitset_words("reaching_closure", seed_bits, words)?;
    crate::dispatch_decode::require_bitset_tail_clear("reaching_closure", seed_bits, node_count)?;
    crate::dispatch_decode::require_csr_shape(
        "reaching_closure",
        node_count,
        edge_offsets,
        edge_targets,
        edge_kind_mask,
    )?;
    scratch.prepare_edge_kind_masks(
        edge_kind_mask,
        vyre_primitives::predicate::edge_kind::CONTROL,
    )?;
    let graph = crate::fixed_point_graph::FixedPointForwardGraph::new_for_kind(
        crate::fixed_point_graph::FixedPointAnalysisKind::Reaching,
        "reaching_closure",
        node_count,
        edge_offsets,
        edge_targets,
        &scratch.edge_kind_masks,
    )?;
    reaching_closure_prepared_borrowed_into_result_with_scratch_via(
        dispatch,
        &graph,
        seed_bits,
        max_iterations,
        scratch,
        result,
    )
}

pub fn reaching_closure_prepared_borrowed_into_with_scratch_via<F>(
    dispatch: &F,
    graph: &crate::fixed_point_graph::FixedPointForwardGraph,
    seed_bits: &[u32],
    max_iterations: u32,
    scratch: &mut crate::fixed_point_scratch::FixedPointScratch,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let mut result = crate::staging_reserve::reserved_vec(
        crate::dispatch_decode::bitset_word_capacity(
            "reaching_closure prepared result",
            graph.node_count,
        )?,
        "fixed-point result word",
    )?;
    reaching_closure_prepared_borrowed_into_result_with_scratch_via(
        dispatch,
        graph,
        seed_bits,
        max_iterations,
        scratch,
        &mut result,
    )?;
    Ok(result)
}

/// GPU reaching-definition closure over prepared invariant graph buffers into
/// caller-owned result storage.
pub fn reaching_closure_prepared_borrowed_into_result_with_scratch_via<F>(
    dispatch: &F,
    graph: &crate::fixed_point_graph::FixedPointForwardGraph,
    seed_bits: &[u32],
    max_iterations: u32,
    scratch: &mut crate::fixed_point_scratch::FixedPointScratch,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let program = reaching_defs_step(
        ProgramGraphShape::new(graph.node_count, graph.edge_count),
        "fin",
        "fout",
    );
    reaching_closure_prepared_program_borrowed_into_result_with_scratch_via(
        dispatch,
        graph,
        &program,
        seed_bits,
        max_iterations,
        scratch,
        result,
    )
}

/// GPU reaching-definition closure over a prepared graph and cached Program.
pub fn reaching_closure_plan_borrowed_into_with_scratch_via<F>(
    dispatch: &F,
    plan: &crate::fixed_point_graph::FixedPointAnalysisPlan,
    seed_bits: &[u32],
    max_iterations: u32,
    scratch: &mut crate::fixed_point_scratch::FixedPointScratch,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let mut result = crate::staging_reserve::reserved_vec(
        crate::dispatch_decode::bitset_word_capacity(
            "reaching_closure planned result",
            plan.graph().node_count,
        )?,
        "fixed-point result word",
    )?;
    reaching_closure_plan_borrowed_into_result_with_scratch_via(
        dispatch,
        plan,
        seed_bits,
        max_iterations,
        scratch,
        &mut result,
    )?;
    Ok(result)
}

/// GPU reaching-definition closure over a prepared graph and cached Program
/// into caller-owned result storage.
pub fn reaching_closure_plan_borrowed_into_result_with_scratch_via<F>(
    dispatch: &F,
    plan: &crate::fixed_point_graph::FixedPointAnalysisPlan,
    seed_bits: &[u32],
    max_iterations: u32,
    scratch: &mut crate::fixed_point_scratch::FixedPointScratch,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    plan.require_kind(
        crate::fixed_point_graph::FixedPointAnalysisKind::Reaching,
        "reaching_closure_plan_borrowed_into_with_scratch_via",
    )?;
    reaching_closure_prepared_program_borrowed_into_result_with_scratch_via(
        dispatch,
        plan.graph(),
        plan.program(),
        seed_bits,
        max_iterations,
        scratch,
        result,
    )
}

fn reaching_closure_prepared_program_borrowed_into_result_with_scratch_via<F>(
    dispatch: &F,
    graph: &crate::fixed_point_graph::FixedPointForwardGraph,
    program: &vyre::ir::Program,
    seed_bits: &[u32],
    max_iterations: u32,
    scratch: &mut crate::fixed_point_scratch::FixedPointScratch,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    graph.require_kind(
        crate::fixed_point_graph::FixedPointAnalysisKind::Reaching,
        "reaching_closure_prepared_program_borrowed_into_result_with_scratch_via",
    )?;
    crate::dispatch_decode::require_positive_iterations("reaching_closure", max_iterations)?;
    let words = crate::dispatch_decode::bitset_word_capacity("reaching_closure", graph.node_count)?;
    crate::dispatch_decode::require_bitset_words("reaching_closure", seed_bits, words)?;
    crate::dispatch_decode::require_bitset_tail_clear(
        "reaching_closure",
        seed_bits,
        graph.node_count,
    )?;
    scratch.prepare_frontier_words(words, seed_bits)?;
    scratch.begin_frontier_density(graph.node_count);
    scratch.prepare_next_words(words)?;
    for _ in 0..max_iterations {
        crate::dispatch_decode::try_pack_u32_into(&scratch.frontier_words, &mut scratch.frontier_bytes, "weir u32 byte staging")?;
        let inputs = [
            graph.pg_nodes_bytes.as_slice(),
            graph.edge_offsets_bytes.as_slice(),
            graph.edge_targets_bytes.as_slice(),
            graph.edge_kind_mask_bytes.as_slice(),
            graph.pg_node_tags_bytes.as_slice(),
            scratch.frontier_bytes.as_slice(),
            scratch.frontier_bytes.as_slice(),
        ];
        dispatch(
            program,
            &inputs,
            Some([graph.node_count.max(1), 1, 1]),
            &mut scratch.outputs,
        )?;
        crate::dispatch_decode::unpack_only_exact_u32_into(
            &scratch.outputs,
            "reaching_closure",
            "frontier",
            words,
            &mut scratch.next,
        )?;
        crate::dispatch_decode::require_bitset_tail_clear(
            "reaching_closure output",
            &scratch.next,
            graph.node_count,
        )?;
        scratch.record_decoded_frontier_transition();
        if crate::fixed_point_scratch::copy_if_converged(
            &scratch.next,
            &scratch.frontier_words,
            result,
        )? {
            return Ok(());
        }
        std::mem::swap(&mut scratch.frontier_words, &mut scratch.next);
    }
    Err(format!(
        "reaching_closure did not converge within {max_iterations} iterations"
    ))
}