weirflow 0.1.0

//! Compact reaching-definition summary query.
//!
//! This is the GPU-resident form optimizers and rule conditions usually need:
//! not the full reaching-def bitset copied back to the host, but a compact
//! exact summary of the intersection for branch decisions, cache keys, and
//! alias gates.

use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
use vyre_primitives::graph::csr_forward_traverse::bitset_words;

pub(crate) const OP_ID: &str = "weir::reaching_def_summary";
/// Number of independent reduction bins used by the default summary program.
pub const DEFAULT_PARTIAL_BINS: u32 = 1024;
const DEFAULT_PARTIAL_BINS_USIZE: usize = 1024;
/// Hard cap on reduction bins to keep summary outputs bounded.
pub const MAX_PARTIAL_BINS: u32 = DEFAULT_PARTIAL_BINS;
#[cfg(test)]
const MAX_PARTIAL_BINS_USIZE: usize = DEFAULT_PARTIAL_BINS_USIZE;
/// Number of output words produced by the default summary program.
pub const SUMMARY_WORDS: usize = DEFAULT_PARTIAL_BINS_USIZE * 3;
const SUMMARY_WORDS_U32: u32 = DEFAULT_PARTIAL_BINS * 3;

/// Build a compact reaching-def summary Program.
///
/// Output layout is `partial_bins` triples:
/// - `summary[bin * 3]`: partial set-bit count
/// - `summary[bin * 3 + 1]`: `1` iff this bin saw a non-empty word
/// - `summary[bin * 3 + 2]`: commutative checksum for this bin
pub fn reaching_def_summary(
    node_count: u32,
    gen_kill_in: &str,
    use_set: &str,
    summary: &str,
) -> Result<Program, String> {
    reaching_def_summary_sharded(
        node_count,
        gen_kill_in,
        use_set,
        summary,
        DEFAULT_PARTIAL_BINS,
    )
}

/// Build a compact reaching-def summary Program with an explicit shard count.
///
/// Non-power-of-two shard counts are rounded up to the next power of two so
/// the bin mask stays valid without exposing a production panic.
pub fn reaching_def_summary_sharded(
    node_count: u32,
    gen_kill_in: &str,
    use_set: &str,
    summary: &str,
    partial_bins: u32,
) -> Result<Program, String> {
    let partial_bins = normalize_partial_bins(partial_bins);
    let summary_words = partial_bins.checked_mul(3).ok_or_else(|| {
        format!(
            "weir::reaching_def_summary partial_bins={partial_bins} overflows summary word count. Fix: lower the shard count before GPU dispatch."
        )
    })?;
    let summary_bytes = usize::try_from(summary_words)
        .ok()
        .and_then(|words| words.checked_mul(4))
        .ok_or_else(|| {
            format!(
                "weir::reaching_def_summary summary_words={summary_words} overflows output byte range. Fix: lower the shard count before GPU dispatch."
            )
        })?;
    Ok(reaching_def_summary_program_with_layout(
        node_count,
        gen_kill_in,
        use_set,
        summary,
        partial_bins,
        summary_words,
        summary_bytes,
    ))
}

