src/grammar.rs

//! Restricted-grammar generator for pure-Lua functions.
//!
//! Emits a single global `function f(...)` over typed parameters. Generation is
//! type-aware so that programs run cleanly instead of mostly raising
//! nil-arithmetic / index errors — the execution filter then keeps only the
//! terminating, observable ones. Difficulty knobs (DATA.md) are the axes along
//! which global structure grows: nesting depth, def-use count, length, table
//! size, loop bound.

use std::collections::HashSet;

use rand::Rng;
use rand_chacha::ChaCha20Rng;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Type {
    Number,
    Bool,
    List, // list of Number
    Str,
}

/// Structural features of a generated program, recorded as metadata so eval can
/// stratify and the data-efficiency study can group by construct (EXPERIMENTS.md).
#[derive(Debug, Clone, Default)]
pub struct ProgramFeatures {
    pub recursion: bool,
    pub closure: bool,
    pub loops: bool,
    pub table_build: bool,
}

#[derive(Debug, Clone)]
pub struct GenConfig {
    pub n_params: usize,
    pub n_locals: usize,
    pub max_depth: usize,
    pub expr_fuel: usize,
    pub list_len: usize,
    pub loop_bound: i64,
}

impl GenConfig {
    /// Difficulty 0..=4 maps onto concrete knob settings.
    pub fn for_difficulty(d: u8) -> GenConfig {
        match d {
            0 => GenConfig { n_params: 1, n_locals: 1, max_depth: 1, expr_fuel: 2, list_len: 3, loop_bound: 4 },
            1 => GenConfig { n_params: 2, n_locals: 2, max_depth: 2, expr_fuel: 3, list_len: 4, loop_bound: 6 },
            2 => GenConfig { n_params: 2, n_locals: 3, max_depth: 3, expr_fuel: 4, list_len: 5, loop_bound: 8 },
            3 => GenConfig { n_params: 3, n_locals: 4, max_depth: 4, expr_fuel: 5, list_len: 6, loop_bound: 10 },
            _ => GenConfig { n_params: 3, n_locals: 5, max_depth: 5, expr_fuel: 6, list_len: 8, loop_bound: 12 },
        }
    }
}

/// A generated program plus the parameter types needed to sample inputs.
pub struct Program {
    pub source: String,
    pub params: Vec<Type>,
    pub features: ProgramFeatures,
}

pub struct Generator<'a> {
    rng: &'a mut ChaCha20Rng,
    cfg: GenConfig,
    scope: Vec<(String, Type)>,
    /// Names whose value derives (transitively) from a parameter. Used to keep
    /// returns input-dependent so they aren't constant (the noop filter).
    tainted: HashSet<String>,
    feat: ProgramFeatures,
    next_id: usize,
}

impl<'a> Generator<'a> {
    pub fn new(rng: &'a mut ChaCha20Rng, cfg: GenConfig) -> Self {
        Generator {
            rng,
            cfg,
            scope: Vec::new(),
            tainted: HashSet::new(),
            feat: ProgramFeatures::default(),
            next_id: 0,
        }
    }

    pub fn generate(&mut self) -> Program {
        self.scope.clear();
        self.tainted.clear();
        self.feat = ProgramFeatures::default();
        self.next_id = 0;

        // Parameters: at least one Number, the rest random typed. Every
        // parameter is, by definition, input-dependent (tainted).
        let mut params = Vec::new();
        for i in 0..self.cfg.n_params {
            let ty = if i == 0 { Type::Number } else { self.rand_type() };
            let name = self.fresh();
            self.tainted.insert(name.clone());
            self.scope.push((name, ty));
            params.push(ty);
        }
        let param_names: Vec<String> =
            self.scope.iter().map(|(n, _)| n.clone()).collect();

        let mut lines: Vec<String> = Vec::new();

        // Global-structure features that def-use chains alone don't give. Higher
        // difficulty only; each is the kind of long-range dependency AR pays for
        // and bidirectional refinement is meant to exploit (DATA.md).
        if self.cfg.max_depth >= 3 && self.rng.gen_bool(0.6) {
            self.gen_recursive_helper(&mut lines, 1);
        }
        // A closure capturing a local declared earlier: a long-range binding.
        if self.cfg.max_depth >= 2 && self.rng.gen_bool(0.5) {
            self.gen_closure_helper(&mut lines, 1);
        }
        // A table built incrementally where each element depends on the previous
        // one: table coherence + a def-use chain threaded through the table.
        if self.cfg.max_depth >= 2 && self.rng.gen_bool(0.5) {
            self.gen_table_build(&mut lines, 1);
        }

        let n_locals = self.cfg.n_locals;
        for _ in 0..n_locals {
            self.gen_local(&mut lines, 1);
        }
        // Control-flow body adds nesting / global structure.
        self.gen_block(&mut lines, 1, self.cfg.max_depth);
        // Final return, anchored on an input-dependent value.
        let ret = self.gen_return(1);
        lines.push(ret);

        let mut src = format!("function f({})\n", param_names.join(", "));
        for l in &lines {
            src.push_str(l);
            src.push('\n');
        }
        src.push_str("end\n");

        Program { source: src, params, features: self.feat.clone() }
    }

