Pantograph/Frontend/Elab.lean

@@ -123,8 +123,6 @@ private def collectSorrysInTree (t : Elab.InfoTree) : IO (List InfoWithContext)
 def collectSorrys (step: CompilationStep) : IO (List InfoWithContext) := do
   return (← step.trees.mapM collectSorrysInTree).join
 
-
-
 /--
 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

Pantograph/Frontend/InfoTree.lean

@@ -103,7 +103,7 @@ partial def InfoTree.findAllInfo
   | .context inner t => findAllInfo t (inner.mergeIntoOuter? context?) haltOnMatch pred
   | .node i children  =>
     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)
+    let tail := if haltOnMatch ∧ !head.isEmpty then [] else children.toList.bind (fun t => findAllInfo t context? haltOnMatch pred)
     head ++ tail
   | _ => []
 
@@ -118,7 +118,7 @@ partial def InfoTree.findAllInfoM [Monad m]
   | .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)
+    let tail := if haltOnMatch ∧ !head.isEmpty then pure [] else children.toList.mapM (fun t => t.findAllInfoM context? haltOnMatch pred)
     return head ++ (← tail).join
   | _ => return []
 

Test/Frontend.lean

@@ -10,8 +10,6 @@ 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
-    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