weirflow 0.1.0

use super::{DispatchU32Schema, U32_BYTES};

#[cold]
fn only_output_mismatch(stage: &str, output: &str, got: usize) -> String {
    format!(
        "{stage} dispatch returned {got} output buffers, expected exactly one {output} buffer. Fix: backend must return only {output}."
    )
}

#[cold]
fn output_count_mismatch(
    stage: &str,
    expected_count: usize,
    output_names: &str,
    got: usize,
) -> String {
    format!(
        "{stage} dispatch returned {got} output buffers, expected exactly {expected_count}: {output_names}. Fix: backend must return the declared output ABI and no extras."
    )
}

/// Require exactly one GPU output buffer and return it as bytes.
///
/// # Errors
///
/// Returns an actionable backend-contract error when the stage returns no
/// output or extra output buffers.
#[inline]
pub(crate) fn only_output<'a>(
    outputs: &'a [Vec<u8>],
    stage: &str,
    output: &str,
) -> Result<&'a [u8], String> {
    if outputs.len() != 1 {
        // Wrong output count is a cold backend-contract violation.
        return Err(only_output_mismatch(stage, output, outputs.len()));
    }
    Ok(outputs[0].as_slice())
}

/// Require an exact GPU output-buffer count for multi-output stages.
#[inline]
pub(crate) fn require_output_count(
    outputs: &[Vec<u8>],
    stage: &str,
    expected_count: usize,
    output_names: &str,
) -> Result<(), String> {
    if outputs.len() != expected_count {
        // Wrong output count is a cold backend-contract violation.
        return Err(output_count_mismatch(
            stage,
            expected_count,
            output_names,
            outputs.len(),
        ));
    }
    Ok(())
}

/// Require one GPU output buffer and decode exactly `words` little-endian `u32`s into scratch.
///
/// This is the default ABI boundary for single-output Weir primitives: it
/// rejects missing outputs, extra outputs, short buffers, and trailing bytes in
/// one place before any primitive-specific interpretation happens.
#[inline]
pub(crate) fn unpack_only_exact_u32_into(
    outputs: &[Vec<u8>],
    stage: &str,
    output: &str,
    words: usize,
    out: &mut Vec<u32>,
) -> Result<(), String> {
    let bytes = only_output(outputs, stage, output)?;
    unpack_exact_u32_into(bytes, words, stage, out)
}

/// Decode exactly `words` little-endian `u32` values from a GPU output buffer.
///
/// # Errors
///
/// Returns an actionable backend-contract error if the byte length is not
/// exactly `words * 4`.
#[inline]
pub(crate) fn unpack_exact_u32(bytes: &[u8], words: usize, name: &str) -> Result<Vec<u32>, String> {
    DispatchU32Schema::output(name, words).require_output_bytes(bytes.len())?;
    let mut out = crate::staging_reserve::reserved_vec(words, "exact u32 decode")?;
    vyre_primitives::wire::unpack_u32_slice_into(bytes, words, name, &mut out)?;
    Ok(out)
}

/// Decode one little-endian `u32` from a GPU output buffer without allocation.
#[inline]
pub(crate) fn unpack_exact_u32_scalar(bytes: &[u8], name: &str) -> Result<u32, String> {
    DispatchU32Schema::scalar_output(name).require_scalar_output_bytes(bytes.len())?;
    vyre_primitives::wire::read_u32_le_word(bytes, 0, name)
}

/// Decode exactly `words` little-endian `u32` values into caller-owned scratch.
///
/// # Errors
///
/// Returns an actionable backend-contract error if the byte length is not
/// exactly `words * 4`.
#[inline]
pub(crate) fn unpack_exact_u32_into(
    bytes: &[u8],
    words: usize,
    name: &str,
    out: &mut Vec<u32>,
) -> Result<(), String> {
    DispatchU32Schema::output(name, words).require_output_bytes(bytes.len())?;
    crate::staging_reserve::reserve_vec(out, words, "exact u32 decode")?;
    vyre_primitives::wire::unpack_u32_slice_into(bytes, words, name, out)?;
    Ok(())
}

/// Decode all bytes as little-endian `u32` values into caller-owned scratch.
///
/// # Errors
///
/// Returns an error if the byte length is not a multiple of 4.
#[inline]
pub fn unpack_u32_from(bytes: &[u8], out: &mut Vec<u32>) -> Result<(), String> {
    if bytes.len() % U32_BYTES != 0 {
        return Err(format!(
            "unpack_u32_from: byte length {} is not a multiple of 4. Fix: backend output must contain whole u32 words.",
            bytes.len()
        ));
    }
    let words = bytes.len() / U32_BYTES;
    unpack_exact_u32_into(bytes, words, "unpack_u32_from", out)
}

