weirflow 0.1.0

use super::program_cache::ensure_resident_single_query_programs;
use super::result_slots::resize_result_slots;
use crate::ifds_frontier_decode::ifds_encoded_frontier_nodes_into;
use crate::ifds_resident_alloc::{
    allocate_resident_ifds_scratch_with_seed_capacity, free_resident_ifds_scratch,
};
use crate::ifds_resident_types::{
    IfdsResidentDispatch, ResidentIfdsHostScratch, ResidentIfdsScratch, ResidentPreparedIfdsCsr,
};
use crate::ifds_seed_frontiers::seed_frontier_words_into;

/// Solve IFDS reachability using resident CSR/frontier buffers.
///
/// This keeps invariant CSR buffers on device, grows one resident frontier
/// accumulator in place, and polls only a 4-byte changed flag per iteration.
/// The full frontier is downloaded once, after GPU-side convergence.
pub fn solve_resident_prepared_via<D>(
    dispatch: &D,
    prepared: &ResidentPreparedIfdsCsr<D::Resource>,
    seed_facts: &[(u32, u32, u32)],
    max_iterations: u32,
) -> Result<Vec<u32>, String>
where
    D: IfdsResidentDispatch,
{
    let mut results =
        solve_resident_prepared_many_via(dispatch, prepared, &[seed_facts], max_iterations)?;
    results
        .pop()
        .ok_or_else(|| "weir IFDS resident solve produced no result for one seed set".to_string())
}

/// Solve several IFDS seed queries against one resident CSR with one scratch
/// frontier allocation and one scratch changed-flag allocation.
///
/// This is the multi-query hot path for dataflow clients: graph buffers stay
/// resident, scratch buffers are reused across seed sets, each iteration polls
/// one u32 changed flag, and each query downloads its final frontier once.
pub fn solve_resident_prepared_many_via<D>(
    dispatch: &D,
    prepared: &ResidentPreparedIfdsCsr<D::Resource>,
    seed_sets: &[&[(u32, u32, u32)]],
    max_iterations: u32,
) -> Result<Vec<Vec<u32>>, String>
where
    D: IfdsResidentDispatch,
{
    let max_seed_facts = seed_sets
        .iter()
        .map(|seed_facts| seed_facts.len())
        .max()
        .unwrap_or(1)
        .max(1);
    let scratch =
        allocate_resident_ifds_scratch_with_seed_capacity(dispatch, prepared, max_seed_facts)?;
    let result = solve_resident_prepared_many_with_scratch_via(
        dispatch,
        prepared,
        &scratch,
        seed_sets,
        max_iterations,
    );
    let free_result = free_resident_ifds_scratch(dispatch, scratch);
    match (result, free_result) {
        (Ok(results), Ok(())) => Ok(results),
        (Err(error), Ok(())) => Err(error),
        (Ok(_), Err(error)) => Err(error),
        (Err(solve_error), Err(free_error)) => Err(format!(
            "{solve_error}; additionally, weir IFDS resident scratch free failed: {free_error}"
        )),
    }
}

/// Solve several IFDS seed queries against one resident CSR using caller-owned
/// reusable scratch buffers.
pub fn solve_resident_prepared_many_with_scratch_via<D>(
    dispatch: &D,
    prepared: &ResidentPreparedIfdsCsr<D::Resource>,
    scratch: &ResidentIfdsScratch<D::Resource>,
    seed_sets: &[&[(u32, u32, u32)]],
    max_iterations: u32,
) -> Result<Vec<Vec<u32>>, String>
where
    D: IfdsResidentDispatch,
{
    let mut results =
        crate::staging_reserve::reserved_vec(seed_sets.len(), "resident IFDS result row")?;
    solve_resident_prepared_many_with_scratch_into_via(
        dispatch,
        prepared,
        scratch,
        seed_sets,
        max_iterations,
        &mut results,
    )?;
    Ok(results)
}

/// Solve several IFDS seed queries against one resident CSR using caller-owned
/// reusable scratch buffers and caller-owned result storage.
pub fn solve_resident_prepared_many_with_scratch_into_via<D>(
    dispatch: &D,
    prepared: &ResidentPreparedIfdsCsr<D::Resource>,
    scratch: &ResidentIfdsScratch<D::Resource>,
    seed_sets: &[&[(u32, u32, u32)]],
    max_iterations: u32,
    results: &mut Vec<Vec<u32>>,
) -> Result<(), String>
where
    D: IfdsResidentDispatch,
{
    let mut host_scratch = ResidentIfdsHostScratch::default();
    solve_resident_prepared_many_with_scratch_and_host_into_via(
        dispatch,
        prepared,
        scratch,
        &mut host_scratch,
        seed_sets,
        max_iterations,
        results,
    )
}

