vyre-primitives 0.6.5

use std::sync::Arc;

use vyre_foundation::ir::model::expr::{GeneratorRef, Ident};
use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
use vyre_foundation::MemoryOrdering;

use super::layout::{
    CSR_FORWARD_OR_CHANGED_CHANGED_BUFFER, CSR_FORWARD_OR_CHANGED_FRONTIER_BUFFER,
    CSR_FORWARD_OR_CHANGED_PARALLEL_WORKGROUP_SIZE, OP_ID,
};
use crate::graph::edge_scan::csr_edge_expand_nodes;
use crate::graph::program_graph::ProgramGraphShape;

/// Parallel in-place expansion program for production fixed-point drivers.
///
/// Unlike [`csr_forward_or_changed`], this variant gives each source node its
/// own invocation instead of walking the whole CSR from one lane. The pass is
/// monotone: each dispatch may observe only the frontier bits visible at that
/// point in the dispatch, but every newly discovered destination is ORed into
/// the same resident accumulator and sets `changed[0]`. Re-dispatch until the
/// changed flag stays zero to compute the same reachability fixpoint without a
/// full frontier readback per iteration.
#[must_use]
pub fn csr_forward_or_changed_parallel(
    shape: ProgramGraphShape,
    frontier_out: &str,
    changed: &str,
    edge_kind_mask: u32,
) -> Program {
    let words = crate::bitset::bitset_words(shape.node_count);
    let body = csr_forward_or_changed_parallel_body_prefixed(
        shape,
        frontier_out,
        changed,
        edge_kind_mask,
        "",
    );
    let mut buffers = shape.read_only_buffers();
    buffers.push(
        BufferDecl::storage(
            frontier_out,
            CSR_FORWARD_OR_CHANGED_FRONTIER_BUFFER,
            BufferAccess::ReadWrite,
            DataType::U32,
        )
        .with_count(words.max(1)),
    );
    buffers.push(
        BufferDecl::storage(
            changed,
            CSR_FORWARD_OR_CHANGED_CHANGED_BUFFER,
            BufferAccess::ReadWrite,
            DataType::U32,
        )
        .with_count(1),
    );
    Program::wrapped(
        buffers,
        CSR_FORWARD_OR_CHANGED_PARALLEL_WORKGROUP_SIZE,
        vec![Node::Region {
            generator: Ident::from(OP_ID),
            source_region: None,
            body: Arc::new(body),
        }],
    )
}

/// Build the parallel expansion body used by production closure drivers and
/// large persistent-BFS programs.
#[must_use]
pub fn csr_forward_or_changed_parallel_body_prefixed(
    shape: ProgramGraphShape,
    frontier_out: &str,
    changed: &str,
    edge_kind_mask: u32,
    local_prefix: &str,
) -> Vec<Node> {
    csr_forward_or_changed_parallel_body_prefixed_impl(
        shape,
        frontier_out,
        changed,
        edge_kind_mask,
        local_prefix,
        None,
        None,
        None,
    )
}

/// Build one parallel expansion body that snapshots source-node activity
/// before any lane writes newly reached destination bits.
#[must_use]
pub fn csr_forward_or_changed_parallel_snapshot_body_prefixed(
    shape: ProgramGraphShape,
    frontier_out: &str,
    changed: &str,
    edge_kind_mask: u32,
    local_prefix: &str,
) -> Vec<Node> {
    csr_forward_or_changed_parallel_body_prefixed_impl(
        shape,
        frontier_out,
        changed,
        edge_kind_mask,
        local_prefix,
        Some(MemoryOrdering::GridSync),
        None,
        None,
    )
}

/// Build one snapshotting parallel expansion body and skip the expensive edge
/// scan when `active_gate` is zero. Newly discovered nodes set both
/// `changed[0]` and `changed[active_changed_index]`.
#[must_use]
pub(crate) fn csr_forward_or_changed_parallel_snapshot_body_prefixed_with_active(
    shape: ProgramGraphShape,
    frontier_out: &str,
    changed: &str,
    active_gate: Expr,
    active_changed_index: Expr,
    edge_kind_mask: u32,
    local_prefix: &str,
) -> Vec<Node> {
    csr_forward_or_changed_parallel_body_prefixed_impl(
        shape,
        frontier_out,
        changed,
        edge_kind_mask,
        local_prefix,
        Some(MemoryOrdering::GridSync),
        Some(active_gate),
        Some((changed, active_changed_index)),
    )
}

