From 22716597d335257bd8952ace82af1396c485a91b Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Mon, 16 Jun 2025 18:10:11 -0700
Subject: [PATCH] feat(frontend): Write invocation points to file

---
 Pantograph/Protocol.lean |  14 +-
 Repl.lean                |  10 +-
 Test/Integration.lean    | 450 +++++++++++++++++++--------------------
 3 files changed, 238 insertions(+), 236 deletions(-)

diff --git a/Pantograph/Protocol.lean b/Pantograph/Protocol.lean
index d77afd4..07e4777 100644
--- a/Pantograph/Protocol.lean
+++ b/Pantograph/Protocol.lean
@@ -361,8 +361,8 @@ structure FrontendProcess where
   readHeader : Bool := false
   -- Alter the REPL environment after the compilation units.
   inheritEnv : Bool := false
-  -- collect tactic invocations
-  invocations: Bool := false
+  -- collect tactic invocations and output to a given file
+  invocations?: Option String := .none
   -- collect `sorry`s
   sorrys: Bool := false
   -- collect type errors
@@ -383,8 +383,8 @@ structure CompilationUnit where
   -- String boundaries of compilation units
   boundary: (Nat × Nat)
   messages: Array String := #[]
-  -- Tactic invocations
-  invocations?: Option (List InvokedTactic) := .none
+  -- Number of tactic invocations
+  nInvocations?: Option Nat := .none
   goalStateId?: Option Nat := .none
   goals?: Option (Array Goal) := .none
   -- Code segments which generated the goals
@@ -396,6 +396,12 @@ structure CompilationUnit where
 structure FrontendProcessResult where
   units: List CompilationUnit
   deriving Lean.ToJson
+structure FrontendDataUnit where
+  invocations? : Option (List Protocol.InvokedTactic) := .none
+  deriving Lean.ToJson
+structure FrontendData where
+  units : List FrontendDataUnit
+  deriving Lean.ToJson
 
 abbrev FallibleT := ExceptT InteractionError
 
diff --git a/Repl.lean b/Repl.lean
index 2822993..50118cc 100644
--- a/Repl.lean
+++ b/Repl.lean
@@ -125,7 +125,7 @@ def frontend_process (args: Protocol.FrontendProcess): EMainM Protocol.FrontendP
   let frontendM: Elab.Frontend.FrontendM (List CompilationUnit) :=
     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: Option (List Protocol.InvokedTactic) ← if args.invocations?.isSome then
         Frontend.collectTacticsFromCompilationStep step
       else
         pure []
@@ -152,6 +152,10 @@ def frontend_process (args: Protocol.FrontendProcess): EMainM Protocol.FrontendP
     if let .some scope := state'.commandState.scopes.head? then
       -- modify the scope
       set { ← getMainState with scope }
+  if let .some fileName := args.invocations? then
+    let units := li.map λ unit => { invocations? := .some unit.invocations }
+    let data : Protocol.FrontendData := { units }
+    IO.FS.writeFile fileName (toJson data |>.compress)
   let units ← li.mapM λ step => withEnv step.env do
     let newConstants? := if args.newConstants then
         .some $ step.newConstants.toArray.map λ name => name.toString
@@ -167,11 +171,11 @@ def frontend_process (args: Protocol.FrontendProcess): EMainM Protocol.FrontendP
         let stateId ← newGoalState state
         let srcBoundaries := srcBoundaries.toArray.map (λ (b, e) => (b.byteIdx, e.byteIdx))
         pure (.some stateId, .some goals, .some srcBoundaries)