fn reaching_def_summary_program_with_layout(
    node_count: u32,
    gen_kill_in: &str,
    use_set: &str,
    summary: &str,
    partial_bins: u32,
    summary_words: u32,
    summary_bytes: usize,
) -> Program {
    let words = bitset_words(node_count);
    let tid = Expr::gid_x();
    Program::wrapped(
        vec![
            BufferDecl::storage(gen_kill_in, 0, BufferAccess::ReadOnly, DataType::U32)
                .with_count(words),
            BufferDecl::storage(use_set, 1, BufferAccess::ReadOnly, DataType::U32)
                .with_count(words),
            BufferDecl::output(summary, 2, DataType::U32)
                .with_count(summary_words)
                .with_output_byte_range(0..summary_bytes),
        ],
        [256, 1, 1],
        vec![vyre_harness::region::wrap_anonymous(
            OP_ID,
            vec![Node::if_then(
                Expr::lt(tid.clone(), Expr::u32(words)),
                vec![
                    Node::let_bind(
                        "summary_bin",
                        Expr::bitand(tid.clone(), Expr::u32(partial_bins - 1)),
                    ),
                    Node::let_bind(
                        "summary_base",
                        Expr::mul(Expr::var("summary_bin"), Expr::u32(3)),
                    ),
                    Node::let_bind(
                        "word",
                        Expr::bitand(
                            Expr::load(gen_kill_in, tid.clone()),
                            Expr::load(use_set, tid.clone()),
                        ),
                    ),
                    Node::let_bind(
                        "_summary_count_prev",
                        Expr::atomic_add(
                            summary,
                            Expr::var("summary_base"),
                            Expr::popcount(Expr::var("word")),
                        ),
                    ),
                    Node::if_then(
                        Expr::ne(Expr::var("word"), Expr::u32(0)),
                        vec![
                            Node::let_bind(
                                "_summary_any_prev",
                                Expr::atomic_or(
                                    summary,
                                    Expr::add(Expr::var("summary_base"), Expr::u32(1)),
                                    Expr::u32(1),
                                ),
                            ),
                            Node::let_bind(
                                "mixed",
                                Expr::bitxor(
                                    Expr::var("word"),
                                    Expr::mul(
                                        Expr::add(tid.clone(), Expr::u32(1)),
                                        Expr::u32(0x9E37_79B9),
                                    ),
                                ),
                            ),
                            Node::let_bind(
                                "_summary_hash_prev",
                                Expr::atomic_xor(
                                    summary,
                                    Expr::add(Expr::var("summary_base"), Expr::u32(2)),
                                    Expr::var("mixed"),
                                ),
                            ),
                        ],
                    ),
                ],
            )],
        )],
    )
}

/// reference oracle for [`reaching_def_summary`].
#[must_use]
#[cfg(any(test, feature = "cpu-parity"))]
#[deprecated(
    note = "reference oracle only; production code must dispatch the Weir Program on a concrete GPU backend or use weir::oracle for parity evidence"
)]
pub(crate) fn cpu_ref(gen_kill_in: &[u32], use_set: &[u32]) -> Vec<u32> {
    reference_sharded(gen_kill_in, use_set, DEFAULT_PARTIAL_BINS)
}

/// reference oracle for [`reaching_def_summary_sharded`].
#[must_use]
#[cfg(any(test, feature = "cpu-parity"))]
#[deprecated(
    note = "reference oracle only; production code must dispatch the Weir Program on a concrete GPU backend or use weir::oracle for parity evidence"
)]
pub(crate) fn cpu_ref_sharded(gen_kill_in: &[u32], use_set: &[u32], partial_bins: u32) -> Vec<u32> {
    reference_sharded(gen_kill_in, use_set, partial_bins)
}

#[cfg(any(test, feature = "cpu-parity"))]
fn reference_sharded(gen_kill_in: &[u32], use_set: &[u32], partial_bins: u32) -> Vec<u32> {
    try_reference_sharded(gen_kill_in, use_set, partial_bins)
        .expect("reaching-def summary reference oracle must fit normalized fixture dimensions")
}

#[cfg(any(test, feature = "cpu-parity"))]
fn try_reference_sharded(
    gen_kill_in: &[u32],
    use_set: &[u32],
    partial_bins: u32,
) -> Result<Vec<u32>, String> {
    let partial_bins = normalize_partial_bins(partial_bins);
    let partial_bins_usize = crate::dispatch_decode::u32_to_usize(
        partial_bins,
        "reaching-def summary partial bin count",
    )?;
    let partial_slots = partial_bins_usize.checked_mul(3).ok_or_else(|| {
        "reaching-def summary partial slot count overflows usize after normalization".to_string()
    })?;
    let mut partials = vec![0u32; partial_slots];
    for (index, (&left, &right)) in gen_kill_in.iter().zip(use_set.iter()).enumerate() {
        let word = left & right;
        let bin = index & (partial_bins_usize - 1);
        let base = bin * 3;
        partials[base] = partials[base].checked_add(word.count_ones()).ok_or_else(|| {
            "reaching-def summary partial count overflows u32. Fix: shard the summary oracle input."
                .to_string()
        })?;
        if word != 0 {
            let index_u32 = u32::try_from(index).map_err(|error| {
                format!(
                    "reaching-def summary word index {index} does not fit u32: {error}. Fix: shard the summary oracle input."
                )
            })?;
            partials[base + 1] = 1;
            partials[base + 2] ^= word ^ (index_u32.wrapping_add(1).wrapping_mul(0x9E37_79B9));
        }
    }
    Ok(partials)
}

