From d0321e72ddb477a5eea1ebee346c5ee00512d22e Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Sat, 5 Oct 2024 14:49:17 -0700
Subject: [PATCH] feat: Add message diagnostics to frontend.process

---
 Pantograph/Frontend/Basic.lean |  8 +++++++
 Pantograph/Protocol.lean       | 13 +++++++---
 Repl.lean                      | 43 +++++++++++++++++-----------------
 Test/Integration.lean          | 43 ++++++++++++++++++----------------
 4 files changed, 63 insertions(+), 44 deletions(-)

diff --git a/Pantograph/Frontend/Basic.lean b/Pantograph/Frontend/Basic.lean
index 933424c..1074a94 100644
--- a/Pantograph/Frontend/Basic.lean
+++ b/Pantograph/Frontend/Basic.lean
@@ -42,6 +42,14 @@ structure CompilationStep where
   msgs : List Message
   trees : List Elab.InfoTree
 
+namespace CompilationStep
+
+@[export pantograph_frontend_compilation_step_message_strings_m]
+def messageStrings (step: CompilationStep) : IO (Array String) := do
+  List.toArray <$> step.msgs.mapM (·.toString)
+
+end CompilationStep
+
 
 /--
 Process one command, returning a `CompilationStep` and
diff --git a/Pantograph/Protocol.lean b/Pantograph/Protocol.lean
index abfeede..acc2681 100644
--- a/Pantograph/Protocol.lean
+++ b/Pantograph/Protocol.lean
@@ -298,11 +298,18 @@ structure InvokedTactic where
   goalAfter: String
   tactic: String
   deriving Lean.ToJson
-structure FrontendProcessResult where
+
+structure CompilationUnit where
   -- String boundaries of compilation units
-  units: List (Nat × Nat)
+  boundary: (Nat × Nat)
+  -- Tactic invocations
   invocations?: Option (List InvokedTactic) := .none
-  goalStates?: Option (List (Nat × Array Goal)) := .none
+  goalStateId?: Option Nat := .none
+  goals: Array Goal := #[]
+  messages: Array String := #[]
+  deriving Lean.ToJson
+structure FrontendProcessResult where
+  units: List CompilationUnit
   deriving Lean.ToJson
 
 abbrev CR α := Except InteractionError α
diff --git a/Repl.lean b/Repl.lean
index 1277e73..7597164 100644
--- a/Repl.lean
+++ b/Repl.lean
@@ -216,35 +216,36 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
           let env ← Lean.MonadEnv.getEnv
           pure <| .some env
       let (context, state) ← do Frontend.createContextStateFromFile file fileName env? {}
-      let m := Frontend.mapCompilationSteps λ step => do
-        let unitBoundary := (step.src.startPos.byteIdx, step.src.stopPos.byteIdx)
-        let tacticInvocations ← if args.invocations then
-            Frontend.collectTacticsFromCompilationStep step
+      let frontendM := Frontend.mapCompilationSteps λ step => do
+        let boundary := (step.src.startPos.byteIdx, step.src.stopPos.byteIdx)
+        let invocations?: Option (List Protocol.InvokedTactic) ← if args.invocations then
+            let invocations ← Frontend.collectTacticsFromCompilationStep step
+            pure $ .some invocations
           else
-            pure []
+            pure .none
         let sorrys := if args.sorrys then
             Frontend.collectSorrys step
           else
             []
-        return (unitBoundary, tacticInvocations, sorrys)
-      let li ← m.run context |>.run' state
-      let units := li.map λ (unit, _, _) => unit
-      let invocations? := if args.invocations then
-          .some $ li.bind λ (_, invocations, _) => invocations
-        else
-          .none
-      let goalStates? ← if args.sorrys then do
-          let stateIds ← li.filterMapM λ (_, _, sorrys) => do
-            if sorrys.isEmpty then
-              return .none
+        let messages ← step.messageStrings
+        return (boundary, invocations?, sorrys, messages)
+      let li ← frontendM.run context |>.run' state
+      let units ← li.mapM λ (boundary, invocations?, sorrys, messages) => do
+        let (goalStateId?, goals) ← if sorrys.isEmpty then do
+            pure (.none, #[])
+          else do
             let goalState ← runMetaInMainM $ Frontend.sorrysToGoalState sorrys
             let stateId ← newGoalState goalState
             let goals ← goalSerialize goalState options
-            return .some (stateId, goals)
-          pure $ .some stateIds
-        else
-          pure .none
-      return .ok { units, invocations?, goalStates? }
+            pure (.some stateId, goals)
+        return {
+          boundary,
+          invocations?,
+          goalStateId?,
+          goals,
+          messages,
+        }
+      return .ok { units }
     catch e =>
       return .error $ errorI "frontend" (← e.toMessageData.toString)
 
diff --git a/Test/Integration.lean b/Test/Integration.lean
index b3d49fe..413ed1c 100644
--- a/Test/Integration.lean
+++ b/Test/Integration.lean
@@ -176,19 +176,21 @@ def test_frontend_process : Test :=
         ("sorrys", .bool false),
       ]
      ({
-       units := [(0, file.utf8ByteSize)],
-       invocations? := .some [
-         {
-           goalBefore := "⊢ ∀ (p : Prop), p → p",
-           goalAfter := goal1,
-           tactic := "intro p h",
-         },
-         {
-           goalBefore := goal1 ,
-           goalAfter := "",
-           tactic := "exact h",
-         },
-       ]
+       units := [{
+         boundary := (0, file.utf8ByteSize),
+         invocations? := .some [
+           {
+             goalBefore := "⊢ ∀ (p : Prop), p → p",
+             goalAfter := goal1,
+             tactic := "intro p h",
+           },
+           {
+             goalBefore := goal1 ,
+             goalAfter := "",
+             tactic := "exact h",
+           },
+         ]
+       }],
     }: Protocol.FrontendProcessResult),
   ]
 
@@ -212,13 +214,14 @@ def test_frontend_process_sorry : Test :=
         ("sorrys", .bool true),
       ]
      ({
-       units := [
-         (0, solved.utf8ByteSize),
-         (solved.utf8ByteSize, solved.utf8ByteSize + withSorry.utf8ByteSize),
-       ],
-       goalStates? := [
-         (0, #[goal1]),
-       ]
+       units := [{
+         boundary := (0, solved.utf8ByteSize),
+       }, {
+         boundary := (solved.utf8ByteSize, solved.utf8ByteSize + withSorry.utf8ByteSize),
+         goalStateId? := .some 0,
+         goals := #[goal1],
+         messages := #["<anonymous>:2:0: warning: declaration uses 'sorry'\n"],
+       }],
     }: Protocol.FrontendProcessResult),
   ]