File size: 3,635 Bytes
3afc977
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
//! echo-1 M0 data engine.
//!
//! Generates pure-Lua functions, executes them in a deterministic sandbox,
//! filters to the clean-terminating observable ones, and writes a JSONL dataset
//! where each record carries an I/O spec (synthesis), a prefix/hole/suffix spec
//! (infilling), and held-out tests for execution-based verification.

mod grammar;
mod mutate;
mod pipeline;
mod sandbox;
mod spec;
mod value;
mod verify;

use std::io::{BufRead, Write};

use anyhow::Result;
use clap::{Parser, Subcommand};
use serde::{Deserialize, Serialize};

use spec::IoPair;

#[derive(Parser)]
#[command(name = "echo-data", about = "echo-1 M0 data engine")]
struct Cli {
    #[command(subcommand)]
    cmd: Cmd,
}

#[derive(Subcommand)]
enum Cmd {
    /// Generate a dataset of verified pure-Lua programs.
    Gen {
        /// Number of programs to keep.
        #[arg(short, long, default_value_t = 100)]
        count: usize,
        /// Master RNG seed (reproducible).
        #[arg(short, long, default_value_t = 1)]
        seed: u64,
        /// Fix difficulty 0..=4; omit to round-robin.
        #[arg(short, long)]
        difficulty: Option<u8>,
        /// Shown I/O examples per program.
        #[arg(long, default_value_t = 3)]
        examples: usize,
        /// Held-out tests per program.
        #[arg(long, default_value_t = 5)]
        tests: usize,
        /// Instruction budget per execution (deterministic timeout).
        #[arg(long, default_value_t = 200_000)]
        budget: i64,
        /// Output JSONL path.
        #[arg(short, long, default_value = "data/dataset.jsonl")]
        out: String,
        /// Skip the mutation gate (much faster; for bulk training data).
        #[arg(long, default_value_t = false)]
        no_gate: bool,
    },
    /// Verify candidate programs from stdin against their tests (one JSON object
    /// per line: {"source":..., "tests":[{"input":[...],"output":...}]}). Emits
    /// one result object per line. Same sandbox/verifier used to label data.
    VerifyBatch {
        /// Instruction budget per execution.
        #[arg(long, default_value_t = 200_000)]
        budget: i64,
    },
}

#[derive(Deserialize)]
struct Candidate {
    source: String,
    tests: Vec<IoPair>,
}

#[derive(Serialize)]
struct VerifyOut {
    pass: bool,
    verdict: String,
}

fn verify_batch(budget: i64) -> Result<()> {
    let sb = sandbox::Sandbox::new().map_err(|e| anyhow::anyhow!("init sandbox: {e}"))?;
    let stdin = std::io::stdin();
    let stdout = std::io::stdout();
    let mut out = std::io::BufWriter::new(stdout.lock());

    for line in stdin.lock().lines() {
        let line = line?;
        if line.trim().is_empty() {
            continue;
        }
        let cand: Candidate = serde_json::from_str(&line)?;
        let verdict = verify::verify(&sb, &cand.source, &cand.tests, budget);
        let res = VerifyOut {
            pass: verdict == verify::Verdict::Pass,
            verdict: format!("{verdict:?}"),
        };
        writeln!(out, "{}", serde_json::to_string(&res)?)?;
    }
    out.flush()?;
    Ok(())
}

fn main() -> Result<()> {
    let cli = Cli::parse();
    match cli.cmd {
        Cmd::Gen { count, seed, difficulty, examples, tests, budget, out, no_gate } => {
            pipeline::run(pipeline::GenParams {
                count,
                seed,
                fixed_difficulty: difficulty,
                n_examples: examples,
                n_tests: tests,
                budget,
                out,
                gate: !no_gate,
            })
        }
        Cmd::VerifyBatch { budget } => verify_batch(budget),
    }
}