weirflow 0.1.0

use super::{backward_slice, prepare_slice_graph};
use crate::dispatch_input_cache::OwnedDispatchInputs;
use vyre_primitives::graph::program_graph::ProgramGraphShape;

pub fn slice_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();
    slice_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 backward-slice closure using borrowed dispatch inputs.
///
/// Prefer this over [`slice_closure_via`] when the backend exposes borrowed
/// dispatch; it avoids rebuilding invariant graph byte buffers per iteration.
pub fn slice_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("slice_closure", max_iterations)?;
    crate::dispatch_decode::require_bitset_tail_clear("slice_closure", seed_bits, node_count)?;
    let out = slice_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("slice_closure output", &out, node_count)?;
    Ok(out)
}

/// GPU backward-slice closure using borrowed dispatch inputs and caller-owned
/// output storage.
pub fn slice_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();
    slice_closure_borrowed_into_with_scratch_via(
        dispatch,
        node_count,
        edge_offsets,
        edge_targets,
        edge_kind_mask,
        seed_bits,
        max_iterations,
        &mut scratch,
    )
}

pub fn slice_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 graph = prepare_slice_graph(node_count, edge_offsets, edge_targets, edge_kind_mask)?;
    slice_closure_prepared_borrowed_into_with_scratch_via(
        dispatch,
        &graph,
        seed_bits,
        max_iterations,
        scratch,
    )
}

/// GPU backward-slice closure into caller-owned final result storage.
pub fn slice_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>,
{
    let graph = prepare_slice_graph(node_count, edge_offsets, edge_targets, edge_kind_mask)?;
    slice_closure_prepared_borrowed_into_result_with_scratch_via(
        dispatch,
        &graph,
        seed_bits,
        max_iterations,
        scratch,
        result,
    )
}

/// GPU backward-slice closure against prepared graph buffers.
pub fn slice_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 program = backward_slice(
        ProgramGraphShape::new(graph.node_count(), graph.edge_count()),
        "fin",
        "fout",
    );
    slice_closure_prepared_program_borrowed_into_with_scratch_via(
        dispatch,
        graph,
        &program,
        seed_bits,
        max_iterations,
        scratch,
    )
}

/// GPU backward-slice closure against prepared graph buffers into caller-owned
/// final result storage.
pub fn slice_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 = backward_slice(
        ProgramGraphShape::new(graph.node_count(), graph.edge_count()),
        "fin",
        "fout",
    );
    slice_closure_prepared_program_borrowed_into_result_with_scratch_via(
        dispatch,
        graph,
        &program,
        seed_bits,
        max_iterations,
        scratch,
        result,
    )
}

/// GPU backward-slice closure against a prepared graph+Program plan.
pub fn slice_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>,
{
    plan.require_kind(
        crate::fixed_point_graph::FixedPointAnalysisKind::Slice,
        "slice_closure_plan_borrowed_into_with_scratch_via",
    )?;
    slice_closure_prepared_program_borrowed_into_with_scratch_via(
        dispatch,
        plan.graph(),
        plan.program(),
        seed_bits,
        max_iterations,
        scratch,
    )
}

/// GPU backward-slice closure against a prepared graph+Program plan into
/// caller-owned final result storage.
pub fn slice_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::Slice,
        "slice_closure_plan_borrowed_into_result_with_scratch_via",
    )?;
    slice_closure_prepared_program_borrowed_into_result_with_scratch_via(
        dispatch,
        plan.graph(),
        plan.program(),
        seed_bits,
        max_iterations,
        scratch,
        result,
    )
}

fn slice_closure_prepared_program_borrowed_into_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<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(
            "slice_closure prepared result",
            graph.node_count(),
        )?,
        "fixed-point result word",
    )?;
    slice_closure_prepared_program_borrowed_into_result_with_scratch_via(
        dispatch,
        graph,
        program,
        seed_bits,
        max_iterations,
        scratch,
        &mut result,
    )?;
    Ok(result)
}

fn slice_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::Slice,
        "slice_closure_prepared_program_borrowed_into_result_with_scratch_via",
    )?;
    crate::dispatch_decode::require_positive_iterations("slice_closure", max_iterations)?;
    let words = crate::dispatch_decode::bitset_word_capacity("slice_closure", graph.node_count())?;
    crate::dispatch_decode::require_bitset_words("slice_closure", seed_bits, words)?;
    crate::dispatch_decode::require_bitset_tail_clear(
        "slice_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,
            "slice_closure",
            "slice frontier",
            words,
            &mut scratch.next,
        )?;
        crate::dispatch_decode::require_bitset_tail_clear(
            "slice_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!(
        "slice_closure did not converge within {max_iterations} iterations"
    ))
}