    fn fresh(&mut self) -> String {
        self.next_id += 1;
        format!("v{}", self.next_id)
    }

    fn rand_type(&mut self) -> Type {
        match self.rng.gen_range(0..4) {
            0 => Type::Number,
            1 => Type::Bool,
            2 => Type::List,
            _ => Type::Str,
        }
    }

    fn vars_of(&self, ty: Type) -> Vec<String> {
        self.scope
            .iter()
            .filter(|(_, t)| *t == ty)
            .map(|(n, _)| n.clone())
            .collect()
    }

    /// Does an emitted expression reference any input-dependent variable?
    /// Tokenizes on identifier boundaries so `v1` doesn't match `v10`.
    fn refs_tainted(&self, s: &str) -> bool {
        let mut cur = String::new();
        for ch in s.chars().chain(std::iter::once(' ')) {
            if ch.is_alphanumeric() || ch == '_' {
                cur.push(ch);
            } else if !cur.is_empty() {
                if self.tainted.contains(&cur) {
                    return true;
                }
                cur.clear();
            }
        }
        false
    }

    /// A randomly chosen tainted Number variable. Parameter 0 is always one, so
    /// this never fails after parameters are in scope.
    fn anchor_num(&mut self) -> String {
        let cands: Vec<String> = self
            .scope
            .iter()
            .filter(|(n, t)| *t == Type::Number && self.tainted.contains(n))
            .map(|(n, _)| n.clone())
            .collect();
        self.pick(&cands).unwrap_or_else(|| "0".to_string())
    }

    fn indent(depth: usize) -> String {
        "  ".repeat(depth)
    }

    // ---- statements ----

    fn gen_local(&mut self, lines: &mut Vec<String>, depth: usize) {
        let ty = self.rand_type();
        let name = self.fresh();
        let rhs = self.gen_expr(ty, self.cfg.expr_fuel);
        if self.refs_tainted(&rhs) {
            self.tainted.insert(name.clone());
        }
        lines.push(format!("{}local {} = {}", Self::indent(depth), name, rhs));
        self.scope.push((name, ty));
    }

    fn gen_block(&mut self, lines: &mut Vec<String>, depth: usize, fuel: usize) {
        // At the fuel floor emit a single leaf statement rather than nothing, so
        // control-flow bodies are never empty (no dead `if ... then end` in the
        // training data).
        if fuel == 0 {
            if self.rng.gen_bool(0.5) {
                self.gen_assign(lines, depth);
            } else {
                self.gen_local(lines, depth);
            }
            return;
        }
        let n_stmts = self.rng.gen_range(1..=2 + fuel.min(2));
        for _ in 0..n_stmts {
            match self.rng.gen_range(0..4) {
                0 => self.gen_local(lines, depth),
                1 => self.gen_assign(lines, depth),
                2 => self.gen_if(lines, depth, fuel),
                _ => self.gen_for(lines, depth, fuel),
            }
        }
    }

    fn gen_assign(&mut self, lines: &mut Vec<String>, depth: usize) {
        // Reassign an existing Number var to build a def-use chain.
        let nums = self.vars_of(Type::Number);
        if let Some(name) = self.pick(&nums) {
            let rhs = self.gen_expr(Type::Number, self.cfg.expr_fuel);
            if self.refs_tainted(&rhs) {
                self.tainted.insert(name.clone());
            }
            lines.push(format!("{}{} = {}", Self::indent(depth), name, rhs));
        } else {
            self.gen_local(lines, depth);
        }
    }

