weirflow 0.1.0

use super::prepared::{
    solve_prepared_borrowed_many_with_scratch_via, solve_prepared_borrowed_via,
    solve_prepared_borrowed_with_scratch_via,
};
use crate::dispatch_input_cache::OwnedDispatchInputs;
use crate::ifds_csr_prepare::{
    prepare_ifds_csr_borrowed_via, prepare_ifds_csr_borrowed_with_scratch_via,
};
use crate::ifds_gpu::validate_ifds_problem;
use crate::ifds_resident_types::IfdsBorrowedSolveScratch;
use vyre_foundation::ir::Program;

/// Solve IFDS reachability through a caller-provided owned-buffer GPU dispatch.
///
/// This compatibility entry point preserves older call sites. New production
/// integrations should prefer [`solve_borrowed_via`] so invariant CSR and
/// ProgramGraph buffers are borrowed into each dispatch instead of cloned.
#[allow(clippy::too_many_arguments)]
pub fn solve_via<F>(
    dispatch: &F,
    num_procs: u32,
    blocks_per_proc: u32,
    facts_per_proc: u32,
    intra_edges: &[(u32, u32, u32)],
    inter_edges: &[(u32, u32, u32, u32)],
    flow_gen: &[(u32, u32, u32)],
    flow_kill: &[(u32, u32, u32)],
    seed_facts: &[(u32, u32, u32)],
    max_iterations: u32,
) -> Result<Vec<u32>, String>
where
    F: Fn(&Program, &[Vec<u8>], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    let owned_inputs = OwnedDispatchInputs::new();
    let dispatch_borrowed = |program: &Program,
                             inputs: &[&[u8]],
                             grid_override: Option<[u32; 3]>| {
        owned_inputs.dispatch_refreshing_slots(program, inputs, grid_override, &[5, 6], dispatch)
    };
    solve_borrowed_via(
        &dispatch_borrowed,
        num_procs,
        blocks_per_proc,
        facts_per_proc,
        intra_edges,
        inter_edges,
        flow_gen,
        flow_kill,
        seed_facts,
        max_iterations,
    )
}

/// Solve IFDS reachability through a caller-provided borrowed-buffer GPU dispatch.
///
/// This is the production hot path: the exploded CSR is built by GPU dispatch,
/// invariant ProgramGraph buffers are packed once, and each fixpoint iteration
/// borrows those buffers directly into the backend.
#[allow(clippy::too_many_arguments)]
pub fn solve_borrowed_via<F>(
    dispatch: &F,
    num_procs: u32,
    blocks_per_proc: u32,
    facts_per_proc: u32,
    intra_edges: &[(u32, u32, u32)],
    inter_edges: &[(u32, u32, u32, u32)],
    flow_gen: &[(u32, u32, u32)],
    flow_kill: &[(u32, u32, u32)],
    seed_facts: &[(u32, u32, u32)],
    max_iterations: u32,
) -> Result<Vec<u32>, String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    crate::dispatch_decode::require_positive_iterations("weir IFDS GPU solve", max_iterations)?;
    if seed_facts.is_empty() {
        return Ok(Vec::new());
    }
    validate_ifds_problem(
        "weir IFDS GPU solve",
        num_procs,
        blocks_per_proc,
        facts_per_proc,
        intra_edges,
        inter_edges,
        flow_gen,
        flow_kill,
        seed_facts,
    )?;
    let prepared = prepare_ifds_csr_borrowed_via(
        dispatch,
        num_procs,
        blocks_per_proc,
        facts_per_proc,
        intra_edges,
        inter_edges,
        flow_gen,
        flow_kill,
    )?;
    let reached = solve_prepared_borrowed_via(dispatch, &prepared, seed_facts, max_iterations)?;
    for &node in &reached {
        let (proc_id, block_id, fact_id) = vyre_primitives::graph::exploded::decode_node(node);
        if proc_id >= prepared.shape.num_procs
            || block_id >= prepared.shape.blocks_per_proc
            || fact_id >= prepared.shape.facts_per_proc
        {
            return Err(format!(
                "ifds solve_via output encoded node {node} decodes out of domain: proc={proc_id}, block={block_id}, fact={fact_id}. Fix: reject GPU frontier tail bits before returning solver output."
            ));
        }
    }
    Ok(reached)
}

