From c77f14d383a3d46b1f749a3c55204f35aee0f084 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Tue, 28 Jan 2025 17:40:49 -0800
Subject: [PATCH] fix: Use context environment in sorry capture

---
 Pantograph/Frontend/Elab.lean | 49 +++++++++++++++++++----------------
 Repl.lean                     |  4 +--
 2 files changed, 28 insertions(+), 25 deletions(-)

diff --git a/Pantograph/Frontend/Elab.lean b/Pantograph/Frontend/Elab.lean
index 1e16a82..6dc6a28 100644
--- a/Pantograph/Frontend/Elab.lean
+++ b/Pantograph/Frontend/Elab.lean
@@ -148,34 +148,39 @@ structure AnnotatedGoalState where
 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`.
+
+WARNING: Behaviour is unstable when there are multiple `sorry`s. Consider using
+the draft tactic instead.
 -/
 @[export pantograph_frontend_sorrys_to_goal_state_m]
 def sorrysToGoalState (sorrys : List InfoWithContext) : MetaM AnnotatedGoalState := do
+  let env := sorrys.head? >>= (·.context?) |>.map (·.env) |>.getD (← getEnv)
   assert! !sorrys.isEmpty
-  let goalsM := sorrys.mapM λ i => do
-    match i.info with
-    | .ofTermInfo termInfo  => do
-      let mvarId ← MetaTranslate.translateMVarFromTermInfo termInfo i.context?
-      if (← mvarId.getType).hasSorry then
-        throwError s!"Coupling is not allowed in drafting"
-      return [(mvarId, stxByteRange termInfo.stx)]
-    | .ofTacticInfo tacticInfo => do
-      let mvarIds ← MetaTranslate.translateMVarFromTacticInfoBefore tacticInfo i.context?
-      for mvarId in mvarIds do
+  withEnv env do
+    let goalsM := sorrys.mapM λ i => do
+      match i.info with
+      | .ofTermInfo termInfo  => do
+        let mvarId ← MetaTranslate.translateMVarFromTermInfo termInfo i.context?
         if (← mvarId.getType).hasSorry then
           throwError s!"Coupling is not allowed in drafting"
-      let range := stxByteRange tacticInfo.stx
-      return mvarIds.map (·, range)
-    | _ => panic! "Invalid info"
-  let annotatedGoals := List.flatten (← goalsM.run {} |>.run' {})
-  let goals := annotatedGoals.map Prod.fst
-  let srcBoundaries := annotatedGoals.map Prod.snd
-  let root := match goals with
-    | [] => panic! "No MVars generated"
-    | [g] => g
-    | _ => { name := .anonymous }
-  let state ← GoalState.createFromMVars goals root
-  return { state, srcBoundaries }
+        return [(mvarId, stxByteRange termInfo.stx)]
+      | .ofTacticInfo tacticInfo => do
+        let mvarIds ← MetaTranslate.translateMVarFromTacticInfoBefore tacticInfo i.context?
+        for mvarId in mvarIds do
+          if (← mvarId.getType).hasSorry then
+            throwError s!"Coupling is not allowed in drafting"
+        let range := stxByteRange tacticInfo.stx
+        return mvarIds.map (·, range)
+      | _ => panic! "Invalid info"
+    let annotatedGoals := List.flatten (← goalsM.run {} |>.run' {})
+    let goals := annotatedGoals.map Prod.fst
+    let srcBoundaries := annotatedGoals.map Prod.snd
+    let root := match goals with
+      | [] => panic! "No MVars generated"
+      | [g] => g
+      | _ => { name := .anonymous }
+    let state ← GoalState.createFromMVars goals root
+    return { state, srcBoundaries }
 
 
 @[export pantograph_frontend_collect_new_defined_constants_m]
diff --git a/Repl.lean b/Repl.lean
index 0e745a4..c8f4028 100644
--- a/Repl.lean
+++ b/Repl.lean
@@ -44,8 +44,6 @@ structure CompilationUnit where
   messages : Array String
   newConstants : List Name
 
-
-
 def frontend_process_inner (args: Protocol.FrontendProcess): MainM (CR Protocol.FrontendProcessResult) := do
   let options := (← get).options
   let (fileName, file) ← match args.fileName?, args.file? with
@@ -78,7 +76,7 @@ def frontend_process_inner (args: Protocol.FrontendProcess): MainM (CR Protocol.
       else
         pure []
     return {
-      env := step.after,
+      env := step.before,
       boundary,
       invocations,
       sorrys,