release_tools/gen_hints.py

#!/usr/bin/env python3
"""gen_hints.py — static structural-hints JSON for the open TRUE problems.

Precomputes, per open true-implication problem, the structural context that a
proof search / LLM stage wants: constancy classes (variables that leave a side
of the hypothesis invariant), top-k promising h-instantiations, and equation
structure notes. Static file = zero coupling to any live search interface;
congr_kit / the-prover's meet-in-the-middle search / the LLM tail can all seed
from it.

Reuses the opnorm reference solver's analysis (build_constancy_info,
compute_h_instantiations, analyze_equation_structure, compute_match_collapse_hints,
compute_equation_analysis) so the hint shapes match what congr_kit already
consumes — this just materializes them offline for the open set.

Usage:
    python3 gen_hints.py <repo> <open_true_ids.json> <out.json>
      open_true_ids.json: {"hard2":[ids...], "hard3":[ids...], ...} or a flat [ids]
"""
from __future__ import annotations

import importlib.util
import json
import sys
from pathlib import Path


def load_opnorm(repo: Path):
    p = repo / "examples" / "solo" / "demos" / "opnorm" / "solver.py"
    spec = importlib.util.spec_from_file_location("opnorm_solver", p)
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return m


def load_problems(repo: Path) -> dict:
    problems = {}
    for name in ("normal", "hard1", "hard2", "hard3"):
        for line in (repo / "examples" / "problems" / f"{name}.jsonl").read_text().splitlines():
            line = line.strip()
            if line:
                r = json.loads(line)
                problems[r["id"]] = r
    return problems


def _parse(text: str):
    """Parse one side of a ◇-equation into an AST: str = var, ('op', l, r)."""
    text = text.replace("*", "◇")

    def go(s):
        s = s.strip()
        while len(s) >= 2 and s[0] == "(" and s[-1] == ")":
            depth = 0
            ok = True
            for i, c in enumerate(s):
                depth += (c == "(") - (c == ")")
                if depth == 0 and i < len(s) - 1:
                    ok = False
                    break
            if ok:
                s = s[1:-1].strip()
            else:
                break
        depth = 0
        last = -1
        for i, c in enumerate(s):
            depth += (c == "(") - (c == ")")
            if depth == 0 and c == "◇":
                last = i
        if last >= 0:
            return ("op", go(s[:last]), go(s[last + 1:]))
        return s

    return go(text)


def _render(t) -> str:
    if isinstance(t, str):
        return t
    return f"({_render(t[1])} ◇ {_render(t[2])})"


def deepest_disagreement(lhs_text: str, rhs_text: str):
    """L/R paths (from root, ''=root; L=left arg, R=right arg — congr-kit
    addressing) to the DEEPEST positions where eq2's two sides structurally
    diverge, with the disagreeing subterms. These are the priority seeds: the
    const/rewrite move that lands usually rewrites the deepest mismatch."""
    la, ra = _parse(lhs_text), _parse(rhs_text)
    out = []

    def rec(a, b, path):
        if a == b:
            return
        if isinstance(a, tuple) and isinstance(b, tuple):
            rec(a[1], b[1], path + "L")
            rec(a[2], b[2], path + "R")
        else:
            out.append((path, a, b))

    rec(la, ra, "")
    if not out:
        return []
    maxd = max(len(p) for p, _, _ in out)
    return [{"path": (p or "(root)"), "lhs_subterm": _render(a), "rhs_subterm": _render(b)}
            for p, a, b in out if len(p) == maxd]


def _clean(obj):
    """Make opnorm's structures JSON-serializable (sets -> sorted lists)."""
    if isinstance(obj, set):
        return sorted(obj)
    if isinstance(obj, dict):
        return {k: _clean(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_clean(x) for x in obj]
    return obj


def hints_for(op, problem: dict) -> dict:
    e1 = op.normalize_op_to_diamond(problem["equation1"])
    e2 = op.normalize_op_to_diamond(problem["equation2"])
    v1 = op.parse_variables(e1)
    v2 = op.parse_variables(e2)
    out = {"eq1": e1, "eq2": e2, "eq1_vars": v1, "eq2_vars": v2}
    try:
        cinfo, lhs_only, rhs_only = op.build_constancy_info(e1, v1, v2)
        # keep the useful, serializable parts of each constancy lemma
        out["constancy"] = [{
            "quant_vars": ci.get("quant_vars"),
            "lhs_template": ci.get("lhs_template"),
            "rhs_template": ci.get("rhs_template"),
            "have_line": ci.get("have_line"),
        } for ci in cinfo]
        out["lhs_only_vars"] = sorted(lhs_only)
        out["rhs_only_vars"] = sorted(rhs_only)
    except Exception as e:  # noqa: BLE001
        out["constancy_error"] = f"{type(e).__name__}: {e}"
    # priority seed: deepest structural mismatch between eq2's two sides
    try:
        lhs2, rhs2 = e2.split("=", 1)
        out["deepest_disagreement"] = deepest_disagreement(lhs2, rhs2)
    except Exception as e:  # noqa: BLE001
        out["deepest_disagreement"] = f"ERROR {type(e).__name__}: {e}"
    for key, fn, argstyle in (
        ("h_instantiations", op.compute_h_instantiations, "vars"),
        ("structure_notes", op.analyze_equation_structure, "texts"),
        ("match_collapse", op.compute_match_collapse_hints, "texts"),
        ("equation_analysis", op.compute_equation_analysis, "texts"),
    ):
        try:
            out[key] = _clean(fn(e1, v1, v2) if argstyle == "vars" else fn(e1, e2))
        except Exception as e:  # noqa: BLE001
            out[key] = f"ERROR {type(e).__name__}: {e}"
    return out


def main() -> None:
    if len(sys.argv) != 4:
        sys.exit(__doc__)
    repo = Path(sys.argv[1]).resolve()
    ids_spec = json.loads(Path(sys.argv[2]).read_text())
    out_path = Path(sys.argv[3])
    if isinstance(ids_spec, dict):
        ids = [i for tier in ("hard2", "hard3", "hard1", "normal")
               for i in ids_spec.get(tier, [])]
    else:
        ids = list(ids_spec)

    op = load_opnorm(repo)
    problems = load_problems(repo)
    hints = {}
    for n, pid in enumerate(ids, 1):
        if pid not in problems:
            print(f"skip unknown id {pid}", file=sys.stderr)
            continue
        hints[pid] = hints_for(op, problems[pid])
        if n % 25 == 0:
            print(f"  ..{n}/{len(ids)}", file=sys.stderr)
    out_path.write_text(json.dumps(hints, indent=1, ensure_ascii=False))
    print(f"wrote {out_path}: hints for {len(hints)} problems", file=sys.stderr)


if __name__ == "__main__":
    main()