src/pipeline.rs

//! Generate → execute → filter → record. The data engine's main loop. Each kept
//! program is independently reproducible from (master seed, per-program seed).

use std::collections::hash_map::DefaultHasher;
use std::collections::HashSet;
use std::fs::{self, File};
use std::hash::{Hash, Hasher};
use std::io::{BufWriter, Write};
use std::path::Path;

use anyhow::{Context, Result};
use rand::{RngCore, SeedableRng};
use rand_chacha::ChaCha20Rng;

use crate::grammar::{GenConfig, Generator};
use crate::mutate;
use crate::sandbox::Sandbox;
use crate::spec::{self, Features, Record, RejectStats};
use crate::verify::{self, Verdict};

/// Mutation-gate parameters.
const N_PROBES: usize = 24;
const MAX_ADDED_TESTS: usize = 8;

pub struct GenParams {
    pub count: usize,
    pub seed: u64,
    pub fixed_difficulty: Option<u8>,
    pub n_examples: usize,
    pub n_tests: usize,
    pub budget: i64,
    pub out: String,
    /// Run the mutation gate (strengthen tests). Off for fast bulk training data.
    pub gate: bool,
}

pub fn run(p: GenParams) -> Result<()> {
    if let Some(parent) = Path::new(&p.out).parent() {
        if !parent.as_os_str().is_empty() {
            fs::create_dir_all(parent).ok();
        }
    }
    let file = File::create(&p.out).with_context(|| format!("creating {}", p.out))?;
    let mut w = BufWriter::new(file);

    let sb = Sandbox::new().map_err(|e| anyhow::anyhow!("init sandbox: {e}"))?;
    let mut master = ChaCha20Rng::seed_from_u64(p.seed);

    let mut stats = RejectStats::default();
    let mut kept = 0usize;
    let mut attempts = 0usize;
    let mut self_verify_fail = 0usize;
    let mut seen: HashSet<u64> = HashSet::new();
    let mut total_added_tests = 0usize;
    let mut total_mutants = 0usize;
    let mut total_distinct = 0usize;
    let max_attempts = p.count.saturating_mul(50).max(1000);

    while kept < p.count && attempts < max_attempts {
        attempts += 1;
        let diff = p
            .fixed_difficulty
            .unwrap_or(((attempts - 1) % 5) as u8);
        let prog_seed = master.next_u64();
        let mut prng = ChaCha20Rng::seed_from_u64(prog_seed);

        let cfg = GenConfig::for_difficulty(diff);
        let prog = Generator::new(&mut prng, cfg).generate();

        // Uniqueness: never keep the same source twice. Critical for the
        // data-efficiency study, which controls how many UNIQUE programs exist.
        let mut hasher = DefaultHasher::new();
        prog.source.hash(&mut hasher);
        if !seen.insert(hasher.finish()) {
            stats.duplicate += 1;
            continue;
        }

        let extracted = spec::extract(
            &sb, &mut prng, &prog, p.n_examples, p.n_tests, p.budget, &mut stats,
        );
        let (examples, tests) = match extracted {
            Some(x) => x,
            None => continue,
        };

        // Invariant: the reference program must satisfy its own held-out tests.
        if verify::verify(&sb, &prog.source, &tests, p.budget) != Verdict::Pass {
            self_verify_fail += 1;
            continue;
        }

        // Mutation gate: strengthen tests until they catch every behaviorally
        // distinct mutant. Reject if that can't be done within the cap. Skipped
        // for bulk training data (--no-gate), where test strength doesn't matter.
        let tests = if p.gate {
            let probes = spec::sample_inputs(&prog.params, N_PROBES, &mut prng);
            match mutate::strengthen(&sb, &prog.source, tests, &probes, p.budget, MAX_ADDED_TESTS) {
                Some(g) => {
                    total_added_tests += g.added;
                    total_mutants += g.mutants_total;
                    total_distinct += g.mutants_distinct;
                    g.tests
                }
                None => {
                    stats.weak_tests += 1;
                    continue;
                }
            }
        } else {
            tests
        };

        let infill = spec::make_infill(&prog.source, &mut prng);
        let params = prog.params.iter().map(|t| spec::type_name(*t).to_string()).collect();
        let output = examples.first().map(|e| e.output.clone()).unwrap_or(crate::value::LValue::Nil);
        let features = Features::build(&prog.source, &prog.features, &output);

        let rec = Record {
            id: format!("echo1-{:08}", kept),
            difficulty: diff,
            seed: prog_seed,
            source: prog.source,
            params,
            examples,
            tests,
            infill,
            features,
        };

        let line = serde_json::to_string(&rec)?;
        w.write_all(line.as_bytes())?;
        w.write_all(b"\n")?;
        kept += 1;
    }

    w.flush()?;

    eprintln!("kept       {kept} / requested {}", p.count);
    eprintln!("attempts   {attempts}");
    eprintln!("rejects    load_fail={} run_fail={} nonobservable={} nondet={} noop/sparse={} dup={} weak_tests={}",
        stats.load_fail, stats.run_fail, stats.nonobservable, stats.nondeterministic, stats.noop, stats.duplicate, stats.weak_tests);
    eprintln!("self_verify_fail {self_verify_fail} (should be 0)");
    eprintln!("mutation gate: {total_mutants} mutants, {total_distinct} behaviorally distinct, {total_added_tests} discriminating tests added across {kept} kept programs");
    eprintln!("wrote      {}", p.out);

    if kept < p.count {
        eprintln!("WARN: hit attempt cap before reaching requested count");
    }
    Ok(())
}