/// Solve several IFDS seed queries against one resident CSR using caller-owned
/// reusable resident buffers, host packing/readback scratch, and result storage.
pub fn solve_resident_prepared_many_with_scratch_and_host_into_via<D>(
    dispatch: &D,
    prepared: &ResidentPreparedIfdsCsr<D::Resource>,
    scratch: &ResidentIfdsScratch<D::Resource>,
    host_scratch: &mut ResidentIfdsHostScratch,
    seed_sets: &[&[(u32, u32, u32)]],
    max_iterations: u32,
    results: &mut Vec<Vec<u32>>,
) -> Result<(), String>
where
    D: IfdsResidentDispatch,
{
    if seed_sets.is_empty() {
        results.clear();
        return Ok(());
    }
    crate::dispatch_decode::require_positive_iterations(
        "weir IFDS resident prepared GPU solve",
        max_iterations,
    )?;
    if seed_sets.iter().all(|seed_facts| seed_facts.is_empty()) {
        resize_result_slots(results, seed_sets.len())?;
        for result in results.iter_mut() {
            result.clear();
        }
        return Ok(());
    }
    if scratch.words != prepared.words {
        return Err(format!(
            "weir IFDS resident scratch has {} frontier words but prepared CSR requires {}. Fix: allocate scratch from the same resident IFDS CSR before solving.",
            scratch.words, prepared.words
        ));
    }
    ensure_resident_single_query_programs(prepared, host_scratch, max_iterations)?;
    resize_result_slots(results, seed_sets.len())?;
    let step_resources = [
        prepared.pg_nodes.clone(),
        prepared.row_ptr.clone(),
        prepared.col_idx.clone(),
        prepared.pg_edge_kind_mask.clone(),
        prepared.pg_node_tags.clone(),
        scratch.frontier.clone(),
        scratch.changed.clone(),
    ];
    let ResidentIfdsHostScratch {
        frontier,
        frontier_bytes,
        changed_bytes,
        step_program,
        ..
    } = host_scratch;
    crate::dispatch_decode::try_write_zero_words(frontier, prepared.words, "weir zero word staging")?;
    let program = step_program
        .as_ref()
        .ok_or_else(|| "weir IFDS resident single-query step program cache was not initialized. Fix: call ensure_resident_single_query_programs before solving.".to_string())?;
    for (seed_index, seed_facts) in seed_sets.iter().enumerate() {
        results[seed_index].clear();
        if seed_facts.len() > scratch.max_seed_facts {
            return Err(format!(
                "weir IFDS resident scratch seed capacity is {} facts but seed set {seed_index} contains {} facts. Fix: allocate scratch with allocate_resident_ifds_scratch_with_seed_capacity or shard the seed set.",
                scratch.max_seed_facts,
                seed_facts.len()
            ));
        }
        let has_frontier = seed_frontier_words_into(
            prepared.shape,
            prepared.node_count,
            prepared.words,
            seed_facts,
            frontier,
        )
        .map_err(|error| {
            format!("weir IFDS resident seed frontier for seed set {seed_index} failed: {error}")
        })?;
        if !has_frontier {
            continue;
        }
        crate::dispatch_decode::try_pack_u32_into(frontier, frontier_bytes, "weir u32 byte staging")?;
        dispatch
            .upload_resident_at_many(&[(&scratch.frontier, 0, frontier_bytes.as_slice())])
            .map_err(|error| {
                format!(
                    "weir IFDS resident upload seed frontier for seed set {seed_index} failed: {error}"
                )
            })?;
        let mut converged = false;
        for _iteration in 0..max_iterations {
            crate::dispatch_decode::try_pack_u32_into(&[0], changed_bytes, "weir u32 byte staging")?;
            dispatch
                .upload_resident_at_many(&[(&scratch.changed, 0, changed_bytes.as_slice())])
                .map_err(|error| {
                    format!(
                        "weir IFDS resident clear changed flag for seed set {seed_index} failed: {error}"
                    )
                })?;
            dispatch
                .dispatch_resident(
                    program,
                    step_resources.as_slice(),
                    Some([prepared.node_count.max(1), 1, 1]),
                )
                .map_err(|error| {
                    format!(
                        "weir IFDS resident primitive step for seed set {seed_index} failed: {error}"
                    )
                })?;
            dispatch
                .download_resident_range_into(
                    &scratch.changed,
                    0,
                    std::mem::size_of::<u32>(),
                    changed_bytes,
                )
                .map_err(|error| {
                    format!(
                        "weir IFDS resident download changed flag for seed set {seed_index} failed: {error}"
                    )
                })?;
            let changed_word = crate::dispatch_decode::unpack_exact_u32_scalar(
                changed_bytes,
                "ifds resident solve changed flag",
            )?;
            if changed_word == 0 {
                converged = true;
                break;
            }
        }
        if !converged {
            return Err(format!(
                "IFDS resident solve for seed set {seed_index} did not converge within {max_iterations} iterations"
            ));
        }
        dispatch
            .download_resident_range_into(
                &scratch.frontier,
                0,
                scratch.frontier_byte_len,
                frontier_bytes,
            )
            .map_err(|error| {
                format!(
                    "weir IFDS resident download final frontier for seed set {seed_index} failed: {error}"
                )
            })?;
        crate::dispatch_decode::unpack_exact_u32_into(
            frontier_bytes,
            prepared.words,
            "ifds resident solve final frontier",
            frontier,
        )?;
        crate::dispatch_decode::require_bitset_tail_clear(
            "ifds resident solve output",
            frontier,
            prepared.node_count,
        )?;
        ifds_encoded_frontier_nodes_into(
            frontier,
            prepared.node_count,
            prepared.shape,
            &mut results[seed_index],
        )?;
    }
    results.truncate(seed_sets.len());
    Ok(())
}