-    let invocations? := if args.invocations then .some step.invocations else .none
+    let nInvocations? := if args.invocations?.isSome then .some step.invocations.length else .none
     return {
       boundary := step.boundary,
       messages := step.messages,
-      invocations?,
+      nInvocations?,
       goalStateId?,
       goals?,
       goalSrcBoundaries?,
diff --git a/Test/Integration.lean b/Test/Integration.lean
index 9945185..bca5e3a 100644
--- a/Test/Integration.lean
+++ b/Test/Integration.lean
@@ -17,53 +17,59 @@ deriving instance Lean.ToJson for Protocol.GoalStart
 deriving instance Lean.ToJson for Protocol.GoalPrint
 deriving instance Lean.ToJson for Protocol.GoalTactic
 deriving instance Lean.ToJson for Protocol.FrontendProcess
+deriving instance Lean.ToJson for Protocol.FrontendDataUnit
+deriving instance Lean.ToJson for Protocol.FrontendData
+
+abbrev TestM α := TestT MainM α
+abbrev Test := TestM Unit
 
 def step { α β } [Lean.ToJson α] [Lean.ToJson β] (cmd: String) (payload: α)
-    (expected: β) (name? : Option String := .none): MainM LSpec.TestSeq := do
+    (expected: β) (name? : Option String := .none) : TestM Unit := do
   let payload := Lean.toJson payload
   let name := name?.getD s!"{cmd} {payload.compress}"
   let result ← Repl.execute { cmd, payload }
-  return LSpec.test name (result.pretty = (Lean.toJson expected).pretty)
-
-abbrev Test := List (MainM LSpec.TestSeq)
+  checkEq name result.compress (Lean.toJson expected).compress
+def stepFile { α } [Lean.ToJson α] (name : String) (path : String)
+    (expected : α) : TestM Unit := do
+  let content ← IO.FS.readFile path
+  let payload? : Except String Lean.Json := Lean.Json.parse content
+  match payload? with
+  | .ok payload =>
+    let expected := Lean.toJson expected
+    checkEq name payload.compress expected.compress
+  | .error e => fail s!"{name} {e}"
 
 def test_expr_echo : Test :=
-  [
-    step "expr.echo"
-     ({ expr := "λ {α : Sort (u + 1)} => List α", levels? := .some #["u"]}: Protocol.ExprEcho)
-     ({
-       type := { pp? := .some "{α : Type u} → Type u" },
-       expr := { pp? := .some "fun {α} => List α" }
-     }: Protocol.ExprEchoResult),
-  ]
+  step "expr.echo"
+   ({ expr := "λ {α : Sort (u + 1)} => List α", levels? := .some #["u"]}: Protocol.ExprEcho)
+   ({
+     type := { pp? := .some "{α : Type u} → Type u" },
+     expr := { pp? := .some "fun {α} => List α" }
+   }: Protocol.ExprEchoResult)
 
-def test_option_modify : Test :=
+def test_option_modify : Test := do
   let pp? := Option.some "∀ (n : Nat), n + 1 = n.succ"
   let sexp? := Option.some "(:forall n (:c Nat) ((:c Eq) (:c Nat) ((:c HAdd.hAdd) (:c Nat) (:c Nat) (:c Nat) ((:c instHAdd) (:c Nat) (:c instAddNat)) 0 ((:c OfNat.ofNat) (:c Nat) (:lit 1) ((:c instOfNatNat) (:lit 1)))) ((:c Nat.succ) 0)))"
   let module? := Option.some "Init.Data.Nat.Basic"
   let options: Protocol.Options := {}
-  [
-    step "env.inspect" ({ name := "Nat.add_one" } : Protocol.EnvInspect)
-     ({ type := { pp? }, module? }: Protocol.EnvInspectResult),
-    step "options.set" ({ printExprAST? := .some true } : Protocol.OptionsSet)
-     ({ }: Protocol.OptionsSetResult),
-    step "env.inspect" ({ name := "Nat.add_one" } : Protocol.EnvInspect)
-     ({ type := { pp?, sexp? }, module? }: Protocol.EnvInspectResult),
-    step "options.print" ({} : Protocol.OptionsPrint)
-     ({ options with printExprAST := true }: Protocol.Options),
-  ]
-def test_malformed_command : Test :=
+  step "env.inspect" ({ name := "Nat.add_one" } : Protocol.EnvInspect)
+   ({ type := { pp? }, module? }: Protocol.EnvInspectResult)
+  step "options.set" ({ printExprAST? := .some true } : Protocol.OptionsSet)
+   ({ }: Protocol.OptionsSetResult)
+  step "env.inspect" ({ name := "Nat.add_one" } : Protocol.EnvInspect)
+   ({ type := { pp?, sexp? }, module? }: Protocol.EnvInspectResult)
+  step "options.print" ({} : Protocol.OptionsPrint)
+   ({ options with printExprAST := true }: Protocol.Options)
+
+def test_malformed_command : Test := do
   let invalid := "invalid"
-  [
-    step invalid ({ name := "Nat.add_one" }: Protocol.EnvInspect)
-     ({ error := "command", desc := s!"Unknown command {invalid}" }: Protocol.InteractionError)
-     (name? := .some "Invalid Command"),
-    step "expr.echo" (Lean.Json.mkObj [(invalid, .str "Random garbage data")])
-     ({ error := "command", desc := s!"Unable to parse json: Pantograph.Protocol.ExprEcho.expr: String expected" }:
-      Protocol.InteractionError)
-    (name? := .some "JSON Deserialization")
-  ]
-def test_tactic : Test :=
+  step invalid ({ name := "Nat.add_one" }: Protocol.EnvInspect)
+   ({ error := "command", desc := s!"Unknown command {invalid}" }: Protocol.InteractionError)
+   (name? := .some "Invalid Command")
+  step "expr.echo" (Lean.Json.mkObj [(invalid, .str "Random garbage data")])
+   ({ error := "command", desc := s!"Unable to parse json: Pantograph.Protocol.ExprEcho.expr: String expected" }:
+    Protocol.InteractionError) (name? := .some "JSON Deserialization")
+def test_tactic : Test := do
   let varX := { name := "_uniq.10", userName := "x", type? := .some { pp? := .some "Prop" }}
   let goal1: Protocol.Goal := {
     name := "_uniq.11",
@@ -78,43 +84,39 @@ def test_tactic : Test :=
       { name := "_uniq.13", userName := "y", type? := .some { pp? := .some "Prop" }}
     ],
   }
-  [
-    step "goal.start" ({ expr := "∀ (p q: Prop), p ∨ q → q ∨ p" }: Protocol.GoalStart)
-     ({ stateId := 0, root := "_uniq.9" }: Protocol.GoalStartResult),
-    step "goal.tactic" ({ stateId := 0, tactic? := .some "intro x" }: Protocol.GoalTactic)
-     ({ nextStateId? := .some 1, goals? := #[goal1], }: Protocol.GoalTacticResult),
-    step "goal.print" ({ stateId := 1, parentExpr? := .some true, rootExpr? := .some true }: Protocol.GoalPrint)
-     ({
-       root? := .some { pp? := "fun x => ?m.11"},
-       parent? := .some { pp? := .some "fun x => ?m.11" },
-     }: Protocol.GoalPrintResult),
-    step "goal.tactic" ({ stateId := 1, tactic? := .some "intro y" }: Protocol.GoalTactic)
-     ({ nextStateId? := .some 2, goals? := #[goal2], }: Protocol.GoalTacticResult),
-    step "goal.tactic" ({ stateId := 1, tactic? := .some "apply Nat.le_of_succ_le" }: Protocol.GoalTactic)
-     ({ messages? := .some #["tactic 'apply' failed, failed to unify\n  ∀ {m : Nat}, Nat.succ ?n ≤ m → ?n ≤ m\nwith\n  ∀ (q : Prop), x ∨ q → q ∨ x\nx : Prop\n⊢ ∀ (q : Prop), x ∨ q → q ∨ x"] }:
-      Protocol.GoalTacticResult),
-    step "goal.tactic" ({ stateId := 0, tactic? := .some "sorry" }: Protocol.GoalTactic)
-     ({ nextStateId? := .some 3, goals? := .some #[], hasSorry := true }: Protocol.GoalTacticResult),
-  ]
+  step "goal.start" ({ expr := "∀ (p q: Prop), p ∨ q → q ∨ p" }: Protocol.GoalStart)
+   ({ stateId := 0, root := "_uniq.9" }: Protocol.GoalStartResult)
+  step "goal.tactic" ({ stateId := 0, tactic? := .some "intro x" }: Protocol.GoalTactic)
+   ({ nextStateId? := .some 1, goals? := #[goal1], }: Protocol.GoalTacticResult)
+  step "goal.print" ({ stateId := 1, parentExpr? := .some true, rootExpr? := .some true }: Protocol.GoalPrint)
+   ({
+     root? := .some { pp? := "fun x => ?m.11"},
+     parent? := .some { pp? := .some "fun x => ?m.11" },
+   }: Protocol.GoalPrintResult)
+  step "goal.tactic" ({ stateId := 1, tactic? := .some "intro y" }: Protocol.GoalTactic)
+   ({ nextStateId? := .some 2, goals? := #[goal2], }: Protocol.GoalTacticResult)
+  step "goal.tactic" ({ stateId := 1, tactic? := .some "apply Nat.le_of_succ_le" }: Protocol.GoalTactic)
+   ({ messages? := .some #["tactic 'apply' failed, failed to unify\n  ∀ {m : Nat}, Nat.succ ?n ≤ m → ?n ≤ m\nwith\n  ∀ (q : Prop), x ∨ q → q ∨ x\nx : Prop\n⊢ ∀ (q : Prop), x ∨ q → q ∨ x"] }:
+    Protocol.GoalTacticResult)
+  step "goal.tactic" ({ stateId := 0, tactic? := .some "sorry" }: Protocol.GoalTactic)
+   ({ nextStateId? := .some 3, goals? := .some #[], hasSorry := true }: Protocol.GoalTacticResult)
 example : (1 : Nat) + (2 * 3) = 1 + (4 - 3) + (6 - 4) + 3 := by
   simp
-def test_tactic_timeout : Test :=
-  [
-    step "goal.start" ({ expr := "(1 : Nat) + (2 * 3) = 1 + (4 - 3) + (6 - 4) + 3" }: Protocol.GoalStart)
-     ({ stateId := 0, root := "_uniq.319" }: Protocol.GoalStartResult),
-    -- timeout of 10 milliseconds
-    step "options.set" ({ timeout? := .some 10 } : Protocol.OptionsSet)
-     ({ }: Protocol.OptionsSetResult),
-    step "goal.tactic" ({ stateId := 0, expr? := .some "by\nsleep 1000; simp" }: Protocol.GoalTactic)
-     ({ error := "internal", desc := "interrupt" }: Protocol.InteractionError),
-    -- ensure graceful recovery
-    step "options.set" ({ timeout? := .some 0 } : Protocol.OptionsSet)
-     ({ }: Protocol.OptionsSetResult),
-    step "goal.tactic" ({ stateId := 0, tactic? := .some "simp" }: Protocol.GoalTactic)
-     ({ nextStateId? := .some 1, goals? := .some #[], }: Protocol.GoalTacticResult),
-  ]
+def test_tactic_timeout : Test := do
+  step "goal.start" ({ expr := "(1 : Nat) + (2 * 3) = 1 + (4 - 3) + (6 - 4) + 3" }: Protocol.GoalStart)
+   ({ stateId := 0, root := "_uniq.319" }: Protocol.GoalStartResult)
+  -- timeout of 10 milliseconds
+  step "options.set" ({ timeout? := .some 10 } : Protocol.OptionsSet)
+   ({ }: Protocol.OptionsSetResult)
+  step "goal.tactic" ({ stateId := 0, expr? := .some "by\nsleep 1000; simp" }: Protocol.GoalTactic)
+   ({ error := "internal", desc := "interrupt" }: Protocol.InteractionError)
+  -- ensure graceful recovery
+  step "options.set" ({ timeout? := .some 0 } : Protocol.OptionsSet)
+   ({ }: Protocol.OptionsSetResult)
+  step "goal.tactic" ({ stateId := 0, tactic? := .some "simp" }: Protocol.GoalTactic)
+   ({ nextStateId? := .some 1, goals? := .some #[], }: Protocol.GoalTacticResult)
 
-def test_automatic_mode (automatic: Bool): Test :=
+def test_automatic_mode (automatic: Bool): Test := do
   let varsPQ := #[
     { name := "_uniq.10", userName := "p", type? := .some { pp? := .some "Prop" }},
     { name := "_uniq.13", userName := "q", type? := .some { pp? := .some "Prop" }}
@@ -150,196 +152,186 @@ def test_automatic_mode (automatic: Bool): Test :=
       { name := "_uniq.49", userName := "h✝", type? := .some { pp? := .some "p" }, isInaccessible := true}
     ],
   }
-  [
-    step "options.set" ({automaticMode? := .some automatic}: Protocol.OptionsSet)
-     ({}: Protocol.OptionsSetResult),
-    step "goal.start" ({ expr := "∀ (p q: Prop), p ∨ q → q ∨ p"} : Protocol.GoalStart)
-     ({ stateId := 0, root := "_uniq.9" }: Protocol.GoalStartResult),
-    step "goal.tactic" ({ stateId := 0, tactic? := .some "intro p q h" }: Protocol.GoalTactic)
-     ({ nextStateId? := .some 1, goals? := #[goal1], }: Protocol.GoalTacticResult),
-    step "goal.tactic" ({ stateId := 1, tactic? := .some "cases h" }: Protocol.GoalTactic)
-     ({ nextStateId? := .some 2, goals? := #[goal2l, goal2r], }: Protocol.GoalTacticResult),
-    let goals? := if automatic then #[goal3l, goal2r] else #[goal3l]
-    step "goal.tactic" ({ stateId := 2, tactic? := .some "apply Or.inr" }: Protocol.GoalTactic)
-     ({ nextStateId? := .some 3, goals?, }: Protocol.GoalTacticResult),
-  ]
+  step "options.set" ({automaticMode? := .some automatic}: Protocol.OptionsSet)
+   ({}: Protocol.OptionsSetResult)
+  step "goal.start" ({ expr := "∀ (p q: Prop), p ∨ q → q ∨ p"} : Protocol.GoalStart)
+   ({ stateId := 0, root := "_uniq.9" }: Protocol.GoalStartResult)
+  step "goal.tactic" ({ stateId := 0, tactic? := .some "intro p q h" }: Protocol.GoalTactic)
+   ({ nextStateId? := .some 1, goals? := #[goal1], }: Protocol.GoalTacticResult)
+  step "goal.tactic" ({ stateId := 1, tactic? := .some "cases h" }: Protocol.GoalTactic)
+   ({ nextStateId? := .some 2, goals? := #[goal2l, goal2r], }: Protocol.GoalTacticResult)
+  let goals? := if automatic then #[goal3l, goal2r] else #[goal3l]
+  step "goal.tactic" ({ stateId := 2, tactic? := .some "apply Or.inr" }: Protocol.GoalTactic)
+   ({ nextStateId? := .some 3, goals?, }: Protocol.GoalTacticResult)
 
-def test_env_add_inspect : Test :=
+def test_env_add_inspect : Test := do
   let name1 := "Pantograph.mystery"
   let name2 := "Pantograph.mystery2"
   let name3 := "Pantograph.mystery3"
-  [
-    step "env.add"
-      ({
-        name := name1,
-        value := "λ (a b: Prop) => Or a b",
-        isTheorem := false
-      }: Protocol.EnvAdd)
-     ({}: Protocol.EnvAddResult),
-    step "env.inspect" ({name := name1, value? := .some true} : Protocol.EnvInspect)
-     ({
-       value? := .some { pp? := .some "fun a b => a ∨ b" },
-       type := { pp? := .some "Prop → Prop → Prop" },
-     }:
-      Protocol.EnvInspectResult),
-    step "env.add"
-      ({
-        name := name2,
-        type? := "Nat → Int",
-        value := "λ (a: Nat) => a + 1",
-        isTheorem := false
-      }: Protocol.EnvAdd)
-     ({}: Protocol.EnvAddResult),
-    step "env.inspect" ({name := name2, value? := .some true} : Protocol.EnvInspect)
-     ({
-       value? := .some { pp? := .some "fun a => ↑a + 1" },
-       type := { pp? := .some "Nat → Int" },
-     }:
-      Protocol.EnvInspectResult),
-    step "env.add"
-      ({
-        name := name3,
-        levels? := .some #["u"]
-        type? := "(α : Type u) → α → (α × α)",
-        value := "λ (α : Type u) (x : α) => (x, x)",
-        isTheorem := false
-      }: Protocol.EnvAdd)
-     ({}: Protocol.EnvAddResult),
-    step "env.inspect" ({name := name3} : Protocol.EnvInspect)
-     ({
-       type := { pp? := .some "(α : Type u) → α → α × α" },
-     }:
-      Protocol.EnvInspectResult),
-  ]
+  step "env.add"
+    ({
+      name := name1,
+      value := "λ (a b: Prop) => Or a b",
+      isTheorem := false
+    }: Protocol.EnvAdd)
+   ({}: Protocol.EnvAddResult)
+  step "env.inspect" ({name := name1, value? := .some true} : Protocol.EnvInspect)
+   ({
+     value? := .some { pp? := .some "fun a b => a ∨ b" },
+     type := { pp? := .some "Prop → Prop → Prop" },
+   }: Protocol.EnvInspectResult)
+  step "env.add"
+    ({
+      name := name2,
+      type? := "Nat → Int",
+      value := "λ (a: Nat) => a + 1",
+      isTheorem := false
+    }: Protocol.EnvAdd)
+   ({}: Protocol.EnvAddResult)
+  step "env.inspect" ({name := name2, value? := .some true} : Protocol.EnvInspect)
+   ({
+     value? := .some { pp? := .some "fun a => ↑a + 1" },
+     type := { pp? := .some "Nat → Int" },
+   }: Protocol.EnvInspectResult)
+  step "env.add"
+    ({
+      name := name3,
+      levels? := .some #["u"]
+      type? := "(α : Type u) → α → (α × α)",
+      value := "λ (α : Type u) (x : α) => (x, x)",
+      isTheorem := false
+    }: Protocol.EnvAdd)
+   ({}: Protocol.EnvAddResult)
+  step "env.inspect" ({name := name3} : Protocol.EnvInspect)
+   ({
+     type := { pp? := .some "(α : Type u) → α → α × α" },
+   }: Protocol.EnvInspectResult)
 
 example : ∀ (p: Prop), p → p := by
   intro p h
   exact h
 
-def test_frontend_process : Test :=
+def test_frontend_process : Test := do
   let file := "example : ∀ (p q: Prop), p → p ∨ q := by\n  intro p q h\n  exact Or.inl h"
   let goal1 := "p q : Prop\nh : p\n⊢ p ∨ q"
-  [
-    step "frontend.process"
-      ({
-        file? := .some file,
-        invocations := true,
-      }: Protocol.FrontendProcess)
-     ({
-       units := [{
-         boundary := (0, file.utf8ByteSize),
-         invocations? := .some [
-           {
-             goalBefore := "⊢ ∀ (p q : Prop), p → p ∨ q",
-             goalAfter := goal1,
-             tactic := "intro p q h",
-             usedConstants := #[],
-           },
-           {
-             goalBefore := goal1 ,
-             goalAfter := "",
-             tactic := "exact Or.inl h",
-             usedConstants := #["Or.inl"],
-           },
-         ]
-       }],
-    }: Protocol.FrontendProcessResult),
-  ]
+  IO.FS.withTempDir λ tempdir => do
+  let filename := s!"{tempdir}/invocations.jsonl"
+  step "frontend.process"
+    ({
+      file? := .some file,
+      invocations? := .some filename,
+    }: Protocol.FrontendProcess)
+   ({
+     units := [{
+       boundary := (0, file.utf8ByteSize),
+       nInvocations? := .some 2,
+     }],
+   }: Protocol.FrontendProcessResult)
+  stepFile (α := Protocol.FrontendData) "invocations" filename
+    { units := [{
+        invocations? := .some [
+          {
+            goalBefore := "⊢ ∀ (p q : Prop), p → p ∨ q",
+            goalAfter := goal1,
+            tactic := "intro p q h",
+            usedConstants := #[],
+          },
+          {
+            goalBefore := goal1 ,
+            goalAfter := "",
+            tactic := "exact Or.inl h",
+            usedConstants := #["Or.inl"],
+          },
+        ]
+    } ] }
 
 example : 1 + 2 = 3 := rfl
 example (p: Prop): p → p := by simp
 
-def test_frontend_process_sorry : Test :=
+def test_frontend_process_sorry : Test := do
   let solved := "example : 1 + 2 = 3 := rfl\n"
   let withSorry := "example (p: Prop): p → p := sorry"
-  [
-    let file := s!"{solved}{withSorry}"
-    let goal1: Protocol.Goal := {
-      name := "_uniq.6",
-      target := { pp? := .some "p → p" },
-      vars := #[{ name := "_uniq.4", userName := "p", type? := .some { pp? := .some "Prop" }}],
-    }
-    step "frontend.process"
-      ({
-        file? := .some file,
-        sorrys := true,
-      }: Protocol.FrontendProcess)
-     ({
-       units := [{
-         boundary := (0, solved.utf8ByteSize),
-       }, {
-         boundary := (solved.utf8ByteSize, solved.utf8ByteSize + withSorry.utf8ByteSize),
-         goalStateId? := .some 0,
-         goals? := .some #[goal1],
-         goalSrcBoundaries? := .some #[(57, 62)],
-         messages := #["<anonymous>:2:0: warning: declaration uses 'sorry'\n"],
-       }],
-     }: Protocol.FrontendProcessResult),
-  ]
+  let file := s!"{solved}{withSorry}"
+  let goal1: Protocol.Goal := {
+    name := "_uniq.6",
+    target := { pp? := .some "p → p" },
+    vars := #[{ name := "_uniq.4", userName := "p", type? := .some { pp? := .some "Prop" }}],
+  }
+  step "frontend.process"
+    ({
+      file? := .some file,
+      sorrys := true,
+    }: Protocol.FrontendProcess)
+   ({
+     units := [{
+       boundary := (0, solved.utf8ByteSize),
+     }, {
+       boundary := (solved.utf8ByteSize, solved.utf8ByteSize + withSorry.utf8ByteSize),
+       goalStateId? := .some 0,
+       goals? := .some #[goal1],
+       goalSrcBoundaries? := .some #[(57, 62)],
+       messages := #["<anonymous>:2:0: warning: declaration uses 'sorry'\n"],
+     }],
+   }: Protocol.FrontendProcessResult)
 
-def test_import_open : Test :=
+def test_import_open : Test := do
   let header := "import Init\nopen Nat\nuniverse u"
   let goal1: Protocol.Goal := {
     name := "_uniq.67",
     target := { pp? := .some "n + 1 = n.succ" },
     vars := #[{ name := "_uniq.66", userName := "n", type? := .some { pp? := .some "Nat" }}],
   }
-  [
-    step "frontend.process"
-      ({
-        file? := .some header,
-        readHeader := true,
-        inheritEnv := true,
-      }: Protocol.FrontendProcess)
-     ({
-       units := [
-         { boundary := (12, 21) },
-         { boundary := (21, header.utf8ByteSize) },
-       ],
-     }: Protocol.FrontendProcessResult),
-    step "goal.start" ({ expr := "∀ (n : Nat), n + 1 = Nat.succ n"} : Protocol.GoalStart)
-     ({ stateId := 0, root := "_uniq.65" }: Protocol.GoalStartResult),
-    step "goal.tactic" ({ stateId := 0, tactic? := .some "intro n" }: Protocol.GoalTactic)
-     ({ nextStateId? := .some 1, goals? := #[goal1], }: Protocol.GoalTacticResult),
-    step "goal.tactic" ({ stateId := 1, tactic? := .some "apply add_one" }: Protocol.GoalTactic)
-     ({ nextStateId? := .some 2, goals? := .some #[], }: Protocol.GoalTacticResult),
-    step "goal.start" ({ expr := "∀ (x : Sort u), Sort (u + 1)"} : Protocol.GoalStart)
-     ({ stateId := 3, root := "_uniq.5" }: Protocol.GoalStartResult),
-  ]
+  step "frontend.process"
+    ({
+      file? := .some header,
+      readHeader := true,
+      inheritEnv := true,
+    }: Protocol.FrontendProcess)
+   ({
+     units := [
+       { boundary := (12, 21) },
+       { boundary := (21, header.utf8ByteSize) },
+     ],
+   }: Protocol.FrontendProcessResult)
+  step "goal.start" ({ expr := "∀ (n : Nat), n + 1 = Nat.succ n"} : Protocol.GoalStart)
+   ({ stateId := 0, root := "_uniq.65" }: Protocol.GoalStartResult)
+  step "goal.tactic" ({ stateId := 0, tactic? := .some "intro n" }: Protocol.GoalTactic)
+   ({ nextStateId? := .some 1, goals? := #[goal1], }: Protocol.GoalTacticResult)
+  step "goal.tactic" ({ stateId := 1, tactic? := .some "apply add_one" }: Protocol.GoalTactic)
+   ({ nextStateId? := .some 2, goals? := .some #[], }: Protocol.GoalTacticResult)
+  step "goal.start" ({ expr := "∀ (x : Sort u), Sort (u + 1)"} : Protocol.GoalStart)
+   ({ stateId := 3, root := "_uniq.5" }: Protocol.GoalStartResult)
 
 /-- Ensure there cannot be circular references -/
-def test_frontend_process_circular : Test :=
+def test_frontend_process_circular : Test := do
   let withSorry := "theorem mystery : 1 + 2 = 2 + 3 := sorry"
-  [
-    let goal1: Protocol.Goal := {
-      name := "_uniq.2",
-      target := { pp? := .some "1 + 2 = 2 + 3" },
-      vars := #[],
-    }
-    step "frontend.process"
-      ({
-        file? := .some withSorry,
-        sorrys := true,
-      }: Protocol.FrontendProcess)
-     ({
-       units := [{
-         boundary := (0, withSorry.utf8ByteSize),
-         goalStateId? := .some 0,
-         goals? := .some #[goal1],
-         goalSrcBoundaries? := .some #[(35, 40)],
-         messages := #["<anonymous>:1:8: warning: declaration uses 'sorry'\n"],
-       }],
-     }: Protocol.FrontendProcessResult),
-    step "goal.tactic" ({ stateId := 0, tactic? := .some "exact?" }: Protocol.GoalTactic)
-     ({ messages? := .some #["`exact?` could not close the goal. Try `apply?` to see partial suggestions."] }
-     : Protocol.GoalTacticResult),
-  ]
+  let goal1: Protocol.Goal := {
+    name := "_uniq.2",
+    target := { pp? := .some "1 + 2 = 2 + 3" },
+    vars := #[],
+  }
+  step "frontend.process"
+    ({
+      file? := .some withSorry,
+      sorrys := true,
+    }: Protocol.FrontendProcess)
+   ({
+     units := [{
+       boundary := (0, withSorry.utf8ByteSize),
+       goalStateId? := .some 0,
+       goals? := .some #[goal1],
+       goalSrcBoundaries? := .some #[(35, 40)],
+       messages := #["<anonymous>:1:8: warning: declaration uses 'sorry'\n"],
+     }],
+   }: Protocol.FrontendProcessResult)
+  step "goal.tactic" ({ stateId := 0, tactic? := .some "exact?" }: Protocol.GoalTactic)
+   ({ messages? := .some #["`exact?` could not close the goal. Try `apply?` to see partial suggestions."] }
+   : Protocol.GoalTacticResult)
 
-def runTest (env: Lean.Environment) (steps: Test): IO LSpec.TestSeq := do
+def runTestSuite (env : Lean.Environment) (steps : Test): IO LSpec.TestSeq := do
   -- Setup the environment for execution
   let coreContext ← createCoreContext #[]
-  let mainM : MainM LSpec.TestSeq :=
-    steps.foldlM (λ suite step => do return suite ++ (← step)) LSpec.TestSeq.done
+  let mainM : MainM LSpec.TestSeq := runTest steps
   mainM.run { coreContext } |>.run' { env }
 
 def suite (env : Lean.Environment): List (String × IO LSpec.TestSeq) :=
@@ -357,7 +349,7 @@ def suite (env : Lean.Environment): List (String × IO LSpec.TestSeq) :=
     ("frontend.process import", test_import_open),
     ("frontend.process circular", test_frontend_process_circular),
   ]
-  tests.map (fun (name, test) => (name, runTest env test))
+  tests.map (fun (name, test) => (name, runTestSuite env test))
 
 
 end Pantograph.Test.Integration