fn csr_forward_or_changed_parallel_body_prefixed_impl(
    shape: ProgramGraphShape,
    frontier_out: &str,
    changed: &str,
    edge_kind_mask: u32,
    local_prefix: &str,
    snapshot_barrier: Option<MemoryOrdering>,
    active_gate: Option<Expr>,
    extra_changed: Option<(&str, Expr)>,
) -> Vec<Node> {
    let local = |name: &str| -> String {
        if local_prefix.is_empty() {
            name.to_string()
        } else {
            format!("{local_prefix}_{name}")
        }
    };
    let src = Expr::gid_x();
    let in_bounds = local("in_bounds");
    let word_idx = local("word_idx");
    let bit_mask = local("bit_mask");
    let src_word = local("src_word");
    let src_active = local("src_active");
    let changed_old = local("changed_old");
    let extra_changed_old = local("extra_changed_old");

    let mark_changed = || {
        let mut nodes = vec![Node::let_bind(
            changed_old.as_str(),
            Expr::atomic_or(changed, Expr::u32(0), Expr::u32(1)),
        )];
        if let Some((extra_changed_buffer, extra_changed_index)) = &extra_changed {
            nodes.push(Node::let_bind(
                extra_changed_old.as_str(),
                Expr::atomic_or(
                    *extra_changed_buffer,
                    extra_changed_index.clone(),
                    Expr::u32(1),
                ),
            ));
        }
        nodes
    };

    // The neighbor expansion is the ONE canonical CSR edge-scan (identity frontier
    // index, a single bitset, and `mark_changed` on each newly-set bit). This is
    // the EDGE-WALK-ONLY level: the source-activity read is done by THIS function
    // (inline in the non-snapshot path, or snapshotted before a barrier), so it wraps
    // `csr_edge_expand_nodes` rather than the full `csr_edge_scan_nodes`. Byte-identical
    // to the previous hand-written `edge_scan` closure. `mark_changed` is borrowed so
    // both call sites can reuse it.
    let edge_scan = || {
        csr_edge_expand_nodes(
            shape,
            frontier_out,
            src.clone(),
            |word| word,
            &mark_changed,
            edge_kind_mask,
            local_prefix,
        )
    };

    if let Some(ordering) = snapshot_barrier {
        let ungated_src_active = Expr::select(
            Expr::var(in_bounds.as_str()),
            Expr::bitand(Expr::var(src_word.as_str()), Expr::var(bit_mask.as_str())),
            Expr::u32(0),
        );
        let src_active_expr = if let Some(active_gate) = active_gate {
            Expr::select(
                Expr::ne(active_gate, Expr::u32(0)),
                ungated_src_active,
                Expr::u32(0),
            )
        } else {
            ungated_src_active
        };
        return vec![
            Node::let_bind(
                in_bounds.as_str(),
                Expr::lt(src.clone(), Expr::u32(shape.node_count)),
            ),
            Node::let_bind(
                word_idx.as_str(),
                Expr::select(
                    Expr::var(in_bounds.as_str()),
                    Expr::shr(src.clone(), Expr::u32(5)),
                    Expr::u32(0),
                ),
            ),
            Node::let_bind(
                bit_mask.as_str(),
                Expr::shl(Expr::u32(1), Expr::bitand(src.clone(), Expr::u32(31))),
            ),
            Node::let_bind(
                src_word.as_str(),
                Expr::load(frontier_out, Expr::var(word_idx.as_str())),
            ),
            Node::let_bind(src_active.as_str(), src_active_expr),
            Node::barrier_with_ordering(ordering),
            Node::if_then(
                Expr::ne(Expr::var(src_active.as_str()), Expr::u32(0)),
                edge_scan(),
            ),
        ];
    }

    let body = vec![
        Node::let_bind(word_idx.as_str(), Expr::shr(src.clone(), Expr::u32(5))),
        Node::let_bind(
            bit_mask.as_str(),
            Expr::shl(Expr::u32(1), Expr::bitand(src.clone(), Expr::u32(31))),
        ),
        Node::let_bind(
            src_word.as_str(),
            Expr::load(frontier_out, Expr::var(word_idx.as_str())),
        ),
        Node::if_then(
            Expr::ne(
                Expr::bitand(Expr::var(src_word.as_str()), Expr::var(bit_mask.as_str())),
                Expr::u32(0),
            ),
            edge_scan(),
        ),
    ];

    vec![Node::if_then(
        Expr::lt(Expr::gid_x(), Expr::u32(shape.node_count)),
        body,
    )]
}

/// Wrap a parallel expansion body as a child Region of `parent_op_id`.
#[must_use]
pub fn csr_forward_or_changed_parallel_child_prefixed(
    parent_op_id: &str,
    shape: ProgramGraphShape,
    frontier_out: &str,
    changed: &str,
    edge_kind_mask: u32,
    local_prefix: &str,
) -> Node {
    Node::Region {
        generator: Ident::from(OP_ID),
        source_region: Some(GeneratorRef {
            name: parent_op_id.to_string(),
        }),
        body: Arc::new(csr_forward_or_changed_parallel_body_prefixed(
            shape,
            frontier_out,
            changed,
            edge_kind_mask,
            local_prefix,
        )),
    }
}

/// Wrap a snapshotting parallel expansion body as a child Region.
#[must_use]
pub fn csr_forward_or_changed_parallel_snapshot_child_prefixed(
    parent_op_id: &str,
    shape: ProgramGraphShape,
    frontier_out: &str,
    changed: &str,
    edge_kind_mask: u32,
    local_prefix: &str,
) -> Node {
    Node::Region {
        generator: Ident::from(OP_ID),
        source_region: Some(GeneratorRef {
            name: parent_op_id.to_string(),
        }),
        body: Arc::new(csr_forward_or_changed_parallel_snapshot_body_prefixed(
            shape,
            frontier_out,
            changed,
            edge_kind_mask,
            local_prefix,
        )),
    }
}

/// Wrap an active-gated snapshotting parallel expansion body as a child Region.
#[must_use]
pub fn csr_forward_or_changed_parallel_snapshot_child_prefixed_with_active(
    parent_op_id: &str,
    shape: ProgramGraphShape,
    frontier_out: &str,
    changed: &str,
    active_gate: Expr,
    active_changed_index: Expr,
    edge_kind_mask: u32,
    local_prefix: &str,
) -> Node {
    Node::Region {
        generator: Ident::from(OP_ID),
        source_region: Some(GeneratorRef {
            name: parent_op_id.to_string(),
        }),
        body: Arc::new(
            csr_forward_or_changed_parallel_snapshot_body_prefixed_with_active(
                shape,
                frontier_out,
                changed,
                active_gate,
                active_changed_index,
                edge_kind_mask,
                local_prefix,
            ),
        ),
    }
}