weirflow 0.1.0

use super::solve_core::{
    solve_prepared_borrowed_with_parts_into_via, solve_prepared_borrowed_with_parts_via,
};
use crate::ifds_resident_types::{IfdsSolveScratch, PreparedIfdsCsr};
use vyre_foundation::ir::Program;

/// Solve IFDS reachability over a CSR prepared by [`prepare_ifds_csr_borrowed_via`].
pub fn solve_prepared_borrowed_via<F>(
    dispatch: &F,
    prepared: &PreparedIfdsCsr,
    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>,
{
    if seed_facts.is_empty() {
        crate::dispatch_decode::require_positive_iterations(
            "weir IFDS prepared borrowed solve",
            max_iterations,
        )?;
        return Ok(Vec::new());
    }
    let mut outputs = crate::staging_reserve::reserved_vec(1, "borrowed IFDS output slot")?;
    solve_prepared_borrowed_into_via(
        &|program, inputs, grid, outputs| {
            let result = dispatch(program, inputs, grid)?;
            crate::output_scratch::replace_outputs_preserving_slots(outputs, result);
            Ok(())
        },
        prepared,
        seed_facts,
        max_iterations,
        &mut outputs,
    )
}

/// Solve IFDS reachability over a prepared CSR while adapting a borrowed
/// Vec-returning dispatch into caller-owned solve scratch.
pub fn solve_prepared_borrowed_with_adapter_scratch_via<F>(
    dispatch: &F,
    prepared: &PreparedIfdsCsr,
    seed_facts: &[(u32, u32, u32)],
    max_iterations: u32,
    scratch: &mut IfdsSolveScratch,
) -> Result<Vec<u32>, String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    if seed_facts.is_empty() {
        crate::dispatch_decode::require_positive_iterations(
            "weir IFDS prepared borrowed with adapter scratch solve",
            max_iterations,
        )?;
        return Ok(Vec::new());
    }
    solve_prepared_borrowed_with_scratch_via(
        &|program, inputs, grid, outputs| {
            let result = dispatch(program, inputs, grid)?;
            crate::output_scratch::replace_outputs_preserving_slots(outputs, result);
            Ok(())
        },
        prepared,
        seed_facts,
        max_iterations,
        scratch,
    )
}

/// Solve IFDS reachability over a prepared CSR while reusing caller-owned
/// dispatch output storage across fixpoint iterations.
pub fn solve_prepared_borrowed_into_via<F>(
    dispatch: &F,
    prepared: &PreparedIfdsCsr,
    seed_facts: &[(u32, u32, u32)],
    max_iterations: u32,
    outputs: &mut Vec<Vec<u8>>,
) -> Result<Vec<u32>, String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let frontier_byte_capacity = prepared
        .words
        .checked_mul(std::mem::size_of::<u32>())
        .ok_or_else(|| {
            "weir IFDS prepared borrowed frontier byte capacity overflows usize. Fix: shard the IFDS problem before GPU solving.".to_string()
        })?;
    let mut frontier =
        crate::staging_reserve::reserved_vec(prepared.words, "borrowed IFDS frontier word")?;
    let mut frontier_bytes = crate::staging_reserve::reserved_vec(
        frontier_byte_capacity,
        "borrowed IFDS frontier byte",
    )?;
    let mut next =
        crate::staging_reserve::reserved_vec(prepared.words, "borrowed IFDS next frontier word")?;
    solve_prepared_borrowed_with_parts_via(
        dispatch,
        prepared,
        seed_facts,
        max_iterations,
        &mut frontier,
        &mut frontier_bytes,
        &mut next,
        outputs,
    )
}