#[inline]
fn normalize_partial_bins(partial_bins: u32) -> u32 {
    let requested = partial_bins.max(1);
    if requested >= MAX_PARTIAL_BINS {
        MAX_PARTIAL_BINS
    } else {
        requested
            .checked_next_power_of_two()
            .unwrap_or(MAX_PARTIAL_BINS)
    }
}

/// Fold partial summary triples `(count, any, checksum)` into `[count, any, checksum]`.
///
/// - `count` is the sum of every triple's first lane (saturating to `u32::MAX`
///   on overflow via the fallback path).
/// - `any` is `1` if any triple's middle lane is non-zero, else `0`: a
///   reachability boolean, not a bit-OR.
/// - `checksum` is the XOR of `chunk[2]` for every triple whose `any` lane is
///   non-zero.
///
/// # Example
///
/// ```
/// use weir::reaching_def_summary::fold_summary_partials;
///
/// // Three triples (count, any, checksum):
/// //   (3, 0b101, 0xAB), contributes to any (0b101 != 0) + xor 0xAB
/// //   (2, 0,     0xCD), does not contribute to any/checksum
/// //   (5, 0b110, 0xEF), contributes to any + xor 0xEF
/// let partials = [3, 0b101, 0xAB, 2, 0, 0xCD, 5, 0b110, 0xEF];
/// let [count, any, checksum] = fold_summary_partials(&partials);
/// assert_eq!(count, 10);            // 3 + 2 + 5
/// assert_eq!(any, 1);               // some triple's any-lane was non-zero
/// assert_eq!(checksum, 0xAB ^ 0xEF);
/// ```
#[inline]
#[must_use]
pub fn fold_summary_partials(partials: &[u32]) -> [u32; 3] {
    match try_fold_summary_partials(partials) {
        Ok(folded) => folded,
        Err(_error) => [u32::MAX, folded_any(partials), folded_checksum(partials)],
    }
}

/// Fold partial summary triples into `[count, any, checksum]` with exact count
/// overflow reporting.
#[inline]
pub fn try_fold_summary_partials(partials: &[u32]) -> Result<[u32; 3], String> {
    if partials.len() % 3 != 0 {
        return Err(format!(
            "reaching-def summary partial length {} is not divisible by 3. Fix: reject truncated summary output before folding.",
            partials.len()
        ));
    }
    let mut count = 0u32;
    let mut any = 0u32;
    let mut checksum = 0u32;
    for chunk in partials.chunks_exact(3) {
        count = count.checked_add(chunk[0]).ok_or_else(|| {
            "reaching-def summary folded count overflowed u32. Fix: shard the summary domain before folding."
                .to_string()
        })?;
        if chunk[1] != 0 {
            any = 1;
            checksum ^= chunk[2];
        }
    }
    Ok([count, any, checksum])
}

#[inline]
fn folded_any(partials: &[u32]) -> u32 {
    for chunk in partials.chunks_exact(3) {
        if chunk[1] != 0 {
            return 1;
        }
    }
    0
}

#[inline]
fn folded_checksum(partials: &[u32]) -> u32 {
    let mut checksum = 0u32;
    for chunk in partials.chunks_exact(3) {
        if chunk[1] != 0 {
            checksum ^= chunk[2];
        }
    }
    checksum
}