/// Solve IFDS reachability through a borrowed-buffer GPU dispatch while
/// reusing all host-side preparation and fixpoint scratch across calls.
#[allow(clippy::too_many_arguments)]
pub fn solve_borrowed_with_scratch_via<F>(
    dispatch: &F,
    num_procs: u32,
    blocks_per_proc: u32,
    facts_per_proc: u32,
    intra_edges: &[(u32, u32, u32)],
    inter_edges: &[(u32, u32, u32, u32)],
    flow_gen: &[(u32, u32, u32)],
    flow_kill: &[(u32, u32, u32)],
    seed_facts: &[(u32, u32, u32)],
    max_iterations: u32,
    scratch: &mut IfdsBorrowedSolveScratch,
) -> Result<Vec<u32>, String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    crate::dispatch_decode::require_positive_iterations("weir IFDS GPU solve", max_iterations)?;
    if seed_facts.is_empty() {
        return Ok(Vec::new());
    }
    validate_ifds_problem(
        "weir IFDS GPU solve",
        num_procs,
        blocks_per_proc,
        facts_per_proc,
        intra_edges,
        inter_edges,
        flow_gen,
        flow_kill,
        seed_facts,
    )?;
    let prepared = prepare_ifds_csr_borrowed_with_scratch_via(
        dispatch,
        num_procs,
        blocks_per_proc,
        facts_per_proc,
        intra_edges,
        inter_edges,
        flow_gen,
        flow_kill,
        &mut scratch.prepare,
    )?;
    solve_prepared_borrowed_with_scratch_via(
        dispatch,
        &prepared,
        seed_facts,
        max_iterations,
        &mut scratch.solve,
    )
}

/// Solve several IFDS seed sets through one borrowed-buffer GPU preparation,
/// reusing all host-side preparation and fixpoint scratch across queries.
#[allow(clippy::too_many_arguments)]
pub fn solve_borrowed_many_with_scratch_via<F>(
    dispatch: &F,
    num_procs: u32,
    blocks_per_proc: u32,
    facts_per_proc: u32,
    intra_edges: &[(u32, u32, u32)],
    inter_edges: &[(u32, u32, u32, u32)],
    flow_gen: &[(u32, u32, u32)],
    flow_kill: &[(u32, u32, u32)],
    seed_sets: &[&[(u32, u32, u32)]],
    max_iterations: u32,
    scratch: &mut IfdsBorrowedSolveScratch,
) -> Result<Vec<Vec<u32>>, String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    if seed_sets.is_empty() {
        return Ok(Vec::new());
    }
    crate::dispatch_decode::require_positive_iterations(
        "weir IFDS GPU many solve",
        max_iterations,
    )?;
    if seed_sets.iter().all(|seed_facts| seed_facts.is_empty()) {
        validate_ifds_problem(
            "weir IFDS GPU many solve",
            num_procs,
            blocks_per_proc,
            facts_per_proc,
            intra_edges,
            inter_edges,
            flow_gen,
            flow_kill,
            &[],
        )?;
        return crate::staging_reserve::reserved_result_rows(
            seed_sets.len(),
            "borrowed IFDS empty result row",
        );
    }
    validate_ifds_problem(
        "weir IFDS GPU many solve",
        num_procs,
        blocks_per_proc,
        facts_per_proc,
        intra_edges,
        inter_edges,
        flow_gen,
        flow_kill,
        &[],
    )?;
    let prepared = prepare_ifds_csr_borrowed_with_scratch_via(
        dispatch,
        num_procs,
        blocks_per_proc,
        facts_per_proc,
        intra_edges,
        inter_edges,
        flow_gen,
        flow_kill,
        &mut scratch.prepare,
    )?;
    solve_prepared_borrowed_many_with_scratch_via(
        dispatch,
        &prepared,
        seed_sets,
        max_iterations,
        &mut scratch.solve,
    )
}

/// Solve several IFDS seed sets through one borrowed-buffer GPU preparation.
#[allow(clippy::too_many_arguments)]
pub fn solve_borrowed_many_via<F>(
    dispatch: &F,
    num_procs: u32,
    blocks_per_proc: u32,
    facts_per_proc: u32,
    intra_edges: &[(u32, u32, u32)],
    inter_edges: &[(u32, u32, u32, u32)],
    flow_gen: &[(u32, u32, u32)],
    flow_kill: &[(u32, u32, u32)],
    seed_sets: &[&[(u32, u32, u32)]],
    max_iterations: u32,
) -> Result<Vec<Vec<u32>>, String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    let mut scratch = IfdsBorrowedSolveScratch::default();
    solve_borrowed_many_with_scratch_via(
        &|program, inputs, grid, outputs| {
            let result = dispatch(program, inputs, grid)?;
            crate::output_scratch::replace_outputs_preserving_slots(outputs, result);
            Ok(())
        },
        num_procs,
        blocks_per_proc,
        facts_per_proc,
        intra_edges,
        inter_edges,
        flow_gen,
        flow_kill,
        seed_sets,
        max_iterations,
        &mut scratch,
    )
}