    fn gen_if(&mut self, lines: &mut Vec<String>, depth: usize, fuel: usize) {
        let cond = self.gen_expr(Type::Bool, self.cfg.expr_fuel);
        lines.push(format!("{}if {} then", Self::indent(depth), cond));
        let mark = self.scope.len();
        self.gen_block(lines, depth + 1, fuel - 1);
        self.scope.truncate(mark); // locals leave scope at block end
        if self.rng.gen_bool(0.5) {
            lines.push(format!("{}else", Self::indent(depth)));
            self.gen_block(lines, depth + 1, fuel - 1);
            self.scope.truncate(mark);
        }
        lines.push(format!("{}end", Self::indent(depth)));
    }

    fn gen_for(&mut self, lines: &mut Vec<String>, depth: usize, fuel: usize) {
        // Accumulator pattern: declare a Number before the loop, mutate inside.
        let acc = self.fresh();
        lines.push(format!("{}local {} = 0", Self::indent(depth), acc));
        self.scope.push((acc.clone(), Type::Number));

        self.feat.loops = true;
        let n = self.rng.gen_range(1..=self.cfg.loop_bound);
        let i = self.fresh();
        lines.push(format!("{}for {} = 1, {} do", Self::indent(depth), i, n));
        let mark = self.scope.len();
        self.scope.push((i.clone(), Type::Number));
        // Body always advances the accumulator.
        let inc = self.gen_expr(Type::Number, self.cfg.expr_fuel.min(3));
        if self.refs_tainted(&inc) {
            self.tainted.insert(acc.clone());
        }
        lines.push(format!("{}{} = {} + ({})", Self::indent(depth + 1), acc, acc, inc));
        self.gen_block(lines, depth + 1, fuel.saturating_sub(2));
        self.scope.truncate(mark);
        lines.push(format!("{}end", Self::indent(depth)));
    }

    fn gen_return(&mut self, depth: usize) -> String {
        let fuel = self.cfg.expr_fuel;
        let e = match self.rng.gen_range(0..4) {
            // String: reuse a tainted string, else concatenate an input-dependent
            // number's text with more string computation.
            3 => {
                let tstrs: Vec<String> = self
                    .scope
                    .iter()
                    .filter(|(n, t)| *t == Type::Str && self.tainted.contains(n))
                    .map(|(n, _)| n.clone())
                    .collect();
                if !tstrs.is_empty() && self.rng.gen_bool(0.5) {
                    self.pick(&tstrs).unwrap()
                } else {
                    let a = self.anchor_num();
                    format!("(tostring({}) .. {})", a, self.gen_str(fuel.min(2)))
                }
            }
            // Bool: comparison anchored on an input-dependent number.
            0 => {
                let a = self.anchor_num();
                let cmp = ["==", "~=", "<", "<=", ">", ">="][self.rng.gen_range(0..6)];
                format!("({} {} ({}))", a, cmp, self.gen_num(fuel))
            }
            // List: reuse a tainted list, else a literal that includes a tainted number.
            1 => {
                let tlists: Vec<String> = self
                    .scope
                    .iter()
                    .filter(|(n, t)| *t == Type::List && self.tainted.contains(n))
                    .map(|(n, _)| n.clone())
                    .collect();
                if !tlists.is_empty() && self.rng.gen_bool(0.5) {
                    self.pick(&tlists).unwrap()
                } else {
                    let a = self.anchor_num();
                    let extra = self.rng.gen_range(0..self.cfg.list_len.max(1));
                    let mut elems = vec![a];
                    for _ in 0..extra {
                        elems.push(self.gen_num(fuel.min(2)));
                    }
                    format!("{{{}}}", elems.join(", "))
                }
            }
            // Number: combine an input-dependent number with more computation.
            _ => {
                let a = self.anchor_num();
                let op = ["+", "-", "*"][self.rng.gen_range(0..3)];
                format!("({} {} ({}))", a, op, self.gen_num(fuel))
            }
        };
        format!("{}return {}", Self::indent(depth), e)
    }

    /// Emit a well-founded recursive helper and bind its result to a tainted
    /// Number local. The argument is `math.abs(<input number>) % bound`, so it
    /// is non-negative and bounded; the `n <= 0` base case guarantees
    /// termination in at most `bound` calls.
    fn gen_recursive_helper(&mut self, lines: &mut Vec<String>, depth: usize) {
        self.feat.recursion = true;
        let g = self.fresh(); // function name; deliberately not added to scope
        let base = self.rng.gen_range(0..=3);
        let op = ["+", "-", "*"][self.rng.gen_range(0..3)];
        let ind = Self::indent(depth);
        let ind1 = Self::indent(depth + 1);
        lines.push(format!("{ind}local function {g}(n)"));
        lines.push(format!("{ind1}if n <= 0 then return {base} end"));
        lines.push(format!("{ind1}return (n {op} {g}((n) - 1))"));
        lines.push(format!("{ind}end"));

        let bound = self.cfg.loop_bound.max(2);
        let arg = self.anchor_num();
        let res = self.fresh();
        lines.push(format!(
            "{ind}local {res} = {g}((math.abs({arg}) % {bound}))"
        ));
        self.tainted.insert(res.clone());
        self.scope.push((res, Type::Number));
    }

