feat: Monadic info collection

Pantograph/Frontend/Elab.lean

@@ -67,7 +67,7 @@ end TacticInvocation
 /-- Return all `TacticInfo` nodes in an `InfoTree` corresponding to tactics,
 each equipped with its relevant `ContextInfo`, and any children info trees. -/
 private def collectTacticNodes (t : Elab.InfoTree) : List TacticInvocation :=
-  let infos := t.findAllInfo none fun i => match i with
+  let infos := t.findAllInfo none false fun i => match i with
     | .ofTacticInfo _ => true
     | _ => false
   infos.filterMap fun p => match p with
@@ -101,21 +101,27 @@ structure InfoWithContext where
   info: Elab.Info
   context?: Option Elab.ContextInfo := .none
 
-private def collectSorrysInTree (t : Elab.InfoTree) : List InfoWithContext :=
-  let infos := t.findAllInfo none fun i => match i with
-    | .ofTermInfo { expectedType?, expr, stx, .. } =>
-      expr.isSorry ∧ expectedType?.isSome ∧ stx.isOfKind `Lean.Parser.Term.sorry
+private def collectSorrysInTree (t : Elab.InfoTree) : IO (List InfoWithContext) := do
+  let infos ← t.findAllInfoM none true fun i ctx? => match i with
+    | .ofTermInfo { expectedType?, expr, stx, lctx, .. } => do
+      let .some expectedType := expectedType? | return false
+      let .some ctx := ctx? | return false
+      if expr.isSorry ∧ stx.isOfKind `Lean.Parser.Term.sorry then
+        return true
+      ctx.runMetaM lctx do
+        let type ← Meta.inferType expr
+        Bool.not <$> Meta.isDefEq type expectedType
     | .ofTacticInfo { stx, goalsBefore, .. } =>
       -- The `sorry` term is distinct from the `sorry` tactic
       let isSorry := stx.isOfKind `Lean.Parser.Tactic.tacticSorry
-      isSorry ∧ !goalsBefore.isEmpty
-    | _ => false
-  infos.map fun (info, context?, _) => { info, context? }
+      return isSorry ∧ !goalsBefore.isEmpty
+    | _ => return false
+  return infos.map fun (info, context?, _) => { info, context? }
 
 -- NOTE: Plural deliberately not spelled "sorries"
 @[export pantograph_frontend_collect_sorrys_m]
-def collectSorrys (step: CompilationStep) : List InfoWithContext :=
-  step.trees.bind collectSorrysInTree
+def collectSorrys (step: CompilationStep) : IO (List InfoWithContext) := do
+  return (← step.trees.mapM collectSorrysInTree).join
 
 
 

Pantograph/Frontend/InfoTree.lean

@@ -93,14 +93,35 @@ partial def InfoTree.filter (p : Info → Bool) (m : MVarId → Bool := fun _ =>
   | .hole mvar => if m mvar then [.hole mvar] else []
 
 /-- Analogue of `Lean.Elab.InfoTree.findInfo?`, but that returns a list of all results. -/
-partial def InfoTree.findAllInfo (t : InfoTree) (context?: Option Elab.ContextInfo) (pred : Elab.Info → Bool) :
-    List (Elab.Info × Option Elab.ContextInfo × PersistentArray Elab.InfoTree) :=
+partial def InfoTree.findAllInfo
+    (t : InfoTree)
+    (context?: Option Elab.ContextInfo)
+    (haltOnMatch : Bool := false)
+    (pred : Elab.Info → Bool)
+    : List (Elab.Info × Option Elab.ContextInfo × PersistentArray Elab.InfoTree) :=
   match t with
-  | .context inner t => findAllInfo t (inner.mergeIntoOuter? context?) pred
+  | .context inner t => findAllInfo t (inner.mergeIntoOuter? context?) haltOnMatch pred
   | .node i children  =>
-      (if pred i then [(i, context?, children)] else []) ++ children.toList.bind (fun t => findAllInfo t context? pred)
+    let head := if pred i then [(i, context?, children)] else []
+    let tail := if haltOnMatch then [] else children.toList.bind (fun t => findAllInfo t context? haltOnMatch pred)
+    head ++ tail
   | _ => []
 
+/-- Monadic analogue of `findAllInfo` -/
+partial def InfoTree.findAllInfoM [Monad m]
+    (t : InfoTree)
+    (context?: Option Elab.ContextInfo)
+    (haltOnMatch : Bool)
+    (pred : Elab.Info → Option Elab.ContextInfo → m 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
+  | .node i children  =>
+    let head := if ← pred i context? then [(i, context?, children)] else []
+    let tail := if haltOnMatch then pure [] else children.toList.mapM (fun t => t.findAllInfoM context? haltOnMatch pred)
+    return head ++ (← tail).join
+  | _ => return []
+
 @[export pantograph_infotree_to_string_m]
 partial def InfoTree.toString (t : InfoTree) (ctx?: Option Elab.ContextInfo := .none) : IO String := do
   match t with

Repl.lean

@@ -257,10 +257,10 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
             pure $ .some invocations
           else
             pure .none
-        let sorrys := if args.sorrys then
+        let sorrys ← if args.sorrys then
             Frontend.collectSorrys step
           else
-            []
+            pure []
         let messages ← step.messageStrings
         return (step.before, boundary, invocations?, sorrys, messages)
       let li ← frontendM.run context |>.run' state

Test/Frontend.lean

@@ -10,7 +10,9 @@ def collectSorrysFromSource (source: String) : MetaM (List GoalState) := do
   let filename := "<anonymous>"
   let (context, state) ← do Frontend.createContextStateFromFile source filename (← getEnv) {}
   let m := Frontend.mapCompilationSteps λ step => do
-    return (step.before, Frontend.collectSorrys step)
+    for tree in step.trees do
+      IO.println s!"{← tree.toString}"
+    return (step.before, ← Frontend.collectSorrys step)
   let li ← m.run context |>.run' state
   let goalStates ← li.filterMapM λ (env, sorrys) => withEnv env do
     if sorrys.isEmpty then
@@ -177,6 +179,13 @@ example (n: Nat) : mystery n + 1 = n + 2 := sorry
     }
   ])
 
+def test_capture_type_mismatch : TestT MetaM Unit := do
+  let input := "
+def mystery : Nat := true
+  "
+  let goalStates ← (collectSorrysFromSource input).run' {}
+  let [goalState] := goalStates | panic! s!"Incorrect number of states: {goalStates.length}"
+
 
 def suite (env : Environment): List (String × IO LSpec.TestSeq) :=
   let tests := [
@@ -185,6 +194,7 @@ def suite (env : Environment): List (String × IO LSpec.TestSeq) :=
     ("sorry_in_induction", test_sorry_in_induction),
     ("sorry_in_coupled", test_sorry_in_coupled),
     ("environment_capture", test_environment_capture),
+    ("capture_type_mismatch", test_capture_type_mismatch),
   ]
   tests.map (fun (name, test) => (name, runMetaMSeq env $ runTest test))