space/program_generator.py

"""
Random Toy-IR list-of-dicts for GRPO / Deliverable2 / runtime_core.

Schema matches `runtime_core.SAMPLE_PROGRAM` (op, args, dest, type) — the same
shape `Deliverable2_Formatter` and `CompilerOptimizationEnv` expect.

Note: `metahack1 (1).ipynb` uses a different TAC shape (CONST/src1/STORE).
Use that notebook's generators only if you add a separate converter; this
module is self-contained for training stack compatibility.
"""

from __future__ import annotations

import copy
import json
import random
from typing import List

# Hand-written seeds (same as Colab / train.py short list) — kept in sync for regression.
_BUILTIN_EXTRA: List[List[dict]] = [
    [
        {"op": "const", "dest": "a", "args": ["1"], "type": "int"},
        {"op": "add", "dest": "b", "args": ["a", "a"], "type": "int"},
        {"op": "ret", "args": ["b"]},
    ],
    [
        {"op": "const", "dest": "x", "args": ["2"], "type": "int"},
        {"op": "const", "dest": "y", "args": ["4"], "type": "int"},
        {"op": "mul", "dest": "z", "args": ["x", "y"], "type": "int"},
        {"op": "const", "dest": "k", "args": ["1"], "type": "int"},
        {"op": "add", "dest": "w", "args": ["z", "k"], "type": "int"},
        {"op": "ret", "args": ["w"]},
    ],
]


def random_toy_ir_program(rng: random.Random) -> List[dict]:
    """
    One valid program: consts v0.., then a chain of add/mul on existing names, then ret.
    All ops use the mock-engine-friendly list schema.
    """
    n_const = rng.randint(2, 5)
    n_arith = rng.randint(1, 5)
    progs: List[dict] = []
    for i in range(n_const):
        progs.append(
            {
                "op": "const",
                "dest": f"v{i}",
                "args": [str(rng.randint(0, 20))],
                "type": "int",
            }
        )
    available = [f"v{i}" for i in range(n_const)]
    nxt = n_const
    for _j in range(n_arith):
        a = rng.choice(available)
        b = rng.choice(available)
        opn = rng.choice(["add", "mul"])
        d = f"v{nxt}"
        nxt += 1
        progs.append({"op": opn, "dest": d, "args": [a, b], "type": "int"})
        available.append(d)
    progs.append({"op": "ret", "args": [available[-1]]})
    return progs


def build_training_program_corpus(
    n_total: int = 120,
    seed: int = 42,
    *,
    include_builtins: bool = True,
) -> List[List[dict]]:
    """
    Return `n_total` programs for GRPO. Optionally prepend SAMPLE_PROGRAM + 2 hand-written IRs
    (when include_builtins), then fill with random_toy_ir_program, deduplicating by JSON key.

    Typical range: set `n_total` between 50 and 200 in the notebook.
    """
    if n_total < 1:
        raise ValueError("n_total must be >= 1")

    rng = random.Random(seed)
    out: List[List[dict]] = []
    seen: set[str] = set()

    def _add(p: List[dict]) -> None:
        k = json.dumps(p, sort_keys=True)
        if k in seen:
            return
        seen.add(k)
        out.append(copy.deepcopy(p))

    if include_builtins:
        from runtime_core import SAMPLE_PROGRAM

        for p in (SAMPLE_PROGRAM, *_BUILTIN_EXTRA):
            if len(out) >= n_total:
                break
            _add(p)

    # Fill with random programs (dedupe by full JSON; allow dup if generator keeps colliding)
    guard = 0
    while len(out) < n_total:
        guard += 1
        if guard > 200_000:
            out.append(random_toy_ir_program(rng))
            continue
        cand = random_toy_ir_program(rng)
        k = json.dumps(cand, sort_keys=True)
        if k in seen:
            continue
        seen.add(k)
        out.append(cand)

    return out