    /// Emit a closure that captures a Number local declared earlier (a long-range
    /// binding: the captured variable lives outside the closure body). Bind its
    /// result to a tainted Number local.
    fn gen_closure_helper(&mut self, lines: &mut Vec<String>, depth: usize) {
        let captured = self.anchor_num(); // a tainted Number to close over
        let c = self.fresh(); // closure name; not added to the value scope
        let op = ["+", "-", "*"][self.rng.gen_range(0..3)];
        let ind = Self::indent(depth);
        let ind1 = Self::indent(depth + 1);
        lines.push(format!("{ind}local function {c}(y)"));
        lines.push(format!("{ind1}return ({captured} {op} y)"));
        lines.push(format!("{ind}end"));

        self.feat.closure = true;
        let arg = self.gen_num(self.cfg.expr_fuel.min(3));
        let res = self.fresh();
        let tainted = self.refs_tainted(&arg); // captured is tainted anyway
        lines.push(format!("{ind}local {res} = {c}({arg})"));
        if tainted || self.tainted.contains(&captured) {
            self.tainted.insert(res.clone());
        }
        self.scope.push((res, Type::Number));
    }

    /// Build a table where each element depends on the previous one (table
    /// coherence). Threads a def-use chain through the table and binds it as a
    /// tainted List.
    fn gen_table_build(&mut self, lines: &mut Vec<String>, depth: usize) {
        self.feat.table_build = true;
        let t = self.fresh();
        let ind = Self::indent(depth);
        let ind1 = Self::indent(depth + 1);
        let seed = self.anchor_num(); // first element input-dependent
        lines.push(format!("{ind}local {t} = {{{seed}}}"));

        let n = self.rng.gen_range(2..=self.cfg.loop_bound.max(2));
        let i = self.fresh();
        let op = ["+", "-", "*"][self.rng.gen_range(0..3)];
        let step = self.rng.gen_range(1..=4);
        lines.push(format!("{ind}for {i} = 2, {n} do"));
        // Each element is a function of its predecessor: t[i] = (t[i-1] op step).
        lines.push(format!("{ind1}{t}[{i}] = ({t}[{i} - 1] {op} {step})"));
        lines.push(format!("{ind}end"));

        self.tainted.insert(t.clone());
        self.scope.push((t, Type::List));
    }

    // ---- expressions (type-safe) ----

    fn pick(&mut self, opts: &[String]) -> Option<String> {
        if opts.is_empty() {
            None
        } else {
            let idx = self.rng.gen_range(0..opts.len());
            Some(opts[idx].clone())
        }
    }

    fn gen_expr(&mut self, ty: Type, fuel: usize) -> String {
        match ty {
            Type::Number => self.gen_num(fuel),
            Type::Bool => self.gen_bool(fuel),
            Type::List => self.gen_list(fuel),
            Type::Str => self.gen_str(fuel),
        }
    }

