From 0b4ded10493585aeba2e9b345ddc9ed81d4a02ad Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Mon, 9 Dec 2024 20:15:53 -0800
Subject: [PATCH] fix: Collect sorrys and type mismatches

---
 Pantograph/Frontend/Elab.lean     | 21 +++++++++++++--------
 Pantograph/Frontend/InfoTree.lean | 12 ++++++------
 Test/Frontend.lean                | 11 ++++++++++-
 3 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/Pantograph/Frontend/Elab.lean b/Pantograph/Frontend/Elab.lean
index 4ace608..2eff084 100644
--- a/Pantograph/Frontend/Elab.lean
+++ b/Pantograph/Frontend/Elab.lean
@@ -103,20 +103,25 @@ structure InfoWithContext where
   context?: Option Elab.ContextInfo := .none
 
 private def collectSorrysInTree (t : Elab.InfoTree) : IO (List InfoWithContext) := do
-  let infos ← t.findAllInfoM none true fun i ctx? => match i with
+  let infos ← t.findAllInfoM none fun i ctx? => match i with
     | .ofTermInfo { expectedType?, expr, stx, lctx, .. } => do
-      let .some expectedType := expectedType? | return false
-      let .some ctx := ctx? | return false
+      let .some expectedType := expectedType? | return (false, true)
+      let .some ctx := ctx? | return (false, true)
       if expr.isSorry ∧ stx.isOfKind `Lean.Parser.Term.sorry then
-        return true
-      ctx.runMetaM lctx do
+        return (true, false)
+      let typeMatch ← ctx.runMetaM lctx do
         let type ← Meta.inferType expr
-        Bool.not <$> Meta.isExprDefEqGuarded type expectedType
+        Meta.isExprDefEqGuarded type expectedType
+      return match typeMatch, expr.hasSorry with
+      | false, true => (true, false) -- Types mismatch but has sorry -> collect, halt
+      | false, false => (true, false) -- Types mistmatch but no sorry -> collect, halt
+      | true, true => (false, true) -- Types match but has sorry -> continue
+      | true, false => (false, false) -- Types match but no sorries -> halt
     | .ofTacticInfo { stx, goalsBefore, .. } =>
       -- The `sorry` term is distinct from the `sorry` tactic
       let isSorry := stx.isOfKind `Lean.Parser.Tactic.tacticSorry
-      return isSorry ∧ !goalsBefore.isEmpty
-    | _ => return false
+      return (isSorry ∧ !goalsBefore.isEmpty, ¬ isSorry)
+    | _ => return (false, true)
   return infos.map fun (info, context?, _) => { info, context? }
 
 -- NOTE: Plural deliberately not spelled "sorries"
diff --git a/Pantograph/Frontend/InfoTree.lean b/Pantograph/Frontend/InfoTree.lean
index 50b0965..cfef621 100644
--- a/Pantograph/Frontend/InfoTree.lean
+++ b/Pantograph/Frontend/InfoTree.lean
@@ -107,18 +107,18 @@ partial def InfoTree.findAllInfo
     head ++ tail
   | _ => []
 
-/-- Monadic analogue of `findAllInfo` -/
+/-- Monadic analogue of `findAllInfo`, but predicate controls whether to recurse. -/
 partial def InfoTree.findAllInfoM [Monad m]
     (t : InfoTree)
     (context?: Option Elab.ContextInfo)
-    (haltOnMatch : Bool)
-    (pred : Elab.Info → Option Elab.ContextInfo → m Bool)
+    (pred : Elab.Info → Option Elab.ContextInfo → m (Bool × Bool))
     : m (List (Elab.Info × Option Elab.ContextInfo × PersistentArray Elab.InfoTree)) := do
   match t with
-  | .context inner t => t.findAllInfoM (inner.mergeIntoOuter? context?) haltOnMatch pred
+  | .context inner t => t.findAllInfoM (inner.mergeIntoOuter? context?) pred
   | .node i children  =>
-    let head := if ← pred i context? then [(i, context?, children)] else []
-    let tail := if haltOnMatch ∧ !head.isEmpty then pure [] else children.toList.mapM (fun t => t.findAllInfoM context? haltOnMatch pred)
+    let (flagCollect, flagRecurse) ← pred i context?
+    let head := if flagCollect then [(i, context?, children)] else []
+    let tail := if ¬ flagRecurse then pure [] else children.toList.mapM (fun t => t.findAllInfoM context? pred)
     return head ++ (← tail).join
   | _ => return []
 
diff --git a/Test/Frontend.lean b/Test/Frontend.lean
index 68aaf94..2259029 100644
--- a/Test/Frontend.lean
+++ b/Test/Frontend.lean
@@ -179,10 +179,19 @@ example (n: Nat) : mystery n + 1 = n + 2 := sorry
 
 def test_capture_type_mismatch : TestT MetaM Unit := do
   let input := "
-def mystery : Nat := true
+def mystery (k: Nat) : Nat := true
   "
   let goalStates ← (collectSorrysFromSource input).run' {}
   let [goalState] := goalStates | panic! s!"Incorrect number of states: {goalStates.length}"
+  checkEq "goals" ((← goalState.serializeGoals (options := {})).map (·.devolatilize)) #[
+    {
+      target := { pp? := "Nat" },
+      vars := #[{
+         userName := "k",
+         type? := .some { pp? := "Nat" },
+      }],
+    }
+  ]
 
 def collectNewConstants (source: String) : MetaM (List (List Name)) := do
   let filename := "<anonymous>"