weirflow 0.1.0

//! Fallible host-staging reservation helpers for Weir GPU paths.

/// Reserve total capacity for caller-owned host staging without silently
/// panicking on allocator failure.
pub(crate) fn reserve_vec<T>(
    values: &mut Vec<T>,
    capacity: usize,
    field: &'static str,
) -> Result<(), String> {
    if values.capacity() >= capacity {
        return Ok(());
    }
    values
        .try_reserve_exact(capacity - values.capacity())
        .map_err(|error| {
            format!(
                "Weir GPU host staging could not reserve {capacity} {field} slot(s): {error}. Fix: shard the dataflow graph, split the resident batch, or lower fixed-point scratch fan-out before dispatch."
            )
        })
}

/// Build a vector with a fallibly reserved capacity and a uniform Weir GPU
/// staging error contract.
pub(crate) fn reserved_vec<T>(capacity: usize, field: &'static str) -> Result<Vec<T>, String> {
    let mut values = Vec::new();
    reserve_vec(&mut values, capacity, field)?;
    Ok(values)
}

/// Build a fixed-point result vector with the canonical bitset word capacity
/// and staging error contract.
pub(crate) fn reserved_fixed_point_result(
    stage: &str,
    node_count: u32,
) -> Result<Vec<u32>, String> {
    reserved_vec(
        crate::dispatch_decode::bitset_word_capacity(stage, node_count)?,
        "fixed-point result word",
    )
}

/// Ensure caller-owned vector slots can address at least `slot_count` slots.
pub(crate) fn ensure_vec_slots_at_least<T>(
    slots: &mut Vec<Vec<T>>,
    slot_count: usize,
    field: &'static str,
) -> Result<(), String> {
    reserve_vec(slots, slot_count, field)?;
    if slots.len() < slot_count {
        slots.resize_with(slot_count, Vec::new);
    }
    Ok(())
}

/// Resize caller-owned vector slots to an exact count without losing existing
/// per-slot allocation.
pub(crate) fn resize_vec_slots<T>(
    slots: &mut Vec<Vec<T>>,
    slot_count: usize,
    field: &'static str,
) -> Result<(), String> {
    ensure_vec_slots_at_least(slots, slot_count, field)?;
    if slots.len() > slot_count {
        slots.truncate(slot_count);
    }
    Ok(())
}

/// Resize caller-owned result rows without losing existing per-row allocation.
pub(crate) fn resize_result_rows<T>(
    rows: &mut Vec<Vec<T>>,
    row_count: usize,
    field: &'static str,
) -> Result<(), String> {
    resize_vec_slots(rows, row_count, field)
}

/// Build exact-count result rows with fallible outer-row reservation.
pub(crate) fn reserved_result_rows<T>(
    row_count: usize,
    field: &'static str,
) -> Result<Vec<Vec<T>>, String> {
    let mut rows = reserved_vec(row_count, field)?;
    resize_result_rows(&mut rows, row_count, field)?;
    Ok(rows)
}

/// Clear all retained vector-slot payloads while preserving their allocations.
pub(crate) fn clear_vec_slots<T>(slots: &mut [Vec<T>]) {
    for slot in slots {
        slot.clear();
    }
}

/// Clear all retained row payloads while preserving their allocations.
pub(crate) fn clear_result_rows<T>(rows: &mut [Vec<T>]) {
    clear_vec_slots(rows);
}

#[cfg(test)]
mod tests {
    use super::{
        clear_result_rows, ensure_vec_slots_at_least, reserved_fixed_point_result,
        reserved_result_rows, resize_result_rows, resize_vec_slots,
    };

    #[test]
    fn resize_result_rows_preserves_existing_row_capacity() {
        let mut rows = vec![Vec::with_capacity(8)];
        rows[0].extend_from_slice(&[1u32, 2, 3]);

        resize_result_rows(&mut rows, 3, "test result row").unwrap();

        assert_eq!(rows.len(), 3);
        assert!(rows[0].capacity() >= 8);
        assert_eq!(rows[0], [1, 2, 3]);
    }

    #[test]
    fn clear_result_rows_preserves_slots_and_capacity() {
        let mut rows = vec![Vec::with_capacity(4), Vec::with_capacity(2)];
        rows[0].extend_from_slice(&[1u32, 2]);
        rows[1].push(3);

        clear_result_rows(&mut rows);

        assert_eq!(rows.len(), 2);
        assert!(rows[0].capacity() >= 4);
        assert!(rows[1].capacity() >= 2);
        assert!(rows.iter().all(Vec::is_empty));
    }

    #[test]
    fn ensure_vec_slots_at_least_does_not_truncate_extra_slots() {
        let mut slots = vec![vec![1u8], vec![2], vec![3]];

        ensure_vec_slots_at_least(&mut slots, 2, "test byte slot").unwrap();

        assert_eq!(slots.len(), 3);
        assert_eq!(slots[2], [3]);
    }

    #[test]
    fn resize_vec_slots_truncates_to_exact_slot_count() {
        let mut slots = vec![vec![1u8], vec![2], vec![3]];

        resize_vec_slots(&mut slots, 2, "test exact slot").unwrap();

        assert_eq!(slots.len(), 2);
        assert_eq!(slots[0], [1]);
        assert_eq!(slots[1], [2]);
    }

    #[test]
    fn reserved_result_rows_builds_exact_empty_rows() {
        let rows: Vec<Vec<u32>> = reserved_result_rows(3, "test reserved result row").unwrap();

        assert_eq!(rows.len(), 3);
        assert!(rows.iter().all(Vec::is_empty));
    }

    #[test]
    fn reserved_fixed_point_result_uses_bitset_word_capacity() {
        let result = reserved_fixed_point_result("test fixed-point result", 65).unwrap();

        assert!(result.is_empty());
        assert!(result.capacity() >= 3);
    }

    #[test]
    fn reserved_fixed_point_result_covers_word_boundaries() {
        for (node_count, expected_words) in [(0, 0), (1, 1), (32, 1), (33, 2), (64, 2), (65, 3)] {
            let result = reserved_fixed_point_result("test fixed-point boundary", node_count)
                .expect("Fix: fixed-point result reservation must accept canonical bitset widths");

            assert!(result.is_empty());
            assert!(result.capacity() >= expected_words);
        }
    }
}