    fn gen_num(&mut self, fuel: usize) -> String {
        if fuel == 0 || self.rng.gen_bool(0.35) {
            // terminal
            let vars = self.vars_of(Type::Number);
            if !vars.is_empty() && self.rng.gen_bool(0.7) {
                return self.pick(&vars).unwrap();
            }
            return format!("{}", self.rng.gen_range(-9..=9));
        }
        match self.rng.gen_range(0..7) {
            0 => {
                let op = ["+", "-", "*"][self.rng.gen_range(0..3)];
                format!("({} {} {})", self.gen_num(fuel - 1), op, self.gen_num(fuel - 1))
            }
            // Parenthesize the operand: `-` on a negative literal would emit
            // `--4`, which Lua lexes as a comment.
            1 => format!("(-({}))", self.gen_num(fuel - 1)),
            2 => {
                // division by a nonzero literal keeps it total
                let d = self.nonzero();
                format!("({} // {})", self.gen_num(fuel - 1), d)
            }
            3 => {
                let d = self.nonzero();
                format!("({} % {})", self.gen_num(fuel - 1), d)
            }
            4 => {
                // length or safe index of a list, if available
                let lists = self.vars_of(Type::List);
                if let Some(l) = self.pick(&lists) {
                    if self.rng.gen_bool(0.5) {
                        format!("#{}", l)
                    } else {
                        // wrap index into 1..#l (l is always non-empty by construction)
                        let idx = self.gen_num(fuel - 1);
                        format!("{}[(math.abs({}) % #{}) + 1]", l, idx, l)
                    }
                } else {
                    self.gen_num(fuel - 1)
                }
            }
            5 => {
                // string length, a cross-type dependency string -> number
                let strs = self.vars_of(Type::Str);
                if let Some(s) = self.pick(&strs) {
                    format!("#{}", s)
                } else {
                    self.gen_num(fuel - 1)
                }
            }
            _ => {
                let fns = ["math.abs", "math.floor", "math.ceil"];
                let f = fns[self.rng.gen_range(0..fns.len())];
                format!("{}({})", f, self.gen_num(fuel - 1))
            }
        }
    }

    fn gen_str(&mut self, fuel: usize) -> String {
        if fuel == 0 || self.rng.gen_bool(0.4) {
            let vars = self.vars_of(Type::Str);
            if !vars.is_empty() && self.rng.gen_bool(0.6) {
                return self.pick(&vars).unwrap();
            }
            return self.str_literal();
        }
        match self.rng.gen_range(0..4) {
            0 => format!("({} .. {})", self.gen_str(fuel - 1), self.gen_str(fuel - 1)),
            1 => {
                // string.sub is total: out-of-range indices yield "".
                let a = self.rng.gen_range(1..=4);
                let b = self.rng.gen_range(a..=a + 4);
                format!("string.sub({}, {}, {})", self.gen_str(fuel - 1), a, b)
            }
            2 => {
                // bounded repetition to avoid string blow-up
                let k = self.rng.gen_range(0..=3);
                format!("string.rep({}, {})", self.gen_str(fuel - 1), k)
            }
            _ => {
                let f = ["string.upper", "string.lower", "string.reverse"]
                    [self.rng.gen_range(0..3)];
                format!("{}({})", f, self.gen_str(fuel - 1))
            }
        }
    }

    fn str_literal(&mut self) -> String {
        // Fixed small pool: keeps the token-level string vocabulary tiny and
        // closed (no unseen strings at eval, which would break reconstruction).
        const POOL: &[&str] = &[
            "a", "b", "ab", "abc", "x", "y", "hi", "ok", "lua", "foo", "bar",
            "baz", "key", "val", "one", "two",
        ];
        format!("\"{}\"", POOL[self.rng.gen_range(0..POOL.len())])
    }

    fn gen_bool(&mut self, fuel: usize) -> String {
        if fuel == 0 || self.rng.gen_bool(0.3) {
            let vars = self.vars_of(Type::Bool);
            if !vars.is_empty() && self.rng.gen_bool(0.5) {
                return self.pick(&vars).unwrap();
            }
            return if self.rng.gen_bool(0.5) { "true".into() } else { "false".into() };
        }
        match self.rng.gen_range(0..3) {
            0 => {
                let op = ["==", "~=", "<", "<=", ">", ">="][self.rng.gen_range(0..6)];
                format!("({} {} {})", self.gen_num(fuel - 1), op, self.gen_num(fuel - 1))
            }
            1 => {
                let op = ["and", "or"][self.rng.gen_range(0..2)];
                format!("({} {} {})", self.gen_bool(fuel - 1), op, self.gen_bool(fuel - 1))
            }
            _ => format!("(not {})", self.gen_bool(fuel - 1)),
        }
    }

    fn gen_list(&mut self, fuel: usize) -> String {
        let vars = self.vars_of(Type::List);
        if !vars.is_empty() && self.rng.gen_bool(0.5) {
            return self.pick(&vars).unwrap();
        }
        // Non-empty literal so that indexing/# are always safe.
        let len = 1 + self.rng.gen_range(0..self.cfg.list_len.max(1));
        let elems: Vec<String> = (0..len)
            .map(|_| self.gen_num(fuel.saturating_sub(1).min(2)))
            .collect();
        format!("{{{}}}", elems.join(", "))
    }

    fn nonzero(&mut self) -> i64 {
        let mut d = self.rng.gen_range(-9..=9);
        if d == 0 {
            d = 1;
        }
        d
    }
}