weirflow 0.1.0

use super::{live_step, prepare_live_graph_with_scratch};
use crate::dispatch_input_cache::OwnedDispatchInputs;
use vyre_primitives::graph::program_graph::ProgramGraphShape;

pub fn live_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();
    live_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 live-variable closure using borrowed dispatch inputs.
///
/// This is the allocation-aware hot path. The owning [`live_closure_via`]
/// wrapper exists for compatibility with older dispatch adapters.
pub fn live_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("live_closure", max_iterations)?;
    crate::dispatch_decode::require_bitset_tail_clear("live_closure", seed_bits, node_count)?;
    let out = live_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("live_closure output", &out, node_count)?;
    Ok(out)
}

/// GPU live-variable closure using borrowed dispatch inputs and caller-owned
/// output storage.
///
/// This is the zero-output-vector-churn hot path for backends with `_into`
/// dispatch. The owning [`live_closure_via`] and borrowed
/// [`live_closure_borrowed_via`] wrappers are compatibility adapters.
pub fn live_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();
    live_closure_borrowed_into_with_scratch_via(
        dispatch,
        node_count,
        edge_offsets,
        edge_targets,
        edge_kind_mask,
        seed_bits,
        max_iterations,
        &mut scratch,
    )
}

/// GPU live-variable closure using caller-owned fixed-point scratch.
pub fn live_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("live_closure result", node_count)?,
        "fixed-point result word",
    )?;
    live_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 live-variable closure using caller-owned fixed-point scratch and final
/// result storage.
pub fn live_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("live_closure", max_iterations)?;
    let words = crate::dispatch_decode::bitset_word_capacity("live_closure", node_count)?;
    crate::dispatch_decode::require_bitset_words("live_closure", seed_bits, words)?;
    crate::dispatch_decode::require_bitset_tail_clear("live_closure", seed_bits, node_count)?;
    let graph = prepare_live_graph_with_scratch(
        node_count,
        edge_offsets,
        edge_targets,
        edge_kind_mask,
        scratch,
    )?;
    live_closure_prepared_borrowed_into_result_with_scratch_via(
        dispatch,
        &graph,
        seed_bits,
        max_iterations,
        scratch,
        result,
    )
}

pub fn live_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(
            "live_closure prepared result",
            graph.node_count,
        )?,
        "fixed-point result word",
    )?;
    live_closure_prepared_borrowed_into_result_with_scratch_via(
        dispatch,
        graph,
        seed_bits,
        max_iterations,
        scratch,
        &mut result,
    )?;
    Ok(result)
}

/// GPU live-variable closure over prepared reversed-control graph buffers into
/// caller-owned result storage.
pub fn live_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 = live_step(
        ProgramGraphShape::new(graph.node_count, graph.edge_count),
        "fin",
        "fout",
    );
    live_closure_prepared_program_borrowed_into_result_with_scratch_via(
        dispatch,
        graph,
        &program,
        seed_bits,
        max_iterations,
        scratch,
        result,
    )
}

/// GPU live-variable closure over prepared graph buffers and cached Program.
pub fn live_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(
            "live_closure planned result",
            plan.graph().node_count,
        )?,
        "fixed-point result word",
    )?;
    live_closure_plan_borrowed_into_result_with_scratch_via(
        dispatch,
        plan,
        seed_bits,
        max_iterations,
        scratch,
        &mut result,
    )?;
    Ok(result)
}

/// GPU live-variable closure over a prepared graph+Program plan into
/// caller-owned result storage.
pub fn live_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::Live,
        "live_closure_plan_borrowed_into_result_with_scratch_via",
    )?;
    live_closure_prepared_program_borrowed_into_result_with_scratch_via(
        dispatch,
        plan.graph(),
        plan.program(),
        seed_bits,
        max_iterations,
        scratch,
        result,
    )
}

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