weirflow 0.1.0

//! Shared GPU dispatch-output decoders for Weir via paths.
//!
//! Production GPU wrappers must reject malformed backend output instead of
//! accepting prefixes or silently ignoring trailing bytes. Keep the contract in
//! one place so every dataflow primitive reports the same failure mode.

use super::{byte_len_for_u32_words, DispatchU32Schema};

/// Encode every input word as little-endian bytes for GPU dispatch.
#[inline]
#[must_use]
pub(crate) fn pack_u32(values: &[u32]) -> Vec<u8> {
    #[cfg(target_endian = "little")]
    {
        bytemuck::cast_slice(values).to_vec()
    }
    #[cfg(target_endian = "big")]
    {
        let mut bytes = Vec::with_capacity(values.len().saturating_mul(4));
        for value in values {
            bytes.extend_from_slice(&value.to_le_bytes());
        }
        bytes
    }
}

/// Encode every input word as little-endian bytes for GPU dispatch with an
/// explicit allocator failure contract.
#[inline]
pub(crate) fn try_pack_u32(values: &[u32], field: &'static str) -> Result<Vec<u8>, String> {
    let byte_len = byte_len_for_u32_words(values.len(), field)?;
    #[cfg(target_endian = "little")]
    {
        let bytes = bytemuck::cast_slice(values);
        if bytes.len() != byte_len {
            // Cast width mismatch is a cold malformed-input path.
            return Err(pack_u32_cast_mismatch(
                field,
                values.len(),
                bytes.len(),
                byte_len,
            ));
        }
        let mut out = crate::staging_reserve::reserved_vec(byte_len, field)?;
        out.extend_from_slice(bytes);
        Ok(out)
    }
    #[cfg(target_endian = "big")]
    let mut bytes = crate::staging_reserve::reserved_vec(byte_len, field)?;
    #[cfg(target_endian = "big")]
    for value in values {
        bytes.extend_from_slice(&value.to_le_bytes());
    }
    #[cfg(target_endian = "big")]
    Ok(bytes)
}

/// Encode one repeated input word as little-endian bytes for GPU dispatch.
#[inline]
pub(crate) fn pack_repeated_u32(value: u32, words: usize) -> Result<Vec<u8>, String> {
    let byte_len = byte_len_for_u32_words(words, "pack_repeated_u32")?;
    if value == 0 {
        let mut bytes = Vec::new();
        try_write_zero_bytes(&mut bytes, byte_len, "repeated zero u32 byte pack")?;
        return Ok(bytes);
    }
    let mut bytes = crate::staging_reserve::reserved_vec(byte_len, "repeated u32 byte pack")?;
    pack_repeated_u32_into(value, words, &mut bytes)?;
    Ok(bytes)
}

/// Encode every input word as little-endian bytes into caller-owned scratch.
#[inline]
#[cfg(any(test, feature = "legacy-infallible"))]
pub fn pack_u32_into(values: &[u32], bytes: &mut Vec<u8>) {
    try_pack_u32_into(values, bytes, "u32 byte pack")
        .expect("u32 byte pack allocation failed in legacy infallible caller");
}

/// Encode every input word as little-endian bytes into caller-owned scratch
/// with an explicit allocator failure contract.
#[inline]
pub(crate) fn try_pack_u32_into(
    values: &[u32],
    bytes: &mut Vec<u8>,
    field: &'static str,
) -> Result<(), String> {
    let byte_len = byte_len_for_u32_words(values.len(), field)?;
    let stable_len = bytes.len() == byte_len;
    crate::staging_reserve::reserve_vec(bytes, byte_len, field)?;
    #[cfg(target_endian = "little")]
    {
        let packed = bytemuck::cast_slice(values);
        if packed.len() != byte_len {
            // Cast width mismatch is a cold malformed-input path.
            return Err(pack_u32_cast_mismatch(
                field,
                values.len(),
                packed.len(),
                byte_len,
            ));
        }
        if stable_len {
            bytes.copy_from_slice(packed);
        } else {
            bytes.clear();
            bytes.extend_from_slice(packed);
        }
    }
    #[cfg(target_endian = "big")]
    {
        if stable_len {
            for (slot, value) in bytes.chunks_exact_mut(4).zip(values) {
                slot.copy_from_slice(&value.to_le_bytes());
            }
        } else {
            bytes.clear();
            for value in values {
                bytes.extend_from_slice(&value.to_le_bytes());
            }
        }
    }
    Ok(())
}

