fix: Extracting sorrys from sketches

Pantograph/Frontend/Elab.lean

@@ -159,26 +159,100 @@ def collectTacticsFromCompilationStep (step : CompilationStep) : IO (List Protoc
       return t.pretty
     return { goalBefore, goalAfter, tactic }
 
-private def collectSorrysInTree (t : Elab.InfoTree) : List Elab.TermInfo :=
+private def collectSorrysInTree (t : Elab.InfoTree) : List Elab.Info :=
   let infos := findAllInfo t none fun i => match i with
     | .ofTermInfo { expectedType?, expr, stx, .. } =>
-        expr.isSorry ∧ expectedType?.isSome ∧ stx.isOfKind `Lean.Parser.Term.sorry
+      expr.isSorry ∧ expectedType?.isSome ∧ stx.isOfKind `Lean.Parser.Term.sorry
+    | .ofTacticInfo { stx, .. } =>
+      -- The `sorry` term is distinct from the `sorry` tactic
+      stx.isOfKind `Lean.Parser.Tactic.tacticSorry
     | _ => false
-  infos.filterMap fun p => match p with
-    | (.ofTermInfo i, _, _) => .some i
-    | _ => none
+  infos.map fun (i, _, _) => i
 
+-- NOTE: Plural deliberately not spelled "sorries"
 @[export pantograph_frontend_collect_sorrys_m]
-def collectSorrys (step: CompilationStep) : List Elab.TermInfo :=
+def collectSorrys (step: CompilationStep) : List Elab.Info :=
   step.trees.bind collectSorrysInTree
 
-@[export pantograph_frontend_sorrys_to_goal_state]
-def sorrysToGoalState (sorrys : List Elab.TermInfo) : MetaM GoalState := do
-  assert! !sorrys.isEmpty
-  let goals ← sorrys.mapM λ termInfo => Meta.withLCtx termInfo.lctx #[] do
+
+namespace MetaTranslate
+
+structure Context where
+  sourceMCtx : MetavarContext := {}
+  sourceLCtx : LocalContext := {}
+
+/-
+Monadic state for translating a frozen meta state. The underlying `MetaM`
+operates in the "target" context and state.
+-/
+abbrev MetaTranslateM := ReaderT Context MetaM
+
+def getSourceLCtx : MetaTranslateM LocalContext := do pure (← read).sourceLCtx
+def getSourceMCtx : MetaTranslateM MetavarContext := do pure (← read).sourceMCtx
+
+private def translateExpr (expr: Expr) : MetaTranslateM Expr := do
+  let (expr, _) := instantiateMVarsCore (mctx := ← getSourceMCtx) expr
+  return expr
+
+def translateLocalDecl (frozenLocalDecl: LocalDecl) : MetaTranslateM LocalDecl := do
+  let fvarId ← mkFreshFVarId
+  match frozenLocalDecl with
+  | .cdecl index _ userName type bi kind =>
+    return .cdecl index fvarId userName type bi kind
+  | .ldecl index _ userName type value nonDep kind =>
+    return .ldecl index fvarId userName type value nonDep  kind
+
+def translateMVarId (mvarId: MVarId) : MetaTranslateM MVarId := do
+  let shadowDecl := (← getSourceMCtx).findDecl? mvarId |>.get!
+  let target ← translateExpr shadowDecl.type
+  let mvar ← withTheReader Context (λ ctx => { ctx with sourceLCtx := shadowDecl.lctx }) do
+    let lctx ← MonadLCtx.getLCtx
+    let lctx ← (← getSourceLCtx).foldlM (λ lctx frozenLocalDecl => do
+        let localDecl ← translateLocalDecl frozenLocalDecl
+        let lctx := lctx.addDecl localDecl
+        pure lctx
+      ) lctx
+    withTheReader Meta.Context (fun ctx => { ctx with lctx }) do
+      Meta.mkFreshExprSyntheticOpaqueMVar target
+  return mvar.mvarId!
+
+def translateTermInfo (termInfo: Elab.TermInfo) : MetaM MVarId := do
+  let trM : MetaTranslateM MVarId := do
     let type := termInfo.expectedType?.get!
-    let mvar ← Meta.mkFreshExprSyntheticOpaqueMVar type
+    let lctx ← getSourceLCtx
+    let mvar ← withTheReader Meta.Context (fun ctx => { ctx with lctx }) do
+      Meta.mkFreshExprSyntheticOpaqueMVar type
     return mvar.mvarId!
+  trM.run { sourceLCtx := termInfo.lctx }
+
+
+def translateTacticInfoBefore (tacticInfo: Elab.TacticInfo) : MetaM (List MVarId) := do
+  let trM : MetaTranslateM (List MVarId) := do
+    tacticInfo.goalsBefore.mapM translateMVarId
+  trM.run { sourceMCtx := tacticInfo.mctxBefore }
+
+
+end MetaTranslate
+
+export MetaTranslate (MetaTranslateM)
+
+/--
+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]
+def sorrysToGoalState (sorrys : List Elab.Info) : MetaM GoalState := do
+  assert! !sorrys.isEmpty
+  let goals ← sorrys.mapM λ info => Meta.withLCtx info.lctx #[] do
+    match info with
+    | .ofTermInfo termInfo  => do
+      let mvarId ← MetaTranslate.translateTermInfo termInfo
+      return [mvarId]
+    | .ofTacticInfo tacticInfo => do
+      MetaTranslate.translateTacticInfoBefore tacticInfo
+    | _ => panic! "Invalid info"
+  let goals := goals.bind id
   let root := match goals with
     | [] => panic! "This function cannot be called on an empty list"
     | [g] => g

Test/Frontend.lean

@@ -0,0 +1,55 @@
+import LSpec
+import Pantograph
+import Repl
+import Test.Common
+
+open Lean Pantograph
+namespace Pantograph.Test.Frontend
+
+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 Frontend.collectSorrys step
+  let li ← m.run context |>.run' state
+  let goalStates ← li.filterMapM λ sorrys => do
+    if sorrys.isEmpty then
+      return .none
+    let goalState ← Frontend.sorrysToGoalState sorrys
+    return .some goalState
+  return goalStates
+
+def test_multiple_sorries_in_proof : TestT MetaM Unit := do
+  let sketch := "
+theorem plus_n_Sm_proved_formal_sketch : ∀ n m : Nat, n + (m + 1) = (n + m) + 1 := by
+   have h_nat_add_succ: ∀ n m : Nat, n = m := sorry
+   sorry
+  "
+  let goalStates ← (collectSorrysFromSource sketch).run' {}
+  let [goalState] := goalStates | panic! "Illegal number of states"
+  addTest $ LSpec.check "plus_n_Sm" ((← goalState.serializeGoals (options := {})) = #[
+    {
+      name := "_uniq.1",
+      target := { pp? := "∀ (n m : Nat), n = m" },
+      vars := #[
+      ]
+    },
+    {
+      name := "_uniq.4",
+      target := { pp? := "∀ (n m : Nat), n + (m + 1) = n + m + 1" },
+      vars := #[{
+        name := "_uniq.3",
+        userName := "h_nat_add_succ",
+        type? := .some { pp? := "∀ (n m : Nat), n = m" },
+      }],
+    }
+  ])
+
+
+def suite (env : Environment): List (String × IO LSpec.TestSeq) :=
+  let tests := [
+    ("multiple_sorrys_in_proof", test_multiple_sorries_in_proof),
+  ]
+  tests.map (fun (name, test) => (name, runMetaMSeq env $ runTest test))
+
+end Pantograph.Test.Frontend

Test/Main.lean

@@ -1,5 +1,6 @@
 import LSpec
 import Test.Environment
+import Test.Frontend
 import Test.Integration
 import Test.Library
 import Test.Metavar
@@ -44,6 +45,7 @@ def main (args: List String) := do
 
   let suites: List (String × List (String × IO LSpec.TestSeq)) := [
     ("Environment", Environment.suite),
+    ("Frontend", Frontend.suite env_default),
     ("Integration", Integration.suite env_default),
     ("Library", Library.suite env_default),
     ("Metavar", Metavar.suite env_default),