weirflow 0.1.0

//! Legoblock composition boundary: Weir must compose Vyre IR through
//! `vyre-primitives`, never constructing raw `Expr::`, `Node::`, or `Region::`
//! directly in production paths.

use std::collections::HashSet;
use std::fs;
use std::path::{Path, PathBuf};

// =============================================================================
// Frozen allowlist of existing raw-IR constructions in `src/`.
//
// Policy: REMOVE-ONLY.  Do not add new entries.  Each entry is an exact
// `trim_end()` match of a line that currently contains `Expr::`, `Node::`, or
// `Region::`.  If a line is reformatted, the allowlist entry must be updated
// to the new exact text (this is intentional friction).
// =============================================================================
static RAW_IR_ALLOWLIST: &[(&str, &[&str])] = &[
    // analysis_family.rs (doc-comment only; describes the boundary policy itself).
    ("analysis_family.rs", &[
        "//!   `ProgramGraphShape` only; it never exposes `Expr::`, `Node::`, or",
        "//!   `Region::` construction.",
    ]),
    // callgraph/program.rs: GPU kernel program builder; raw IR is the payload.
    ("callgraph/program.rs", &[
        "    let w = Expr::InvocationId { axis: 0 };",
        "        Node::let_bind(\"direct\", Expr::load(direct_edges_in, w.clone())),",
        "        Node::let_bind(",
        "            Expr::bitand(",
        "                Expr::load(indirect_sites_in, w.clone()),",
        "                Expr::load(pts_closure_in, w.clone()),",
        "        Node::store(",
        "            Expr::bitor(Expr::var(\"direct\"), Expr::var(\"indirect\")),",
        "        vec![Node::Region {",
        "            body: Arc::new(vec![Node::if_then(",
        "                Expr::lt(w.clone(), Expr::u32(words)),",
    ]),
    // def_use.rs (doc-comment only; references upstream `DataFlow::Node::getUses()`).
    ("def_use.rs", &[
        "//! `DataFlow::Node::getUses()` and that every taint rule reaches for",
    ]),
    // dominators/program.rs: GPU kernel program builder.
    ("dominators/program.rs", &[
        "    let t = Expr::InvocationId { axis: 0 };",
        "        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(",
        "            Expr::select(",
        "                Expr::eq(Expr::var(\"word\"), Expr::u32(entry_word)),",
        "                Expr::u32(entry_bit),",
        "                Expr::u32(0),",
        "        Node::let_bind(",
        "            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(",
        "            Expr::load(pred_offsets, Expr::add(Expr::var(\"node\"), Expr::u32(1))),",
        "        Node::let_bind(\"acc\", Expr::var(\"row_mask\")),",
        "        Node::loop_for(",
        "            Expr::var(\"pred_start\"),",
        "            Expr::var(\"pred_end\"),",
        "                Node::let_bind(\"pred\", Expr::load(pred_targets, Expr::var(\"edge\"))),",
        "                Node::assign(",
        "                    Expr::bitand(",
        "                        Expr::var(\"acc\"),",
        "                        Expr::load(",
        "                            Expr::add(",
        "                                Expr::mul(Expr::var(\"pred\"), Expr::u32(words)),",
        "                                Expr::var(\"word\"),",
        "        Node::let_bind(",
        "            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(",
        "            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(",
        "            Expr::select(",
        "                Expr::eq(Expr::var(\"node\"), Expr::u32(entry)),",
        "                Expr::var(\"entry_only\"),",
        "                Expr::var(\"candidate\"),",
        "        vec![Node::Region {",
        "            body: Arc::new(vec![Node::if_then(",
        "                Expr::lt(t, Expr::u32(matrix_words)),",
    ]),
    // escape.rs: GPU kernel program builder; pattern match at line 131 is deconstruction, not construction.
    ("escape.rs", &[
        "    let w = Expr::InvocationId { axis: 0 };",
        "        Node::let_bind(\"pts\", Expr::load(points_to_in, w.clone())),",
        "        Node::let_bind(\"cg\", Expr::load(callgraph_in, w.clone())),",
        "        Node::let_bind(\"prev\", Expr::load(escape_out, w.clone())),",
        "        Node::store(",
        "            Expr::bitor(",
        "                Expr::bitor(Expr::var(\"pts\"), Expr::var(\"cg\")),",
        "                Expr::var(\"prev\"),",
        "        vec![Node::Region {",
        "            body: Arc::new(vec![Node::if_then(",
        "                Expr::lt(w.clone(), Expr::u32(words)),",
        "                Node::Region { generator, .. } if generator.as_str() == OP_ID",
    ]),
    // fixed_point_resident_sequence.rs. GPU kernel program builder (seed init region).
    ("fixed_point_resident_sequence.rs", &[
        "            vec![Node::Region {",
        "                body: Arc::new(vec![Node::store(\"changed\", Expr::u32(0), Expr::u32(0))]),",
    ]),
    // ifds_resident_direct_prepare.rs. GPU kernel program builder.
    ("ifds_resident_direct_prepare.rs", &[
        "    Expr::add(",
        "        Expr::mul(Expr::gid_y(), Expr::u32(DIRECT_LINEAR_TILE_X)),",
        "        Expr::gid_x(),",
        "        vec![Node::Region {",
        "                Node::if_then(",
        "                    Expr::lt(index.clone(), Expr::u32(node_words)),",
        "                    vec![Node::store(",
        "                        Expr::u32(0),",
        "                Node::if_then(",
        "                    Expr::lt(index.clone(), Expr::u32(edge_words)),",
        "                    vec![Node::store(",
        "                        Expr::u32(1),",
        "                Node::if_then(",
        "                    Expr::lt(index.clone(), Expr::u32(node_words)),",
        "                    vec![Node::store(DIRECT_PG_NODE_TAGS_BUFFER, index, Expr::u32(0))],",
    ]),
    // ifds_seed_frontiers.rs: GPU kernel program builder.
    ("ifds_seed_frontiers.rs", &[
        "    let word = Expr::InvocationId { axis: 0 };",
        "    let query = Expr::InvocationId { axis: 1 };",
        "        vec![Node::Region {",
        "            body: Arc::new(vec![Node::if_then(",
        "                Expr::lt(word.clone(), Expr::u32(words)),",
        "                vec![Node::store(",
        "                    Expr::add(Expr::mul(query, Expr::u32(words)), word),",
        "                    Expr::u32(0),",
        "    let local_seed = Expr::InvocationId { axis: 0 };",
        "    let query = Expr::InvocationId { axis: 1 };",
        "    let seed_idx = Expr::add(",
        "        Expr::load(IFDS_SEED_OFFSETS, query.clone()),",
        "    let seed_base = Expr::mul(Expr::var(\"seed_idx\"), Expr::u32(3));",
        "    let dense = Expr::add(",
        "        Expr::add(",
        "            Expr::mul(",
        "                Expr::load(IFDS_SEED_TRIPLES, seed_base.clone()),",
        "                Expr::u32(block_fact_span),",
        "            Expr::mul(",
        "                Expr::load(",
        "                    Expr::add(seed_base.clone(), Expr::u32(1)),",
        "                Expr::u32(shape.facts_per_proc),",
        "        Expr::load(IFDS_SEED_TRIPLES, Expr::add(seed_base, Expr::u32(2))),",
        "        Node::let_bind(\"seed_start\", Expr::load(IFDS_SEED_OFFSETS, query.clone())),",
        "        Node::let_bind(",
        "            Expr::load(IFDS_SEED_OFFSETS, Expr::add(query.clone(), Expr::u32(1))),",
        "        Node::if_then(",
        "            Expr::lt(",
        "                Expr::sub(Expr::var(\"seed_end\"), Expr::var(\"seed_start\")),",
        "                Node::let_bind(\"seed_idx\", seed_idx),",
        "                Node::let_bind(\"dense\", dense),",
        "                Node::let_bind(",
        "                    Expr::add(",
        "                        Expr::mul(query, Expr::u32(words)),",
        "                        Expr::shr(Expr::var(\"dense\"), Expr::u32(5)),",
        "                Node::let_bind(",
        "                    Expr::shl(",
        "                        Expr::u32(1),",
        "                        Expr::bitand(Expr::var(\"dense\"), Expr::u32(31)),",
        "                Node::let_bind(",
        "                    Expr::atomic_or(",
        "                        Expr::var(\"frontier_word\"),",
        "                        Expr::var(\"frontier_bit\"),",
        "        vec![Node::Region {",
        "            body: Arc::new(vec![Node::if_then(",
        "                Expr::lt(local_seed.clone(), Expr::u32(max_seeds_per_query)),",
    ]),
    // loop_sum/program.rs: GPU kernel program builder.
    ("loop_sum/program.rs", &[
        "    let v = Expr::InvocationId { axis: 0 };",
        "        Node::let_bind(\"lo_idx\", Expr::mul(v.clone(), Expr::u32(2))),",
        "        Node::let_bind(\"hi_idx\", Expr::add(Expr::var(\"lo_idx\"), Expr::u32(1))),",
        "        Node::let_bind(\"prev_lo\", Expr::load(cfg_in, Expr::var(\"lo_idx\"))),",
        "        Node::let_bind(\"prev_hi\", Expr::load(cfg_in, Expr::var(\"hi_idx\"))),",
        "        Node::let_bind(\"new_lo\", Expr::load(ranges_in, Expr::var(\"lo_idx\"))),",
        "        Node::let_bind(\"new_hi\", Expr::load(ranges_in, Expr::var(\"hi_idx\"))),",
        "        Node::let_bind(",
        "            Expr::select(",
        "                Expr::lt(Expr::var(\"new_lo\"), Expr::var(\"prev_lo\")),",
        "                Expr::u32(NEG_INF),",
        "                Expr::var(\"prev_lo\"),",
        "        Node::let_bind(",
        "            Expr::select(",
        "                Expr::lt(Expr::var(\"prev_hi\"), Expr::var(\"new_hi\")),",
        "                Expr::u32(POS_INF),",
        "                Expr::var(\"prev_hi\"),",
        "        Node::store(summary_out, Expr::var(\"lo_idx\"), Expr::var(\"wide_lo\")),",
        "        Node::store(summary_out, Expr::var(\"hi_idx\"), Expr::var(\"wide_hi\")),",
        "        vec![Node::Region {",
        "            body: Arc::new(vec![Node::if_then(",
        "                Expr::lt(v.clone(), Expr::u32(var_count)),",
    ]),
    // may_alias.rs: GPU kernel program builder.
    ("may_alias.rs", &[
        "    let t = Expr::InvocationId { axis: 0 };",
        "        Node::let_bind(",
        "            Expr::bitand(Expr::load(pts_p, t.clone()), Expr::load(pts_q, t.clone())),",
        "        Node::store(intersect_buf, t.clone(), Expr::var(\"intersect\")),",
        "        Node::if_then(",
        "            Expr::ne(Expr::var(\"intersect\"), Expr::u32(0)),",
        "            vec![Node::let_bind(",
        "                Expr::atomic_or(out_scalar, Expr::u32(0), Expr::u32(1)),",
        "        vec![Node::Region {",
        "            body: Arc::new(vec![Node::if_then(",
        "                Expr::lt(t.clone(), Expr::u32(words)),",
    ]),
    // range/program.rs: GPU kernel program builder.
    ("range/program.rs", &[
        "    let v = Expr::InvocationId { axis: 0 };",
        "        Node::let_bind(\"lo_idx\", Expr::mul(v.clone(), Expr::u32(2))),",
        "        Node::let_bind(\"hi_idx\", Expr::add(Expr::var(\"lo_idx\"), Expr::u32(1))),",
        "        Node::let_bind(\"lo\", Expr::load(defs_in, Expr::var(\"lo_idx\"))),",
        "        Node::let_bind(\"hi\", Expr::load(defs_in, Expr::var(\"hi_idx\"))),",
        "        Node::let_bind(\"t_lo\", Expr::load(edges_in, Expr::var(\"lo_idx\"))),",
        "        Node::let_bind(\"t_hi\", Expr::load(edges_in, Expr::var(\"hi_idx\"))),",
        "        Node::store(",
        "            Expr::var(\"lo_idx\"),",
        "            Expr::add(Expr::var(\"lo\"), Expr::var(\"t_lo\")),",
        "        Node::store(",
        "            Expr::var(\"hi_idx\"),",
        "            Expr::add(Expr::var(\"hi\"), Expr::var(\"t_hi\")),",
        "        vec![Node::Region {",
        "            body: Arc::new(vec![Node::if_then(",
        "                Expr::lt(v.clone(), Expr::u32(var_count)),",
    ]),
    // range_check.rs: GPU kernel program builder.
    ("range_check.rs", &[
        "    let t = Expr::InvocationId { axis: 0 };",
        "    let body = vec![Node::store(",
        "        Expr::select(",
        "            Expr::lt(",
        "                Expr::load(interval_hi, t.clone()),",
        "                Expr::load(bound, t.clone()),",
        "            Expr::u32(1),",
        "            Expr::u32(0),",
        "        vec![Node::Region {",
        "            body: Arc::new(vec![Node::if_then(",
        "                Expr::lt(t.clone(), Expr::u32(node_count)),",
    ]),
    // reaching_def_summary.rs: GPU kernel program builder (uses Expr::popcount).
    ("reaching_def_summary.rs", &[
        "    let tid = Expr::gid_x();",
        "            vec![Node::if_then(",
        "                Expr::lt(tid.clone(), Expr::u32(words)),",
        "                    Node::let_bind(",
        "                        Expr::bitand(tid.clone(), Expr::u32(partial_bins - 1)),",
        "                    Node::let_bind(",
        "                        Expr::mul(Expr::var(\"summary_bin\"), Expr::u32(3)),",
        "                    Node::let_bind(",
        "                        Expr::bitand(",
        "                            Expr::load(gen_kill_in, tid.clone()),",
        "                            Expr::load(use_set, tid.clone()),",
        "                    Node::let_bind(",
        "                        Expr::atomic_add(",
        "                            Expr::var(\"summary_base\"),",
        "                            Expr::popcount(Expr::var(\"word\")),",
        "                    Node::if_then(",
        "                        Expr::ne(Expr::var(\"word\"), Expr::u32(0)),",
        "                            Node::let_bind(",
        "                                Expr::atomic_or(",
        "                                    Expr::add(Expr::var(\"summary_base\"), Expr::u32(1)),",
        "                                    Expr::u32(1),",
        "                            Node::let_bind(",
        "                                Expr::bitxor(",
        "                                    Expr::var(\"word\"),",
        "                                    Expr::mul(",
        "                                        Expr::add(tid.clone(), Expr::u32(1)),",
        "                                        Expr::u32(0x9E37_79B9),",
        "                            Node::let_bind(",
        "                                Expr::atomic_xor(",
        "                                    Expr::add(Expr::var(\"summary_base\"), Expr::u32(2)),",
        "                                    Expr::var(\"mixed\"),",
    ]),
    // summary/program.rs: GPU kernel program builder.
    ("summary/program.rs", &[
        "    let w = Expr::InvocationId { axis: 0 };",
        "        Node::let_bind(\"ast\", Expr::load(fn_ast_in, w.clone())),",
        "        Node::let_bind(\"cg\", Expr::load(callgraph_in, w.clone())),",
        "        Node::let_bind(\"cached\", Expr::load(cached_summary_in, w.clone())),",
        "        Node::store(",
        "            Expr::bitor(",
        "                Expr::bitor(Expr::var(\"ast\"), Expr::var(\"cg\")),",
        "                Expr::var(\"cached\"),",
        "        vec![Node::Region {",
        "            body: Arc::new(vec![Node::if_then(",
        "                Expr::lt(w.clone(), Expr::u32(words)),",
    ]),
];