/// Encode one repeated input word as little-endian bytes into caller-owned scratch.
#[inline]
pub(crate) fn pack_repeated_u32_into(
    value: u32,
    words: usize,
    bytes: &mut Vec<u8>,
) -> Result<(), String> {
    let byte_len = byte_len_for_u32_words(words, "pack_repeated_u32_into")?;
    if value == 0 {
        return try_write_zero_bytes(bytes, byte_len, "repeated zero u32 byte pack");
    }
    let stable_len = bytes.len() == byte_len;
    crate::staging_reserve::reserve_vec(bytes, byte_len, "repeated u32 byte pack")?;
    let word = value.to_le_bytes();
    if stable_len {
        for slot in bytes.chunks_exact_mut(4) {
            slot.copy_from_slice(&word);
        }
    } else {
        bytes.clear();
        for _ in 0..words {
            bytes.extend_from_slice(&word);
        }
    }
    Ok(())
}

/// Encode exactly `slots` little-endian words into caller-owned scratch.
///
/// # Errors
///
/// Returns an actionable ABI error when the caller supplies too few or too
/// many words.
#[inline]
pub(crate) fn pack_exact_u32_slots_into(
    values: &[u32],
    slots: usize,
    name: &str,
    bytes: &mut Vec<u8>,
) -> Result<(), String> {
    DispatchU32Schema::input(name, slots).require_input_words(values.len())?;
    try_pack_u32_into(values, bytes, "exact u32 slot pack")?;
    Ok(())
}

/// Extend a semantically exact byte buffer to a backend dispatch minimum.
///
/// This is not semantic padding: callers have already passed
/// [`pack_exact_u32_slots`]. The extra zero bytes exist only for backends that
/// require at least one storage word even when the semantic workload is empty.
#[inline]
pub(crate) fn pad_dispatch_min_words(
    bytes: &mut Vec<u8>,
    dispatch_slots: usize,
) -> Result<(), String> {
    try_pad_zero_bytes_to_len(
        bytes,
        byte_len_for_u32_words(dispatch_slots, "pad_dispatch_min_words")?,
        "dispatch minimum word padding",
    )?;
    Ok(())
}

/// Write an exact zero-filled byte buffer into caller-owned scratch with an
/// explicit allocator failure contract.
#[inline]
pub(crate) fn try_write_zero_bytes(
    bytes: &mut Vec<u8>,
    byte_len: usize,
    field: &'static str,
) -> Result<(), String> {
    if bytes.len() == byte_len {
        bytes.fill(0);
        return Ok(());
    }
    crate::staging_reserve::reserve_vec(bytes, byte_len, field)?;
    bytes.clear();
    extend_zero_bytes(bytes, byte_len);
    Ok(())
}

/// Write an exact zero-filled word buffer into caller-owned scratch.
#[inline]
#[cfg(any(test, feature = "legacy-infallible"))]
pub(crate) fn write_zero_words(words: &mut Vec<u32>, word_len: usize) {
    try_write_zero_words(words, word_len, "zero word staging")
        .expect("zero word staging allocation failed in legacy infallible caller");
}

/// Write an exact zero-filled word buffer into caller-owned scratch with an
/// explicit allocator failure contract.
#[inline]
pub(crate) fn try_write_zero_words(
    words: &mut Vec<u32>,
    word_len: usize,
    field: &'static str,
) -> Result<(), String> {
    if words.len() == word_len {
        words.fill(0);
        return Ok(());
    }
    crate::staging_reserve::reserve_vec(words, word_len, field)?;
    words.clear();
    extend_zero_words(words, word_len);
    Ok(())
}

/// Append zero bytes up to an exact byte length without disturbing the prefix,
/// with an explicit allocator failure contract.
#[inline]
pub(crate) fn try_pad_zero_bytes_to_len(
    bytes: &mut Vec<u8>,
    byte_len: usize,
    field: &'static str,
) -> Result<(), String> {
    crate::staging_reserve::reserve_vec(bytes, byte_len, field)?;
    if bytes.len() < byte_len {
        extend_zero_bytes(bytes, byte_len - bytes.len());
    } else {
        bytes.truncate(byte_len);
    }
    Ok(())
}

#[inline]
fn extend_zero_bytes(bytes: &mut Vec<u8>, additional: usize) {
    bytes.extend(std::iter::repeat_n(0, additional));
}

#[inline]
fn extend_zero_words(words: &mut Vec<u32>, additional: usize) {
    words.extend(std::iter::repeat_n(0, additional));
}

#[cold]
fn pack_u32_cast_mismatch(
    field: &'static str,
    words: usize,
    got_bytes: usize,
    expected_bytes: usize,
) -> String {
    format!(
        "{field} cast {words} u32 words into {got_bytes} bytes, expected {expected_bytes}. Fix: split the Weir GPU dispatch buffer or repair host slice width accounting."
    )
}

#[cfg(test)]
mod packing_tests {
    use super::{
        pack_exact_u32_slots_into, pack_repeated_u32_into, pack_u32, pack_u32_into,
        try_pad_zero_bytes_to_len, try_write_zero_bytes, try_write_zero_words,
    };

