"""Lemma-invention training data ("subgoal invention"). For a hard synthetic TRAINING problem H1 -> ... -> Hn -> C (automation fails on it directly), search for an intermediate statement M such that (1) Z3: H |= M and H, M |= C; (2) Rocq: the automation closers prove H -> M and H -> M -> C. Then the goal after `intros` is paired with the tactic `assert (M)`: the model learns to propose the statement that splits a hard goal into two easy ones. Candidates for M include the generalization of the conclusion over a constant ("generalize, then specialize"), subformulas of the hypotheses and the conclusion, and their instances. Rocq checks the whole cut proof end to end. Only training problems are used (the ids in the given pool files); the test set and the held-out tasks never are. Usage: python gen_cuts.py out.jsonl data/synth_train_hard.jsonl [--workers 6] """ import json, sys, threading from pathlib import Path import z3 import gen_problems as G from rocq_env import Coqtop, Goal import htps HERE = Path(__file__).resolve().parent Z3_LOCK = threading.Lock() # z3's Python API shares one context: not thread-safe def z3c(fs, timeout=500): with Z3_LOCK: return G.z3_check(fs, timeout=timeout) def subformulas(f): yield f if f[0] in ("not",): yield from subformulas(f[1]) elif f[0] in ("and", "or", "imp"): yield from subformulas(f[1]); yield from subformulas(f[2]) elif f[0] in ("all", "ex"): yield from subformulas(f[2]) def free_vars(f, bound=()): op = f[0] if op in ("all", "ex"): return free_vars(f[2], bound + (f[1],)) if op in ("not",): return free_vars(f[1], bound) if op in ("and", "or", "imp"): return free_vars(f[1], bound) | free_vars(f[2], bound) return {a for a in f[1:] if isinstance(a, str) and a not in G.CONSTS and a not in bound and a not in G.PROPS} def subst(f, v, c): op = f[0] if op in ("all", "ex"): return f if f[1] == v else (op, f[1], subst(f[2], v, c)) if op in ("not", "and", "or", "imp"): return (op,) + tuple(subst(g, v, c) for g in f[1:]) return (op,) + tuple(c if a == v else a for a in f[1:]) def mentions_const(f, c): return any(mentions_const(g, c) if isinstance(g, tuple) else g == c for g in f[1:]) def generalize(f, c, v="x"): """forall v, f[c := v] (only if f mentions c and v is unused).""" if not mentions_const(f, c) or G.mentions(f, v): return None def rep(g): if g[0] in ("all", "ex"): return (g[0], g[1], rep(g[2])) if g[0] in ("not", "and", "or", "imp"): return (g[0],) + tuple(rep(h) for h in g[1:]) return (g[0],) + tuple(v if a == c else a for a in g[1:]) return ("all", v, rep(f)) def candidates(hyps, concl): out = [] for c in G.CONSTS: # generalize the conclusion g = generalize(concl, c) if g: out.append(g) for f in list(hyps) + [concl]: for s in subformulas(f): fv = free_vars(s) if not fv: out.append(s) elif len(fv) == 1: v = next(iter(fv)) for c in G.CONSTS: out.append(subst(s, v, c)) seen, res = set(), [] for m in out: k = G.rocq(m) if k in seen or m in hyps or m == concl or G.rocq(m) == G.rocq(concl): continue seen.add(k); res.append(m) return res def main(): out_path, pool = sys.argv[1], sys.argv[2] workers = int(sys.argv[sys.argv.index("--workers") + 1]) if "--workers" in sys.argv else 6 test_ids = {json.loads(l)["id"] for f in ("synth_test_hard.jsonl", "synth_test_easy.jsonl") for l in open(HERE / "data" / f)} want = {json.loads(l)["id"] for l in open(pool)} - test_ids cache = HERE / "data" / "trees_pool_a.json" if not cache.exists(): # regenerate the formula trees of pool_a once (same seeds, same order) allt = {} for seed in (11, 12, 13, 14): for pid, hyps, concl, st, _ in G.iter_problems(1500, seed): allt[pid] = [hyps, concl, st] json.dump(allt, open(cache, "w")) def tup(x): return tuple(tup(y) for y in x) if isinstance(x, list) else x trees = {pid: ([tup(h) for h in v[0]], tup(v[1]), v[2]) for pid, v in json.load(open(cache)).items() if pid in want} todo = list(trees.items()) print("problems", len(todo), flush=True) lock = threading.Lock() fo = open(out_path, "w") stats = {"done": 0, "cuts": 0} closers = lambda g: ["intros; " + t for t in htps.closers_for(g, list(G.DEFS))] def solved(coq, g): return any(coq.apply(g, t) == ("ok", []) for t in closers(g)) coqs = [Coqtop(G.PRELUDE, G.AMBIENT, tac_timeout=2) for _ in range(workers)] # spawn before threads def worker(coq): while True: with lock: if not todo: return pid, (hyps, concl, st) = todo.pop() found = None for m in candidates(hyps, concl)[:60]: if z3c(hyps + [("not", m)]) != z3.unsat: continue if z3c(hyps + [m, ("not", concl)]) != z3.unsat: continue g1 = Goal([], G.statement(hyps, m)) g2 = Goal([], G.statement(hyps + [m], concl)) if solved(coq, g1) and solved(coq, g2): found = m; break rec = None if found is not None: root = Goal([], st) status, kids = coq.apply(root, "intros") if status == "ok" and len(kids) == 1: after = kids[0] tac = f"assert ({G.rocq(found)})" st2, parts = coq.apply(after, tac) # the cut must really split the goal into two automation-easy goals if st2 == "ok" and len(parts) == 2 and all(solved(coq, p) for p in parts): rec = {"id": pid, "goal": after.key, "tactic": tac, "kind": "generalize" if found[0] == "all" and G.rocq(found).count("forall") > G.rocq(concl).count("forall") else "cut"} with lock: stats["done"] += 1 if rec: stats["cuts"] += 1 fo.write(json.dumps(rec) + "\n"); fo.flush() if stats["done"] % 10 == 0: print(stats, flush=True) ts = [threading.Thread(target=worker, args=(c,)) for c in coqs] [t.start() for t in ts]; [t.join() for t in ts] print("final", stats, flush=True) if __name__ == "__main__": main()