weirflow 0.1.0

#![allow(clippy::too_many_arguments)]

use super::{try_range_propagate_with_count, RangePropagateScratch};
use crate::dispatch_input_cache::OwnedDispatchInputs;

pub fn range_propagate_via<F>(
    dispatch: &F,
    defs_in: &[u32],
    edges_in: &[u32],
    var_count: 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();
    range_propagate_borrowed_via(
        &|program, inputs, grid| owned_inputs.dispatch(program, inputs, grid, dispatch),
        defs_in,
        edges_in,
        var_count,
    )
}

/// GPU range propagation using borrowed dispatch inputs.
///
/// The owning [`range_propagate_via`] wrapper is retained for compatibility.
pub fn range_propagate_borrowed_via<F>(
    dispatch: &F,
    defs_in: &[u32],
    edges_in: &[u32],
    var_count: u32,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    let mut scratch = RangePropagateScratch::default();
    range_propagate_borrowed_with_scratch_via(dispatch, defs_in, edges_in, var_count, &mut scratch)
}

/// GPU range propagation using borrowed dispatch inputs and caller-owned
/// staging/output scratch.
pub fn range_propagate_borrowed_with_scratch_via<F>(
    dispatch: &F,
    defs_in: &[u32],
    edges_in: &[u32],
    var_count: u32,
    scratch: &mut RangePropagateScratch,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    let slots = range_slots(var_count)?;
    let mut result = crate::staging_reserve::reserved_vec(slots, "range result slot")?;
    range_propagate_borrowed_into_result_with_scratch_via(
        dispatch,
        defs_in,
        edges_in,
        var_count,
        scratch,
        &mut result,
    )?;
    Ok(result)
}

/// GPU range propagation using caller-owned staging/output scratch and
/// caller-owned decoded result storage.
pub fn range_propagate_borrowed_into_result_with_scratch_via<F>(
    dispatch: &F,
    defs_in: &[u32],
    edges_in: &[u32],
    var_count: u32,
    scratch: &mut RangePropagateScratch,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>) -> Result<Vec<Vec<u8>>, String>,
{
    range_propagate_staged_into_result_via(
        &|program, inputs, grid, outputs| {
            let result = dispatch(program, inputs, grid)?;
            crate::output_scratch::replace_outputs_preserving_slots(outputs, result);
            Ok(())
        },
        defs_in,
        edges_in,
        var_count,
        &mut scratch.outputs,
        &mut scratch.defs_bytes,
        &mut scratch.edges_bytes,
        &mut scratch.out_bytes,
        result,
    )
}

/// GPU range propagation using borrowed dispatch inputs and caller-owned
/// output storage.
pub fn range_propagate_borrowed_into_via<F>(
    dispatch: &F,
    defs_in: &[u32],
    edges_in: &[u32],
    var_count: u32,
    outputs: &mut Vec<Vec<u8>>,
) -> Result<Vec<u32>, String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let slots = range_slots(var_count)?;
    let byte_capacity = slots
        .max(1)
        .checked_mul(std::mem::size_of::<u32>())
        .ok_or_else(|| {
            "range_propagate byte capacity overflows host usize. Fix: shard range propagation before GPU dispatch."
                .to_string()
        })?;
    let mut defs_bytes = crate::staging_reserve::reserved_vec(byte_capacity, "range defs byte")?;
    let mut edges_bytes = crate::staging_reserve::reserved_vec(byte_capacity, "range edge byte")?;
    let mut out_bytes = crate::staging_reserve::reserved_vec(byte_capacity, "range output byte")?;
    let mut result = crate::staging_reserve::reserved_vec(slots, "range result slot")?;
    range_propagate_staged_into_result_via(
        dispatch,
        defs_in,
        edges_in,
        var_count,
        outputs,
        &mut defs_bytes,
        &mut edges_bytes,
        &mut out_bytes,
        &mut result,
    )?;
    Ok(result)
}

fn range_propagate_staged_into_result_via<F>(
    dispatch: &F,
    defs_in: &[u32],
    edges_in: &[u32],
    var_count: u32,
    outputs: &mut Vec<Vec<u8>>,
    defs_bytes: &mut Vec<u8>,
    edges_bytes: &mut Vec<u8>,
    out_bytes: &mut Vec<u8>,
    result: &mut Vec<u32>,
) -> Result<(), String>
where
    F: Fn(&vyre::ir::Program, &[&[u8]], Option<[u32; 3]>, &mut Vec<Vec<u8>>) -> Result<(), String>,
{
    let slots = range_slots(var_count)?;
    let dispatch_slots = slots.max(1);
    let program = try_range_propagate_with_count("defs", "edges", "out", var_count)?;
    crate::dispatch_decode::pack_exact_u32_slots_into(
        defs_in,
        slots,
        "range_propagate defs",
        defs_bytes,
    )?;
    crate::dispatch_decode::pack_exact_u32_slots_into(
        edges_in,
        slots,
        "range_propagate edges",
        edges_bytes,
    )?;
    crate::dispatch_decode::pad_dispatch_min_words(defs_bytes, dispatch_slots)?;
    crate::dispatch_decode::pad_dispatch_min_words(edges_bytes, dispatch_slots)?;
    let out_byte_len = dispatch_slots.checked_mul(4).ok_or_else(|| {
        "range_propagate output byte length overflows host usize. Fix: shard range propagation before GPU dispatch.".to_string()
    })?;
    crate::dispatch_decode::try_write_zero_bytes(
        out_bytes,
        out_byte_len,
        "range_propagate output zero staging",
    )?;
    let inputs = [
        defs_bytes.as_slice(),
        edges_bytes.as_slice(),
        out_bytes.as_slice(),
    ];
    dispatch(&program, &inputs, Some([var_count.max(1), 1, 1]), outputs)?;
    crate::dispatch_decode::unpack_only_exact_u32_into(
        outputs,
        "range_propagate",
        "range",
        slots,
        result,
    )
}

fn range_slots(var_count: u32) -> Result<usize, String> {
    crate::dispatch_decode::u32_to_usize(var_count, "range_propagate variable count")?
        .checked_mul(2)
        .ok_or_else(|| {
            format!(
                "range_propagate var_count={var_count} overflows host slot indexing. Fix: shard range propagation before GPU dispatch."
            )
        })
}