atlas_audit.py

#!/usr/bin/env python3
# Copyright (c) 2026 David Roe. All rights reserved.
# Released under Apache 2.0 license as described in the file LICENSE.
# Authors: David Roe, roed@mit.edu, using Claude Opus-4.8 and Fable-5
"""Render `atlas-audit.md` from a Lean Atlas graph export.

Usage:
    python3 scripts/atlas_audit.py [GRAPH_JSON] [TARGET] [OUTPUT]

Defaults:
    GRAPH_JSON = atlas-graph.json         (produced by `lake exe atlas graph-data -o ...`)
    TARGET     = GQ2.SectionTen.main_surjection_count'
    OUTPUT     = atlas-audit.md

Emits two sections:
  1. the Lean Compass review cone (statement-level nodes whose *semantic* correctness can
     affect TARGET — the full closure minus the theorem-proof value edges the type checker
     already guarantees);
  2. the custom axioms reported by Lean's `#print axioms` command.

The second section deliberately does not infer the trust base from the Atlas graph. Atlas
omits internal names such as private helper declarations; an axiom used only through such a
helper is therefore absent from the graph closure even though it remains in the kernel trust
base of TARGET.

Every reference is a repo-relative Markdown link `path#Lstart-Lend`.  See docs/atlas.md.
"""
import json
import re
import subprocess
import sys
import tempfile
from collections import defaultdict
from pathlib import Path

GRAPH = sys.argv[1] if len(sys.argv) > 1 else "atlas-graph.json"
TARGET = sys.argv[2] if len(sys.argv) > 2 else "GQ2.SectionTen.main_surjection_count'"
OUTPUT = sys.argv[3] if len(sys.argv) > 3 else "atlas-audit.md"

# Compass prunes value dependencies OUT OF THEOREMS (guaranteed correct by the type checker).
PRUNE = {"theorem_value_to_definition", "theorem_value_to_theorem"}
STANDARD_AXIOMS = {"propext", "Classical.choice", "Quot.sound"}

d = json.loads(Path(GRAPH).read_text(encoding="utf-8"))
nodes = {n["id"]: n for n in d["nodes"]}
if TARGET not in nodes:
    sys.exit(f"target {TARGET!r} not found in {GRAPH}")

full, comp = defaultdict(list), defaultdict(list)
for e in d["edges"]:
    full[e["source"]].append(e["target"])
    if e["kind"] not in PRUNE:
        comp[e["source"]].append(e["target"])


def reachable(start, adj):
    seen, stack = set(), [start]
    while stack:
        c = stack.pop()
        if c in seen:
            continue
        seen.add(c)
        for t in adj.get(c, ()):
            if t not in seen:
                stack.append(t)
    return seen


def link(nid, label=None):
    n = nodes[nid]
    a, b = n.get("lineStart"), n.get("lineEnd")
    anchor = f"#L{a}" if a == b else f"#L{a}-L{b}"
    return f"[`{label or nid}`]({n['filePath']}{anchor})"


def lines(nid):
    """Visible line marker (en-dash range), e.g. `L58` or `L58–60`."""
    n = nodes[nid]
    a, b = n.get("lineStart"), n.get("lineEnd")
    return f"L{a}" if a == b else f"L{a}–{b}"


def by_line(nid):
    n = nodes[nid]
    return (n.get("lineStart") or 0, n.get("lineEnd") or 0, nid)


def kernel_axioms(target):
    """Ask Lean for TARGET's authoritative axiom dependency set."""
    with tempfile.NamedTemporaryFile("w", suffix=".lean", delete=False) as source:
        source.write(f"import GQ2\n#print axioms {target}\n")
        source_path = Path(source.name)
    try:
        result = subprocess.run(
            ["lake", "env", "lean", str(source_path)],
            check=True,
            capture_output=True,
            text=True,
        )
    finally:
        source_path.unlink(missing_ok=True)
    match = re.search(r"depends on axioms:\s*\[(.*?)\]", result.stdout, re.DOTALL)
    if match is None:
        raise SystemExit(f"could not parse `#print axioms` output for {target!r}")
    return sorted(
        name.strip() for name in match.group(1).split(",")
        if name.strip() and name.strip() not in STANDARD_AXIOMS
    )


fc = reachable(TARGET, full)
cc = reachable(TARGET, comp)
axioms = kernel_axioms(TARGET)
missing_axiom_nodes = [name for name in axioms if name not in nodes]
if missing_axiom_nodes:
    raise SystemExit(
        "kernel-reported axioms missing from Atlas graph: " + ", ".join(missing_axiom_nodes)
    )

out = []
w = out.append
w(f"# Lean Atlas — audit of `{TARGET}`\n")
w(f"*Generated by `scripts/atlas_audit.py` from the {len(nodes)}-node / {len(d['edges'])}-edge "
  f"dependency graph (`lake exe atlas graph-data`). Census: {d['statistics']['axiomCount']} "
  f"axioms; {d['statistics']['sorryCount']} live sorries. Line numbers are a snapshot — "
  f"regenerate after source edits (see docs/atlas.md).*\n")

w("## 1. Lean Compass review cone\n")
w(f"**Full closure: {len(fc)} nodes → Compass cone: {len(cc)} nodes "
  f"({(len(fc) - len(cc)) / len(fc) * 100:.1f}% reduction).** These {len(cc)} are the "
  "statement-level nodes whose *semantic* correctness can affect the target and are therefore "
  "the recommended human-review set. Everything else in the full closure is proof-level "
  "(type-checker-guaranteed) and pruned.\n")
byfile = defaultdict(list)
for i in cc:
    byfile[nodes[i]["filePath"]].append(i)
for fp in sorted(byfile):
    w(f"### [`{fp}`]({fp})")
    for i in sorted(byfile[fp], key=by_line):  # source order → consecutive runs read as a block
        tag = "thm" if nodes[i]["kind"] == "theorem" else "def"
        w(f"- `{lines(i)}` **[{tag}]** {link(i)}")
    w("")

w(f"## 2. Axioms the proof rests on ({len(axioms)})\n")
w("Lean's kernel-reported trust base (`#print axioms`, minus the "
  "`propext`/`Classical.choice`/`Quot.sound` core). This is queried separately because Lean "
  "Atlas intentionally omits private/internal helper declarations from its graph; deriving "
  "the trust base from that graph can therefore miss axioms reached through private proofs.\n")
for a in sorted(axioms, key=by_line):
    w(f"- `{lines(a)}` {link(a)}")
w("")

Path(OUTPUT).write_text("\n".join(out), encoding="utf-8")
print(f"wrote {OUTPUT}: review cone {len(cc)}/{len(fc)} nodes, {len(axioms)} axioms in trust base")