weirflow 0.1.0

#[cfg(any(test, feature = "cpu-parity"))]
use std::sync::Arc;

#[cfg(any(test, feature = "cpu-parity"))]
use vyre::ir::model::expr::Ident;
use vyre::ir::Program;
#[cfg(any(test, feature = "cpu-parity"))]
use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Node};
use vyre_primitives::bitset::and::bitset_and;

/// Build a dominator query Program. Inputs as documented above.
#[must_use]
pub fn dominates(node_count: u32, dom_set: &str, target_set: &str, out: &str) -> Program {
    let words = crate::graph_layout::LinearDomain::new(node_count).bitset_words();
    vyre_harness::region::tag_program(super::OP_ID, bitset_and(dom_set, target_set, out, words))
}

/// Binding for predecessor CSR offsets in [`dominator_step`].
#[cfg(any(test, feature = "cpu-parity"))]
pub(crate) const BINDING_PRED_OFFSETS: u32 = 0;
/// Binding for predecessor CSR targets in [`dominator_step`].
#[cfg(any(test, feature = "cpu-parity"))]
pub(crate) const BINDING_PRED_TARGETS: u32 = 1;
/// Binding for the previous dominator matrix in [`dominator_step`].
#[cfg(any(test, feature = "cpu-parity"))]
pub(crate) const BINDING_DOM_IN: u32 = 2;
/// Binding for the next dominator matrix in [`dominator_step`].
#[cfg(any(test, feature = "cpu-parity"))]
pub(crate) const BINDING_DOM_OUT: u32 = 3;

/// Build one GPU-resident dominator fixed-point step.
///
/// The dispatch domain is `node_count * bitset_words(node_count)` lanes. Each
/// lane owns one `(node, word)` cell of the row-major dominator matrix and
/// computes:
///
/// - entry row: `{entry}`;
/// - node with predecessors: `self | intersection(dom[p] for p in preds(node))`;
/// - node with no predecessors: `{entry}`.
#[must_use]
#[cfg(any(test, feature = "cpu-parity"))]
pub(crate) fn dominator_step(
    node_count: u32,
    edge_count: u32,
    entry: u32,
    pred_offsets: &str,
    pred_targets: &str,
    dom_in: &str,
    dom_out: &str,
) -> Result<Program, String> {
    let words = crate::graph_layout::LinearDomain::new(node_count).bitset_words();
    let matrix_words = node_count.checked_mul(words).ok_or_else(|| {
        format!(
            "weir::dominators dominator_step node_count={node_count} overflows matrix word count. Fix: shard the CFG before building the GPU program."
        )
    })?;
    let pred_offset_count = node_count.checked_add(1).ok_or_else(|| {
        format!(
            "weir::dominators dominator_step node_count={node_count} overflows predecessor offset count. Fix: shard the CFG before building the GPU program."
        )
    })?;
    let t = Expr::InvocationId { axis: 0 };
    let entry_word = entry / 32;
    let entry_bit = 1u32 << (entry % 32);
    let tail_bits = node_count % 32;
    let last_word_mask = if tail_bits == 0 {
        u32::MAX
    } else {
        (1u32 << tail_bits) - 1
    };

    let body = vec![
        Node::let_bind("node", Expr::div(t.clone(), Expr::u32(words))),
        Node::let_bind("word", Expr::rem(t.clone(), Expr::u32(words))),
        Node::let_bind(
            "entry_only",
            Expr::select(
                Expr::eq(Expr::var("word"), Expr::u32(entry_word)),
                Expr::u32(entry_bit),
                Expr::u32(0),
            ),
        ),
        Node::let_bind(
            "row_mask",
            Expr::select(
                Expr::eq(Expr::var("word"), Expr::u32(words - 1)),
                Expr::u32(last_word_mask),
                Expr::u32(u32::MAX),
            ),
        ),
        Node::let_bind("pred_start", Expr::load(pred_offsets, Expr::var("node"))),
        Node::let_bind(
            "pred_end",
            Expr::load(pred_offsets, Expr::add(Expr::var("node"), Expr::u32(1))),
        ),
        Node::let_bind("acc", Expr::var("row_mask")),
        Node::loop_for(
            "edge",
            Expr::var("pred_start"),
            Expr::var("pred_end"),
            vec![
                Node::let_bind("pred", Expr::load(pred_targets, Expr::var("edge"))),
                Node::assign(
                    "acc",
                    Expr::bitand(
                        Expr::var("acc"),
                        Expr::load(
                            dom_in,
                            Expr::add(
                                Expr::mul(Expr::var("pred"), Expr::u32(words)),
                                Expr::var("word"),
                            ),
                        ),
                    ),
                ),
            ],
        ),
        Node::let_bind(
            "self_bit",
            Expr::select(
                Expr::eq(
                    Expr::var("word"),
                    Expr::shr(Expr::var("node"), Expr::u32(5)),
                ),
                Expr::shl(Expr::u32(1), Expr::bitand(Expr::var("node"), Expr::u32(31))),
                Expr::u32(0),
            ),
        ),
        Node::let_bind(
            "candidate",
            Expr::select(
                Expr::ne(Expr::var("pred_start"), Expr::var("pred_end")),
                Expr::bitor(Expr::var("acc"), Expr::var("self_bit")),
                Expr::var("entry_only"),
            ),
        ),
        Node::store(
            dom_out,
            t.clone(),
            Expr::select(
                Expr::eq(Expr::var("node"), Expr::u32(entry)),
                Expr::var("entry_only"),
                Expr::var("candidate"),
            ),
        ),
    ];

    Ok(Program::wrapped(
        vec![
            BufferDecl::storage(
                pred_offsets,
                BINDING_PRED_OFFSETS,
                BufferAccess::ReadOnly,
                DataType::U32,
            )
            .with_count(pred_offset_count),
            BufferDecl::storage(
                pred_targets,
                BINDING_PRED_TARGETS,
                BufferAccess::ReadOnly,
                DataType::U32,
            )
            .with_count(edge_count),
            BufferDecl::storage(
                dom_in,
                BINDING_DOM_IN,
                BufferAccess::ReadOnly,
                DataType::U32,
            )
            .with_count(matrix_words),
            BufferDecl::storage(
                dom_out,
                BINDING_DOM_OUT,
                BufferAccess::ReadWrite,
                DataType::U32,
            )
            .with_count(matrix_words),
        ],
        [256, 1, 1],
        vec![Node::Region {
            generator: Ident::from(super::OP_ID),
            source_region: None,
            body: Arc::new(vec![Node::if_then(
                Expr::lt(t, Expr::u32(matrix_words)),
                body,
            )]),
        }],
    ))
}

#[cfg(test)]
mod source_contract_tests {
    #[test]
    fn dominator_program_builder_returns_errors_instead_of_panicking_on_shape_overflow() {
        let source = include_str!("program.rs");
        let production = source
            .split("#[cfg(test)]")
            .next()
            .expect("dominator program production source must precede tests");
        assert!(
            production.contains(") -> Result<Program, String>")
                && !production.contains(concat!("panic", "!("))
                && !production.contains(".unwrap_or_else("),
            "Fix: dominator GPU Program construction must return actionable errors instead of panicking on oversized shapes."
        );
    }
}