weirflow 0.1.0

//! Reusable host scratch for GPU-driven fixed-point dataflow loops.
//!
//! This module owns the per-session buffers that are safe to reuse across
//! reaching, liveness, points-to, and future fixed-point analyses. The graph
//! and transfer semantics stay in each analysis module; this file only owns
//! scratch lifetime and reuse.

mod buffers;
mod telemetry;

#[cfg(test)]
mod tests;

pub(crate) use buffers::FixedPointForwardChangedProgramCache;
pub use buffers::{count_changed_domain_bits, count_set_domain_bits, FixedPointScratch};
#[cfg(feature = "gpu-telemetry")]
pub use buffers::{count_changed_domain_bits_gpu, count_set_domain_bits_gpu};
pub use telemetry::{FrontierDensityTelemetry, FrontierExecutionMode};

/// Generic retention pool for reusable host-side scratch vectors.
///
/// Clears logical contents while preserving allocation capacity across
/// dispatches. Derefs to [`Vec<T>`] so existing code can push, extend, and
/// index without extra boilerplate.
#[derive(Clone, Debug, PartialEq)]
pub struct ScratchPool<T> {
    vec: Vec<T>,
}

impl<T> Default for ScratchPool<T> {
    fn default() -> Self {
        Self { vec: Vec::new() }
    }
}

impl<T> ScratchPool<T> {
    /// Create an empty scratch pool.
    #[inline]
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Create a scratch pool with pre-reserved capacity.
    #[inline]
    #[must_use]
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            vec: Vec::with_capacity(capacity),
        }
    }

    /// Drop logical contents while preserving allocation capacity.
    #[inline]
    pub fn clear(&mut self) {
        self.vec.clear();
    }

    /// Current retained capacity.
    #[inline]
    #[must_use]
    pub fn capacity(&self) -> usize {
        self.vec.capacity()
    }
}

impl<T> std::ops::Deref for ScratchPool<T> {
    type Target = Vec<T>;

    fn deref(&self) -> &Self::Target {
        &self.vec
    }
}

impl<T> std::ops::DerefMut for ScratchPool<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.vec
    }
}

/// Build a `vyre_primitives::bitset::equal` program for `words`.
#[inline]
pub(crate) fn bitset_equal_program(words: usize) -> vyre::ir::Program {
    vyre_primitives::bitset::equal::bitset_equal("lhs", "rhs", "out", words as u32)
}

/// GPU convergence check using a generic dispatch closure.
///
/// Dispatches `bitset_equal(next_bytes, current_bytes, out_scalar)` and
/// returns `true` iff `out_scalar[0] == 1`.
///
/// `next_bytes` and `current_bytes` must each be exactly `words * 4`
/// little-endian bytes.
#[allow(dead_code)]
pub(crate) fn converged_via_gpu_into<F>(
    next_bytes: &[u8],
    current_bytes: &[u8],
    words: usize,
    dispatch: F,
) -> Result<bool, String>
where
    F: FnOnce(
        &vyre::ir::Program,
        &[&[u8]],
        Option<[u32; 3]>,
        &mut Vec<Vec<u8>>,
    ) -> Result<(), String>,
{
    if next_bytes.len() != words * 4 || current_bytes.len() != words * 4 {
        return Ok(false);
    }
    let program = bitset_equal_program(words);
    let scalar_zero = 0u32.to_le_bytes();
    let mut outputs = Vec::new();
    dispatch(
        &program,
        &[next_bytes, current_bytes, &scalar_zero],
        Some([1, 1, 1]),
        &mut outputs,
    )?;
    let scalar = crate::dispatch_decode::unpack_exact_u32_scalar(
        crate::dispatch_decode::only_output(&outputs, "bitset_equal", "scalar")?,
        "bitset_equal",
    )?;
    Ok(scalar == 1)
}

/// GPU convergence check using a [`vyre::VyreBackend`] directly.
///
/// Dispatches `bitset_equal` through `backend.dispatch_borrowed` and returns
/// `true` iff the scalar output is `1`.
pub(crate) fn converged_via_gpu_backend(
    next_bytes: &[u8],
    current_bytes: &[u8],
    words: usize,
    backend: &dyn vyre::VyreBackend,
) -> Result<bool, String> {
    if next_bytes.len() != words * 4 || current_bytes.len() != words * 4 {
        return Ok(false);
    }
    let program = bitset_equal_program(words);
    let scalar_zero = 0u32.to_le_bytes();
    let outputs = backend
        .dispatch_borrowed(
            &program,
            &[next_bytes, current_bytes, &scalar_zero],
            &vyre::DispatchConfig::default(),
        )
        .map_err(|error| error.to_string())?;
    let scalar = crate::dispatch_decode::unpack_exact_u32_scalar(
        crate::dispatch_decode::only_output(&outputs, "bitset_equal", "scalar")?,
        "bitset_equal",
    )?;
    Ok(scalar == 1)
}

#[inline]
pub(crate) fn copy_if_converged(
    next: &[u32],
    current: &[u32],
    result: &mut Vec<u32>,
) -> Result<bool, String> {
    if next != current {
        return Ok(false);
    }
    copy_words_into(next, result)?;
    Ok(true)
}

#[inline]
pub(crate) fn copy_words_into(words: &[u32], result: &mut Vec<u32>) -> Result<(), String> {
    if result.len() == words.len() {
        result.copy_from_slice(words);
        return Ok(());
    }
    // Slow allocation path is cold after scratch warms up.
    copy_words_into_slow_path(words, result)
}

#[cold]
fn copy_words_into_slow_path(words: &[u32], result: &mut Vec<u32>) -> Result<(), String> {
    crate::staging_reserve::reserve_vec(result, words.len(), "fixed-point result word copy")?;
    result.clear();
    result.extend_from_slice(words);
    Ok(())
}