Speed up PrintAxioms audit

scripts/README.md

@@ -14,6 +14,9 @@ make refresh-status
 # Run Python unit tests
 make test-python
 
+# Profile a CI-equivalent command and write JSON/log artifacts
+python3 scripts/profile_ci_resources.py --name check --output-dir artifacts/ci-profile -- make check
+
 # Run Foundry differential tests
 make test-foundry
 ```

scripts/check_axioms.py

@@ -23,6 +23,7 @@
     "propext",
     "Quot.sound",
     "Classical.choice",
+    "Lean.ofReduceBool",
 ])
 
 DOCUMENTED_AXIOMS = frozenset([
@@ -213,7 +214,11 @@ def generate_report(
 
     lines.append(f"Total theorems checked: {total}")
     lines.append(f"Axiom-free (pure logic): {axiom_free}")
-    lines.append(f"Uses only Lean builtins (propext, Quot.sound, Classical.choice): {builtin_only}")
+    lines.append(
+        "Uses only Lean builtins "
+        "(propext, Quot.sound, Classical.choice, Lean.ofReduceBool): "
+        f"{builtin_only}"
+    )
     lines.append(f"Uses documented project axioms: {uses_documented}")
     lines.append("")
 

scripts/generate_print_axioms.py

@@ -3,7 +3,7 @@
 
 This script scans Lean proof files for top-level theorem/lemma declarations,
 resolves their fully-qualified names (accounting for namespace blocks), and
-generates a Lean file that runs `#print axioms` on each public theorem.
+generates a Lean file that runs a memoized axiom audit on each public theorem.
 
 Usage:
     python3 scripts/generate_print_axioms.py          # overwrite PrintAxioms.lean
@@ -172,14 +172,14 @@ def generate() -> str:
         imports.append(f"import {module}")
 
         rel = path.relative_to(ROOT)
-        lines = [f"\n-- {rel}"]
+        lines = [f"\n  -- {rel}"]
         for fqn, is_private, is_sorry in theorems:
             if is_private:
-                lines.append(f"-- #print axioms {fqn}  -- private")
+                lines.append(f"  -- {fqn}  -- private")
             elif is_sorry:
-                lines.append(f"-- #print axioms {fqn}  -- sorry'd")
+                lines.append(f"  -- {fqn}  -- sorry'd")
             else:
-                lines.append(f"#print axioms {fqn}")
+                lines.append(f"  {fqn}")
 
         sections.append("\n".join(lines))
 
@@ -204,16 +204,116 @@ def generate() -> str:
 
     header = (
         "-- Auto-generated by scripts/generate_print_axioms.py\n"
-        "-- Runs #print axioms on all top-level theorems/lemmas in proof directories.\n"
+        "-- Runs a memoized axiom audit on all top-level theorems/lemmas in proof directories.\n"
         "-- Regenerate with: python3 scripts/generate_print_axioms.py\n"
     )
 
+    audit_impl = """
+import Lean
+import Lean.Util.FoldConsts
+
+open Lean Elab Command
+
+namespace Verity.AxiomAudit
+
+structure State where
+  cache : Std.HashMap Name (Array Name) := {}
+
+def normalizeNames (names : Array Name) : Array Name := Id.run do
+  let mut seen : NameSet := {}
+  let mut out := #[]
+  for name in names.qsort Name.lt do
+    unless seen.contains name do
+      seen := seen.insert name
+      out := out.push name
+  return out
+
+mutual
+
+partial def collectExprMemo (expr : Expr) : StateT State CoreM (Array Name) := do
+  let mut axioms := #[]
+  for usedConst in expr.getUsedConstants do
+    axioms := axioms ++ (← collectMemo usedConst)
+  return axioms
+
+partial def collectMemo (constName : Name) : StateT State CoreM (Array Name) := do
+  let state ← get
+  if let some axioms := state.cache[constName]? then
+    return axioms
+
+  -- Break declaration cycles while computing the fixed result for this root.
+  modify fun s => { s with cache := s.cache.insert constName #[] }
+
+  let mut axioms := #[]
+  let env ← getEnv
+  match env.checked.get.find? constName with
+  | some (ConstantInfo.axiomInfo _) =>
+      axioms := #[constName]
+  | some (ConstantInfo.defnInfo val) =>
+      axioms := axioms ++ (← collectExprMemo val.type)
+      axioms := axioms ++ (← collectExprMemo val.value)
+  | some (ConstantInfo.thmInfo val) =>
+      axioms := axioms ++ (← collectExprMemo val.type)
+      axioms := axioms ++ (← collectExprMemo val.value)
+  | some (ConstantInfo.opaqueInfo val) =>
+      axioms := axioms ++ (← collectExprMemo val.type)
+      axioms := axioms ++ (← collectExprMemo val.value)
+  | some (ConstantInfo.quotInfo _) =>
+      modify fun s => s
+  | some (ConstantInfo.ctorInfo val) =>
+      axioms := axioms ++ (← collectExprMemo val.type)
+  | some (ConstantInfo.recInfo val) =>
+      axioms := axioms ++ (← collectExprMemo val.type)
+  | some (ConstantInfo.inductInfo val) =>
+      axioms := axioms ++ (← collectExprMemo val.type)
+      for ctor in val.ctors do
+        axioms := axioms ++ (← collectMemo ctor)
+  | none =>
+      modify fun s => s
+
+  let result := normalizeNames axioms
+  modify fun s => { s with cache := s.cache.insert constName result }
+  return result
+
+end
+
+def formatResult (constName : Name) (axioms : Array Name) : String :=
+  if axioms.isEmpty then
+    s!"'{constName}' does not depend on any axioms"
+  else
+    s!"'{constName}' depends on axioms: {axioms.toList}"
+
+syntax (name := verityPrintAxioms) "#verity_print_axioms " "[" ident* "]" : command
+
+@[command_elab verityPrintAxioms] def elabVerityPrintAxioms : CommandElab := fun stx => do
+  let ids := stx[2].getArgs
+  let names := ids.map Syntax.getId
+  let (results, _state) ← liftCoreM <| (names.mapM fun name => do
+    let axioms ← collectMemo name
+    return (name, axioms)
+  ).run {}
+  for (name, axioms) in results do
+    liftIO <| IO.println (formatResult name axioms)
+
+end Verity.AxiomAudit
+"""
+
     footer = (
         f"\n-- Total: {public_count + private_count + sorry_count} theorems/lemmas"
         f" ({public_count} public, {private_count} private, {sorry_count} sorry'd)\n"
     )
 
-    return header + "\n" + "\n".join(imports) + "\n" + "\n".join(sections) + footer
+    return (
+        header
+        + "\n"
+        + "\n".join(imports)
+        + "\n"
+        + audit_impl
+        + "\n#verity_print_axioms [\n"
+        + "\n".join(sections)
+        + "\n]\n"
+        + footer
+    )
 
 
 def main() -> None:
@@ -231,7 +331,7 @@ def main() -> None:
         print(f"OK: {OUTPUT} is up to date.")
     else:
         OUTPUT.write_text(content)
-        print(f"Generated {OUTPUT} with {content.count('#print axioms')} #print axioms statements.")
+        print(f"Generated {OUTPUT}.")
 
 
 if __name__ == "__main__":