#[cfg(test)]
mod tests {
    use super::{
        only_output, require_output_count, unpack_exact_u32, unpack_exact_u32_into,
        unpack_exact_u32_scalar, unpack_only_exact_u32_into,
    };
    use crate::dispatch_decode::{
        require_bitset_tail_clear, require_csr_offsets_targets, require_csr_shape,
        require_positive_iterations,
    };

    #[test]
    fn require_csr_shape_rejects_nonzero_first_offset() {
        let err = require_csr_shape("test", 2, &[1, 1, 1], &[0], &[0])
            .expect_err("nonzero first CSR offset must be rejected");
        assert!(
            err.contains("edge_offsets[0]"),
            "unexpected diagnostic: {err}"
        );
    }

    #[test]
    fn require_csr_shape_rejects_zero_node_graph_with_edges() {
        let err = require_csr_shape("test", 0, &[1], &[0], &[0])
            .expect_err("zero-node CSR graph must not declare edges");
        assert!(
            err.contains("edge_offsets[0]"),
            "unexpected diagnostic: {err}"
        );
    }

    #[test]
    fn require_csr_offsets_targets_rejects_target_outside_domain() {
        let err = require_csr_offsets_targets("test", 2, &[0, 1, 1], &[2])
            .expect_err("CSR target outside node_count must be rejected");
        assert!(
            err.contains("targets node 2"),
            "unexpected diagnostic: {err}"
        );
    }

    #[test]
    fn require_bitset_tail_clear_rejects_out_of_domain_bits() {
        let err = require_bitset_tail_clear("test", &[0b1000], 3)
            .expect_err("tail bit beyond three-bit domain must be rejected");
        assert!(
            err.contains("outside the declared domain"),
            "unexpected diagnostic: {err}"
        );
    }

    #[test]
    fn require_bitset_tail_clear_accepts_full_words_and_empty_domains() {
        require_bitset_tail_clear("test", &[u32::MAX], 32)
            .expect("full u32 domain has no tail bits");
        require_bitset_tail_clear("test", &[], 0)
            .expect("empty domain with empty bitset has no tail bits");
        require_bitset_tail_clear("test", &[0], 0)
            .expect("empty domain with backend-minimum zero word has no set tail bits");
    }

    #[test]
    fn require_bitset_tail_clear_rejects_set_bits_in_empty_domain() {
        let err = require_bitset_tail_clear("test", &[1], 0)
            .expect_err("nonzero word in empty domain must be rejected");
        assert!(err.contains("empty domain"), "unexpected diagnostic: {err}");
    }

    #[test]
    fn require_positive_iterations_rejects_zero_budget() {
        let err = require_positive_iterations("test", 0)
            .expect_err("zero fixpoint budget must be rejected");
        assert!(
            err.contains("max_iterations=0"),
            "unexpected diagnostic: {err}"
        );
    }

    #[test]
    fn only_output_rejects_missing_and_extra_backend_outputs() {
        let missing: Vec<Vec<u8>> = Vec::new();
        let err = only_output(&missing, "decode-test", "frontier")
            .expect_err("missing backend output must be rejected");
        assert!(
            err.contains("returned 0 output buffers")
                && err.contains("expected exactly one frontier buffer")
                && err.contains("Fix:"),
            "unexpected missing-output diagnostic: {err}"
        );

        let extra = vec![vec![0; 4], vec![0; 4]];
        let err = only_output(&extra, "decode-test", "frontier")
            .expect_err("extra backend output must be rejected");
        assert!(
            err.contains("returned 2 output buffers")
                && err.contains("expected exactly one frontier buffer")
                && err.contains("Fix:"),
            "unexpected extra-output diagnostic: {err}"
        );
    }

    #[test]
    fn require_output_count_rejects_missing_and_extra_multi_outputs() {
        let one = vec![vec![0; 4]];
        let err = require_output_count(&one, "csr-prepare", 4, "row_ptr,col_idx,len,flags")
            .expect_err("missing multi-output buffers must be rejected");
        assert!(
            err.contains("returned 1 output buffers")
                && err.contains("expected exactly 4")
                && err.contains("no extras"),
            "unexpected missing multi-output diagnostic: {err}"
        );

        let five = vec![vec![0; 4]; 5];
        let err = require_output_count(&five, "csr-prepare", 4, "row_ptr,col_idx,len,flags")
            .expect_err("extra multi-output buffers must be rejected");
        assert!(
            err.contains("returned 5 output buffers")
                && err.contains("expected exactly 4")
                && err.contains("no extras"),
            "unexpected extra multi-output diagnostic: {err}"
        );
    }

    #[test]
    fn exact_u32_decode_rejects_short_and_trailing_bytes() {
        let short = unpack_exact_u32(&[1, 2, 3], 1, "decode-test")
            .expect_err("short u32 output must be rejected");
        assert!(
            short.contains("got 3 bytes")
                && short.contains("expected exactly 4")
                && short.contains("no trailing bytes"),
            "unexpected short-output diagnostic: {short}"
        );

        let trailing = unpack_exact_u32(&[1, 2, 3, 4, 5], 1, "decode-test")
            .expect_err("trailing backend bytes must be rejected");
        assert!(
            trailing.contains("got 5 bytes")
                && trailing.contains("expected exactly 4")
                && trailing.contains("no trailing bytes"),
            "unexpected trailing-output diagnostic: {trailing}"
        );
    }