fn collect_rust_files(dir: &Path, out: &mut Vec<PathBuf>) {
    for entry in fs::read_dir(dir).expect("weir src directory must be readable") {
        let entry = entry.expect("weir src entry must be readable");
        let path = entry.path();
        if path.is_dir() {
            collect_rust_files(&path, out);
        } else if path.extension().and_then(|ext| ext.to_str()) == Some("rs") {
            out.push(path);
        }
    }
}

fn is_oracle_path(path: &Path) -> bool {
    path.components()
        .any(|component| component.as_os_str().to_str() == Some("oracle"))
        || path
            .file_name()
            .and_then(|name| name.to_str())
            .is_some_and(|name| name == "oracle.rs" || name.ends_with("_cpu_oracle.rs"))
}

fn is_test_source_path(path: &Path) -> bool {
    path.components()
        .filter_map(|component| component.as_os_str().to_str())
        .any(|name| name == "tests" || name.ends_with("_tests") || name == "test_harness")
        || path
            .file_name()
            .and_then(|name| name.to_str())
            .is_some_and(|name| name == "tests.rs" || name.ends_with("_tests.rs"))
}

#[test]
fn raw_ir_construction_is_allowlisted() {
    let src = Path::new(env!("CARGO_MANIFEST_DIR")).join("src");
    let mut files = Vec::new();
    collect_rust_files(&src, &mut files);

    let mut allowlist_by_file: std::collections::HashMap<&str, HashSet<&str>> =
        std::collections::HashMap::new();
    for (file, lines) in RAW_IR_ALLOWLIST {
        let mut set = HashSet::new();
        for line in *lines {
            set.insert(*line);
        }
        allowlist_by_file.insert(*file, set);
    }

    let mut violations = Vec::new();
    for path in files {
        if is_oracle_path(&path) || is_test_source_path(&path) {
            continue;
        }
        let rel = path.strip_prefix(&src).unwrap_or(&path);
        let rel_str = rel.to_str().unwrap();

        let text = fs::read_to_string(&path).expect("weir source file must be readable");
        let allowed = allowlist_by_file.get(rel_str);

        for (idx, line) in text.lines().enumerate() {
            let trimmed = line.trim_end();
            if trimmed.contains("Expr::")
                || trimmed.contains("Node::")
                || trimmed.contains("Region::")
            {
                if let Some(allowed_lines) = allowed {
                    if allowed_lines.contains(trimmed) {
                        continue;
                    }
                }
                violations.push(format!("{}:{}: {}", rel_str, idx + 1, trimmed));
            }
        }
    }

    assert!(
        violations.is_empty(),
        "Raw Vyre IR construction (`Expr::`, `Node::`, `Region::`) must go through \
         `vyre-primitives`. New violations found:\n{}",
        violations.join("\n")
    );
}