From dd00d803d1fe500f46710bb32bfcc9eeeb2d8062 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Mon, 9 Dec 2024 20:38:27 -0800
Subject: [PATCH] feat: Collect sorry/elab failure boundaries

---
 Pantograph/Frontend/Basic.lean |  7 +++++++
 Pantograph/Frontend/Elab.lean  | 19 ++++++++++++++-----
 Pantograph/Protocol.lean       |  2 ++
 Repl.lean                      | 13 +++++++------
 Test/Frontend.lean             |  4 ++--
 5 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/Pantograph/Frontend/Basic.lean b/Pantograph/Frontend/Basic.lean
index 1074a94..87decd4 100644
--- a/Pantograph/Frontend/Basic.lean
+++ b/Pantograph/Frontend/Basic.lean
@@ -30,6 +30,13 @@ end Lean.PersistentArray
 
 namespace Pantograph.Frontend
 
+@[export pantograph_frontend_stx_byte_range]
+def stxByteRange (stx : Syntax) : String.Pos × String.Pos :=
+  let pos := stx.getPos?.getD 0
+  let endPos := stx.getTailPos?.getD 0
+  (pos, endPos)
+
+
 abbrev FrontendM := Elab.Frontend.FrontendM
 
 structure CompilationStep where
diff --git a/Pantograph/Frontend/Elab.lean b/Pantograph/Frontend/Elab.lean
index 2eff084..a33fbef 100644
--- a/Pantograph/Frontend/Elab.lean
+++ b/Pantograph/Frontend/Elab.lean
@@ -129,28 +129,37 @@ private def collectSorrysInTree (t : Elab.InfoTree) : IO (List InfoWithContext)
 def collectSorrys (step: CompilationStep) : IO (List InfoWithContext) := do
   return (← step.trees.mapM collectSorrysInTree).join
 
+structure AnnotatedGoalState where
+  state : GoalState
+  srcBoundaries : List (String.Pos × String.Pos)
+
 /--
 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_m]
-def sorrysToGoalState (sorrys : List InfoWithContext) : MetaM GoalState := do
+def sorrysToGoalState (sorrys : List InfoWithContext) : MetaM AnnotatedGoalState := do
   assert! !sorrys.isEmpty
   let goalsM := sorrys.mapM λ i => do
     match i.info with
     | .ofTermInfo termInfo  => do
       let mvarId ← MetaTranslate.translateMVarFromTermInfo termInfo i.context?
-      return [mvarId]
+      return [(mvarId, stxByteRange termInfo.stx)]
     | .ofTacticInfo tacticInfo => do
-      MetaTranslate.translateMVarFromTacticInfoBefore tacticInfo i.context?
+      let mvarIds ← MetaTranslate.translateMVarFromTacticInfoBefore tacticInfo i.context?
+      let range := stxByteRange tacticInfo.stx
+      return mvarIds.map (·, range)
     | _ => panic! "Invalid info"
-  let goals := List.join (← goalsM.run {} |>.run' {})
+  let annotatedGoals := List.join (← 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 }
-  GoalState.createFromMVars goals root
+  let state ← GoalState.createFromMVars goals root
+  return { state, srcBoundaries }
 
 
 @[export pantograph_frontend_collect_new_defined_constants_m]
diff --git a/Pantograph/Protocol.lean b/Pantograph/Protocol.lean
index 0cb6cac..5cb24e8 100644
--- a/Pantograph/Protocol.lean
+++ b/Pantograph/Protocol.lean
@@ -329,6 +329,8 @@ structure CompilationUnit where
   invocations?: Option (List InvokedTactic) := .none
   goalStateId?: Option Nat := .none
   goals: Array Goal := #[]
+  -- Code segments which generated the goals
+  goalSrcBoundaries: Array (Nat × Nat) := #[]
   messages: Array String := #[]
   deriving Lean.ToJson
 structure FrontendProcessResult where
diff --git a/Repl.lean b/Repl.lean
index f1c8f42..201a841 100644
--- a/Repl.lean
+++ b/Repl.lean
@@ -265,18 +265,19 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
         return (step.before, boundary, invocations?, sorrys, messages)
       let li ← frontendM.run context |>.run' state
       let units ← li.mapM λ (env, boundary, invocations?, sorrys, messages) => Lean.withEnv env do
-        let (goalStateId?, goals) ← if sorrys.isEmpty then do
-            pure (.none, #[])
+        let (goalStateId?, goals, goalSrcBoundaries) ← if sorrys.isEmpty then do
+            pure (.none, #[], #[])
           else do
-            let goalState ← runMetaInMainM $ Frontend.sorrysToGoalState sorrys
-            let stateId ← newGoalState goalState
-            let goals ← goalSerialize goalState options
-            pure (.some stateId, goals)
+            let { state, srcBoundaries } ← runMetaInMainM $ Frontend.sorrysToGoalState sorrys
+            let stateId ← newGoalState state
+            let goals ← goalSerialize state options
+            pure (.some stateId, goals, srcBoundaries.toArray.map (λ (b, e) => (b.byteIdx, e.byteIdx)))
         return {
           boundary,
           invocations?,
           goalStateId?,
           goals,
+          goalSrcBoundaries,
           messages,
         }
       return .ok { units }
diff --git a/Test/Frontend.lean b/Test/Frontend.lean
index 2259029..a3b73ae 100644
--- a/Test/Frontend.lean
+++ b/Test/Frontend.lean
@@ -15,8 +15,8 @@ def collectSorrysFromSource (source: String) : MetaM (List GoalState) := do
   let goalStates ← li.filterMapM λ (env, sorrys) => withEnv env do
     if sorrys.isEmpty then
       return .none
-    let goalState ← Frontend.sorrysToGoalState sorrys
-    return .some goalState
+    let { state, .. } ← Frontend.sorrysToGoalState sorrys
+    return .some state
   return goalStates
 
 def test_multiple_sorrys_in_proof : TestT MetaM Unit := do