From 5ef2b5c11840890471db123ecce3f26642e38478 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Mon, 9 Dec 2024 19:40:00 -0800
Subject: [PATCH] feat: Collect newly defined constants

---
 Pantograph/Frontend/Elab.lean | 15 ++++++++++++++-
 Test/Frontend.lean            | 27 +++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/Pantograph/Frontend/Elab.lean b/Pantograph/Frontend/Elab.lean
index c4704fc..4ace608 100644
--- a/Pantograph/Frontend/Elab.lean
+++ b/Pantograph/Frontend/Elab.lean
@@ -2,6 +2,7 @@
 import Lean.Elab.Import
 import Lean.Elab.Command
 import Lean.Elab.InfoTree
+import Lean.DeclarationRange
 
 import Pantograph.Frontend.Basic
 import Pantograph.Frontend.MetaTranslate
@@ -128,7 +129,7 @@ Since we cannot directly merge `MetavarContext`s, we have to get creative. This
 function duplicates frozen mvars in term and tactic info nodes, and add them to
 the current `MetavarContext`.
 -/
-@[export pantograph_frontend_sorrys_to_goal_state]
+@[export pantograph_frontend_sorrys_to_goal_state_m]
 def sorrysToGoalState (sorrys : List InfoWithContext) : MetaM GoalState := do
   assert! !sorrys.isEmpty
   let goalsM := sorrys.mapM λ i => do
@@ -147,5 +148,17 @@ def sorrysToGoalState (sorrys : List InfoWithContext) : MetaM GoalState := do
   GoalState.createFromMVars goals root
 
 
+@[export pantograph_frontend_collect_new_defined_constants_m]
+def collectNewDefinedConstants (step : CompilationStep) : IO (List Name) := do
+  step.after.constants.map₂.foldlM (λ acc name _ => do
+    if step.before.contains name then
+      return acc
+    let coreM : CoreM Bool := Option.isSome <$> findDeclarationRanges? name
+    let hasRange ← coreM.run' { fileName := step.fileName, fileMap := step.fileMap } { env := step.after } |>.toBaseIO
+    match hasRange with
+    | .ok true => return name :: acc
+    | .ok false => return acc
+    | .error e => throw $ IO.userError (← e.toMessageData.toString)
+    ) []
 
 end Pantograph.Frontend
diff --git a/Test/Frontend.lean b/Test/Frontend.lean
index 3b765fd..68aaf94 100644
--- a/Test/Frontend.lean
+++ b/Test/Frontend.lean
@@ -184,6 +184,31 @@ def mystery : Nat := true
   let goalStates ← (collectSorrysFromSource input).run' {}
   let [goalState] := goalStates | panic! s!"Incorrect number of states: {goalStates.length}"
 
+def collectNewConstants (source: String) : MetaM (List (List Name)) := do
+  let filename := "<anonymous>"
+  let (context, state) ← do Frontend.createContextStateFromFile source filename (← getEnv) {}
+  let m := Frontend.mapCompilationSteps λ step => do
+    Frontend.collectNewDefinedConstants step
+  m.run context |>.run' state
+
+def test_collect_one_constant : TestT MetaM Unit := do
+  let input := "
+def mystery : Nat := 123
+  "
+  let names ← collectNewConstants input
+  checkEq "constants" names [[`mystery]]
+def test_collect_one_theorem : TestT MetaM Unit := do
+  let input := "
+theorem mystery [SizeOf α] (as : List α) (i : Fin as.length) : sizeOf (as.get i) < sizeOf as := by
+  match as, i with
+  | a::as, ⟨0, _⟩  => simp_arith [get]
+  | a::as, ⟨i+1, h⟩ =>
+    have ih := sizeOf_get as ⟨i, Nat.le_of_succ_le_succ h⟩
+    apply Nat.lt_trans ih
+    simp_arith
+  "
+  let names ← collectNewConstants input
+  checkEq "constants" names [[`mystery]]
 
 def suite (env : Environment): List (String × IO LSpec.TestSeq) :=
   let tests := [
@@ -193,6 +218,8 @@ def suite (env : Environment): List (String × IO LSpec.TestSeq) :=
     ("sorry_in_coupled", test_sorry_in_coupled),
     ("environment_capture", test_environment_capture),
     ("capture_type_mismatch", test_capture_type_mismatch),
+    ("collect_one_constant", test_collect_one_constant),
+    ("collect_one_theorem", test_collect_one_theorem),
   ]
   tests.map (fun (name, test) => (name, runMetaMSeq env $ runTest test))