    #[test]
    fn exact_u32_decode_into_reuses_stable_scratch_capacity() {
        let mut out = Vec::with_capacity(8);
        out.extend([99, 100]);
        let capacity = out.capacity();

        unpack_exact_u32_into(&[1, 0, 0, 0, 2, 0, 0, 0], 2, "decode-test", &mut out)
            .expect("exact output should decode into caller-owned scratch");

        assert_eq!(out, vec![1, 2]);
        assert_eq!(
            out.capacity(),
            capacity,
            "exact decode into stable scratch must not discard reusable allocation"
        );
    }

    #[test]
    fn scalar_and_single_output_decode_share_exact_abi_contract() {
        let scalar = unpack_exact_u32_scalar(&[7, 0, 0, 0], "scalar-test")
            .expect("four-byte scalar output should decode");
        assert_eq!(scalar, 7);
        let err = unpack_exact_u32_scalar(&[7, 0, 0, 0, 0], "scalar-test")
            .expect_err("scalar output with trailing bytes must be rejected");
        assert!(
            err.contains("expected exactly 4") && err.contains("no trailing bytes"),
            "unexpected scalar diagnostic: {err}"
        );

        let outputs = vec![vec![3, 0, 0, 0, 4, 0, 0, 0]];
        let mut decoded = Vec::new();
        unpack_only_exact_u32_into(&outputs, "single-output-test", "frontier", 2, &mut decoded)
            .expect("single output with exact words should decode");
        assert_eq!(decoded, vec![3, 4]);
    }
}

#[cfg(test)]
mod extra_tests {
    use super::{
        only_output, require_output_count, unpack_exact_u32, unpack_exact_u32_scalar,
        unpack_only_exact_u32_into,
    };
    use crate::dispatch_decode::{
        require_bitset_tail_clear, require_csr_offsets_targets, require_csr_shape,
        require_positive_iterations,
    };

    #[test]
    fn require_csr_offsets_targets_rejects_target_equal_to_node_count() {
        let err = require_csr_offsets_targets("test", 2, &[0, 1, 1], &[2])
            .expect_err("target equal to node_count must be rejected");
        assert!(err.contains("targets node 2"), "{err}");
    }

    #[test]
    fn require_csr_shape_rejects_missing_mask() {
        let err = require_csr_shape("test", 2, &[0, 1, 1], &[0], &[])
            .expect_err("missing mask must be rejected");
        assert!(err.contains("edge_kind_mask"), "{err}");
    }

    #[test]
    fn require_positive_iterations_rejects_zero() {
        let err = require_positive_iterations("test", 0).expect_err("zero budget must be rejected");
        assert!(err.contains("max_iterations=0"), "{err}");
    }

    #[test]
    fn require_bitset_tail_clear_rejects_set_bits_in_tail() {
        let err = require_bitset_tail_clear("test", &[0b10000], 3)
            .expect_err("set tail bits must be rejected");
        assert!(err.contains("outside the declared domain"), "{err}");
    }

    #[test]
    fn unpack_exact_u32_rejects_empty_buffer_for_one_word() {
        let err =
            unpack_exact_u32(&[], 1, "decode-test").expect_err("empty buffer must be rejected");
        assert!(err.contains("got 0 bytes"), "{err}");
    }

    #[test]
    fn unpack_exact_u32_scalar_rejects_three_bytes() {
        let err = unpack_exact_u32_scalar(&[1, 2, 3], "scalar-test")
            .expect_err("three bytes must be rejected");
        assert!(err.contains("expected exactly 4"), "{err}");
    }

    #[test]
    fn only_output_rejects_empty() {
        let err = only_output(&[], "test", "frontier").expect_err("empty outputs must be rejected");
        assert!(err.contains("returned 0 output buffers"), "{err}");
    }

    #[test]
    fn require_output_count_rejects_zero_for_one() {
        let err =
            require_output_count(&[], "test", 1, "out").expect_err("zero outputs must be rejected");
        assert!(err.contains("returned 0 output buffers"), "{err}");
    }

    #[test]
    fn unpack_only_exact_u32_into_rejects_extra_outputs() {
        let outputs = vec![vec![1, 0, 0, 0], vec![2, 0, 0, 0]];
        let mut decoded = Vec::new();
        let err = unpack_only_exact_u32_into(&outputs, "test", "frontier", 1, &mut decoded)
            .expect_err("extra outputs must be rejected");
        assert!(err.contains("returned 2 output buffers"), "{err}");
    }
}