/// Solve IFDS reachability over a prepared CSR while reusing all host-side
/// frontier, decode, and backend-output scratch across calls.
pub fn solve_prepared_borrowed_with_scratch_via<F>(
    dispatch: &F,
    prepared: &PreparedIfdsCsr,
    seed_facts: &[(u32, u32, u32)],
    max_iterations: u32,
    scratch: &mut IfdsSolveScratch,
) -> Result<Vec<u32>, String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    if seed_facts.is_empty() {
        crate::dispatch_decode::require_positive_iterations(
            "weir IFDS prepared borrowed with scratch solve",
            max_iterations,
        )?;
        return Ok(Vec::new());
    }
    let result_capacity = crate::dispatch_decode::u32_to_usize(
        prepared.node_count,
        "weir IFDS prepared solve node count",
    )?;
    let mut result =
        crate::staging_reserve::reserved_vec(result_capacity, "borrowed IFDS result fact")?;
    solve_prepared_borrowed_with_scratch_into_via(
        dispatch,
        prepared,
        seed_facts,
        max_iterations,
        scratch,
        &mut result,
    )?;
    Ok(result)
}

/// Solve IFDS reachability over a prepared CSR into caller-owned result
/// storage while reusing all host-side solve scratch across calls.
pub fn solve_prepared_borrowed_with_scratch_into_via<F>(
    dispatch: &F,
    prepared: &PreparedIfdsCsr,
    seed_facts: &[(u32, u32, u32)],
    max_iterations: u32,
    scratch: &mut IfdsSolveScratch,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    if seed_facts.is_empty() {
        result.clear();
        crate::dispatch_decode::require_positive_iterations(
            "weir IFDS prepared borrowed with scratch into solve",
            max_iterations,
        )?;
        return Ok(());
    }
    solve_prepared_borrowed_with_parts_into_via(
        dispatch,
        prepared,
        seed_facts,
        max_iterations,
        &mut scratch.frontier,
        &mut scratch.frontier_bytes,
        &mut scratch.next,
        &mut scratch.outputs,
        result,
    )
}

/// Solve several IFDS seed sets over one prepared CSR while reusing all
/// host-side solve scratch between queries.
pub fn solve_prepared_borrowed_many_with_scratch_via<F>(
    dispatch: &F,
    prepared: &PreparedIfdsCsr,
    seed_sets: &[&[(u32, u32, u32)]],
    max_iterations: u32,
    scratch: &mut IfdsSolveScratch,
) -> Result<Vec<Vec<u32>>, String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let mut results =
        crate::staging_reserve::reserved_vec(seed_sets.len(), "borrowed IFDS result row")?;
    solve_prepared_borrowed_many_with_scratch_into_via(
        dispatch,
        prepared,
        seed_sets,
        max_iterations,
        scratch,
        &mut results,
    )?;
    Ok(results)
}

/// Solve several IFDS seed sets over one prepared CSR into caller-owned result
/// rows while reusing all host-side solve scratch between queries.
pub fn solve_prepared_borrowed_many_with_scratch_into_via<F>(
    dispatch: &F,
    prepared: &PreparedIfdsCsr,
    seed_sets: &[&[(u32, u32, u32)]],
    max_iterations: u32,
    scratch: &mut IfdsSolveScratch,
    results: &mut Vec<Vec<u32>>,
) -> Result<(), String>
where
    F: Fn(&Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    if seed_sets.is_empty() {
        results.clear();
        return Ok(());
    }
    if seed_sets.iter().all(|seed_facts| seed_facts.is_empty()) {
        crate::staging_reserve::resize_result_rows(
            results,
            seed_sets.len(),
            "borrowed IFDS result row",
        )?;
        crate::staging_reserve::clear_result_rows(results);
        return Ok(());
    }
    crate::staging_reserve::resize_result_rows(
        results,
        seed_sets.len(),
        "borrowed IFDS result row",
    )?;
    for (query_index, seed_facts) in seed_sets.iter().enumerate() {
        solve_prepared_borrowed_with_scratch_into_via(
            dispatch,
            prepared,
            seed_facts,
            max_iterations,
            scratch,
            &mut results[query_index],
        )
        .map_err(|error| format!("{error} Query index: {query_index}."))?;
    }
    Ok(())
}