    #[test]
    fn pack_u32_preserves_little_endian_wire_order() {
        assert_eq!(
            pack_u32(&[0x0102_0304, 0xa0b0_c0d0]),
            vec![0x04, 0x03, 0x02, 0x01, 0xd0, 0xc0, 0xb0, 0xa0]
        );
    }

    #[test]
    fn pack_u32_into_reuses_capacity_and_preserves_wire_order() {
        let mut bytes = Vec::with_capacity(64);
        let original_capacity = bytes.capacity();
        pack_u32_into(&[0x1122_3344], &mut bytes);

        assert_eq!(bytes, vec![0x44, 0x33, 0x22, 0x11]);
        assert_eq!(
            bytes.capacity(),
            original_capacity,
            "packing into caller-owned scratch must not shrink reusable capacity"
        );
    }

    #[test]
    fn pack_repeated_u32_into_preserves_wire_order_without_word_staging() {
        let mut bytes = Vec::with_capacity(64);
        let original_capacity = bytes.capacity();
        pack_repeated_u32_into(0x1122_3344, 3, &mut bytes)
            .expect("small repeated u32 pack must fit");

        assert_eq!(
            bytes,
            vec![0x44, 0x33, 0x22, 0x11, 0x44, 0x33, 0x22, 0x11, 0x44, 0x33, 0x22, 0x11]
        );
        assert_eq!(
            bytes.capacity(),
            original_capacity,
            "repeated packing into caller-owned scratch must not allocate a staging word buffer"
        );
    }

    #[test]
    fn pack_exact_u32_slots_into_reuses_capacity_and_rejects_wrong_width() {
        let mut bytes = Vec::with_capacity(64);
        let original_capacity = bytes.capacity();

        pack_exact_u32_slots_into(&[0x1122_3344], 1, "test exact", &mut bytes)
            .expect("matching semantic width should pack");

        assert_eq!(bytes, vec![0x44, 0x33, 0x22, 0x11]);
        assert_eq!(bytes.capacity(), original_capacity);
        let err = pack_exact_u32_slots_into(&[1, 2], 1, "test exact", &mut bytes)
            .expect_err("mismatched semantic width must be rejected");
        assert!(
            err.contains("expected exactly 1"),
            "unexpected diagnostic: {err}"
        );
    }

    #[test]
    fn zero_staging_extends_after_reserve_without_resize_growth() {
        let mut bytes = Vec::with_capacity(16);
        try_write_zero_bytes(&mut bytes, 8, "test zero bytes").expect("zero bytes should stage");
        assert_eq!(bytes, vec![0; 8]);
        bytes[0] = 7;
        try_pad_zero_bytes_to_len(&mut bytes, 12, "test zero pad").expect("zero pad should stage");
        assert_eq!(bytes[0], 7);
        assert_eq!(&bytes[8..], &[0, 0, 0, 0]);

        let mut words = Vec::with_capacity(8);
        try_write_zero_words(&mut words, 3, "test zero words").expect("zero words should stage");
        assert_eq!(words, vec![0; 3]);
    }

    #[test]
    fn dispatch_decode_source_has_no_resize_based_zero_staging() {
        let source = include_str!("pack.rs");
        let production = source
            .split("#[cfg(test)]")
            .next()
            .expect("dispatch decode production source must precede tests");
        assert!(
            !production.contains(".resize(")
                && production.contains("fn extend_zero_bytes(")
                && production.contains("fn extend_zero_words("),
            "Fix: dispatch decode zero staging must extend after fallible reserve instead of resize-driven growth."
        );
        assert!(
            !production.contains("fn byte_len_for_existing_words")
                && !production.contains("debug_assert_eq!(bytes.len(), byte_len)")
                && !production.contains("debug_assert_eq!(packed.len(), byte_len)"),
            "Fix: dispatch decode byte sizing and cast-width checks must run in release builds."
        );
    }
}

#[cfg(test)]
mod panic_tests {
    use super::{pack_u32_into, write_zero_words};

    #[test]
    #[should_panic(expected = "capacity overflow")]
    fn pack_u32_into_panics_on_huge_input() {
        // vec![] itself panics before pack_u32_into is reached; the
        // important invariant is that huge inputs do not silently succeed.
        let huge = vec![0u32; usize::MAX / 8 + 1];
        let mut out = Vec::new();
        pack_u32_into(&huge, &mut out);
    }

    #[test]
    #[should_panic(expected = "zero word staging allocation failed in legacy infallible caller")]
    fn write_zero_words_panics_on_huge_input() {
        let mut out = Vec::new();
        write_zero_words(&mut out, usize::MAX / 8 + 1);
    }
}