#[cfg(test)]
mod source_contract_tests {
    #[test]
    fn reaching_def_summary_fold_has_checked_release_api_without_panic() {
        let source = include_str!("reaching_def_summary.rs");
        let fold_api = source
            .split("pub fn try_fold_summary_partials(")
            .nth(1)
            .expect("reaching-def summary fold API must be present")
            .split("fn folded_any(")
            .next()
            .expect("reaching-def summary fold API must precede fallback helpers");
        assert!(
            !fold_api.contains(concat!("panic", "!("))
                && !fold_api.contains(".unwrap_or_else("),
            "Fix: reaching-def summary folding must expose checked overflow handling and avoid production panics."
        );
    }
}

inventory::submit! {
    vyre_harness::OpEntry::new(
        OP_ID,
        || reaching_def_summary_program_with_layout(
            64,
            "reaching",
            "uses",
            "summary",
            DEFAULT_PARTIAL_BINS,
            SUMMARY_WORDS_U32,
            SUMMARY_WORDS * 4,
        ),
        Some(|| {
            let u32s = crate::dispatch_decode::pack_u32;
            vec![vec![
                u32s(&[0b1111, 0]),
                u32s(&[0b0010, 1]),
            ]]
        }),
        Some(|| {
            let mut expected = vec![0u32; SUMMARY_WORDS];
            expected[0] = 1;
            expected[1] = 1;
            expected[2] = 0b0010 ^ 0x9E37_79B9;
            let bytes = crate::dispatch_decode::pack_u32(&expected);
            vec![vec![bytes]]
        }),
    )
}

#[cfg(test)]
#[allow(deprecated)]
mod tests {
    use super::*;

    #[test]
    fn summary_counts_any_and_checksum() {
        let summary = cpu_ref(&[0b1111, 0x10], &[0b0011, 0]);
        let folded = fold_summary_partials(&summary);
        assert_eq!(folded[0], 2);
        assert_eq!(folded[1], 1);
        assert_ne!(folded[2], 0);
    }

    #[test]
    fn empty_intersection_is_zero_summary() {
        assert_eq!(
            fold_summary_partials(&cpu_ref(&[0b0101], &[0b1010])),
            [0, 0, 0]
        );
    }

    #[test]
    fn non_power_of_two_shards_are_normalized() {
        let normalized = cpu_ref_sharded(&[0b1, 0b10, 0b100], &[0b1, 0b10, 0], 3);
        assert_eq!(normalized.len(), 12);
        assert_eq!(fold_summary_partials(&normalized)[0], 2);
    }

    #[test]
    fn oversized_shards_are_explicitly_capped() {
        let normalized = cpu_ref_sharded(&[0b1], &[0b1], u32::MAX);
        assert_eq!(normalized.len(), MAX_PARTIAL_BINS_USIZE * 3);
        assert_eq!(fold_summary_partials(&normalized)[0], 1);
    }

    #[test]
    fn checked_fold_reports_count_overflow_and_legacy_fold_pins() {
        let partials = [u32::MAX, 1, 0x1234, 1, 1, 0x5678];
        let err =
            try_fold_summary_partials(&partials).expect_err("checked fold must report overflow");
        assert!(err.contains("overflowed u32"));
        assert_eq!(
            fold_summary_partials(&partials),
            [u32::MAX, 1, 0x1234 ^ 0x5678]
        );
    }

    #[test]
    fn checked_fold_rejects_truncated_partial_triples() {
        let err = try_fold_summary_partials(&[7, 1, 0x1234, 99])
            .expect_err("checked fold must reject trailing malformed summary words");
        assert!(
            err.contains("not divisible by 3"),
            "malformed summary length error must identify the triple contract, got {err}"
        );
        assert_eq!(
            fold_summary_partials(&[7, 1, 0x1234, 99]),
            [u32::MAX, 1, 0x1234],
            "legacy fold must remain a loud sentinel while preserving valid triple any/checksum"
        );
    }
}