From 269e5c707c7bb182562d45cd6b57ccc2d6703b8d Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Thu, 26 Oct 2023 22:47:42 -0700
Subject: [PATCH] refactor: Separate goal printing and processing

Added a test for delta proof variables
---
 Pantograph.lean          |  10 +--
 Pantograph/Goal.lean     |  75 +++----------------
 Pantograph/Protocol.lean |   9 ++-
 Pantograph/Serial.lean   |  72 ++++++++++++++++--
 Test/Proofs.lean         | 158 +++++++++++++++++++++++----------------
 5 files changed, 181 insertions(+), 143 deletions(-)

diff --git a/Pantograph.lean b/Pantograph.lean
index 0e74e81..29f9bd5 100644
--- a/Pantograph.lean
+++ b/Pantograph.lean
@@ -110,7 +110,7 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
         printJsonPretty := args.printJsonPretty?.getD options.printJsonPretty,
         printExprPretty := args.printExprPretty?.getD options.printExprPretty,
         printExprAST := args.printExprAST?.getD options.printExprAST,
-        proofVariableDelta := args.proofVariableDelta?.getD options.proofVariableDelta,
+        noRepeat := args.noRepeat?.getD options.noRepeat,
         printAuxDecls := args.printAuxDecls?.getD options.printAuxDecls,
         printImplementationDetailHyps := args.printImplementationDetailHyps?.getD options.printImplementationDetailHyps
       }
@@ -148,14 +148,14 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
     match state.goalStates.get? args.stateId with
     | .none => return .error $ errorIndex s!"Invalid state index {args.stateId}"
     | .some goalState =>
-      let result ← GoalState.execute goalState args.goalId args.tactic |>.run state.options
-      match result with
-      | .success nextGoalState goals =>
+      match ← GoalState.execute goalState args.goalId args.tactic with
+      | .success nextGoalState =>
         let (goalStates, nextStateId) := state.goalStates.insert nextGoalState
         set { state with goalStates }
+        let goals ← nextGoalState.serializeGoals (parent := .some goalState) (options := state.options)
         return .ok {
           nextStateId? := .some nextStateId,
-          goals? := .some goals
+          goals? := .some goals,
         }
       | .parseError message =>
         return .ok { parseError? := .some message }
diff --git a/Pantograph/Goal.lean b/Pantograph/Goal.lean
index 884b8a0..1f3f71a 100644
--- a/Pantograph/Goal.lean
+++ b/Pantograph/Goal.lean
@@ -1,8 +1,6 @@
 import Lean
 
 import Pantograph.Symbol
-import Pantograph.Serial
-import Pantograph.Protocol
 
 def Lean.MessageLog.getErrorMessages (log : MessageLog) : MessageLog :=
   {
@@ -21,6 +19,9 @@ structure GoalState where
   -- New metavariables acquired in this state
   newMVars: SSet MVarId
 
+  -- The id of the goal in the parent
+  parentGoalId: Nat := 0
+
 abbrev M := Elab.TermElabM
 
 protected def GoalState.create (expr: Expr): M GoalState := do
@@ -49,6 +50,7 @@ protected def GoalState.mctx (state: GoalState): MetavarContext :=
 private def GoalState.mvars (state: GoalState): SSet MVarId :=
   state.mctx.decls.foldl (init := .empty) fun acc k _ => acc.insert k
 
+/-- Inner function for executing tactic on goal state -/
 def executeTactic (state: Elab.Tactic.SavedState) (goal: MVarId) (tactic: Syntax) :
     M (Except (Array String) (Elab.Tactic.SavedState × List MVarId)):= do
   let tacticM (stx: Syntax):  Elab.Tactic.TacticM (Except (Array String) (Elab.Tactic.SavedState × List MVarId)) := do
@@ -71,7 +73,7 @@ def executeTactic (state: Elab.Tactic.SavedState) (goal: MVarId) (tactic: Syntax
 /-- Response for executing a tactic -/
 inductive TacticResult where
   -- Goes to next state
-  | success (state: GoalState) (goals: Array Protocol.Goal)
+  | success (state: GoalState)
   -- Tactic failed with messages
   | failure (messages: Array String)
   -- Could not parse tactic
@@ -81,7 +83,7 @@ inductive TacticResult where
 
 /-- Execute tactic on given state -/
 protected def GoalState.execute (state: GoalState) (goalId: Nat) (tactic: String):
-    Protocol.OptionsT M TacticResult := do
+    M TacticResult := do
   let goal ← match state.savedState.tactic.goals.get? goalId with
     | .some goal => pure $ goal
     | .none => return .indexError goalId
@@ -92,7 +94,6 @@ protected def GoalState.execute (state: GoalState) (goalId: Nat) (tactic: String
     (fileName := "<stdin>") with
     | .ok stx => pure $ stx
     | .error error => return .parseError error
-  let options ← read
   match (← executeTactic (state := state.savedState) (goal := goal) (tactic := tactic)) with
   | .error errors =>
     return .failure errors
@@ -110,20 +111,12 @@ protected def GoalState.execute (state: GoalState) (goalId: Nat) (tactic: String
       else
         return acc.insert mvarId
       ) SSet.empty
-    let nextState: GoalState := {
+    return .success {
       savedState := nextSavedState
       root := state.root,
       newMVars,
+      parentGoalId := goalId,
     }
-    nextSavedState.term.restore
-    let parentDecl? := (← MonadMCtx.getMCtx).findDecl? goal
-    let goals ← nextGoals.mapM fun nextGoal => do
-      match (← MonadMCtx.getMCtx).findDecl? nextGoal with
-      | .some mvarDecl =>
-        let serializedGoal ← serialize_goal options mvarDecl (parentDecl? := parentDecl?)
-        return serializedGoal
-      | .none => throwError s!"Parent mvar id does not exist {nextGoal.name}"
-    return .success nextState goals.toArray
 
 /-- After finishing one branch of a proof (`graftee`), pick up from the point where the proof was left off (`target`) -/
 protected def GoalState.continue (target: GoalState) (graftee: GoalState): Except String GoalState :=
@@ -150,57 +143,11 @@ protected def GoalState.rootExpr (goalState: GoalState): Option Expr :=
   let expr := goalState.mctx.eAssignment.find! goalState.root
   let (expr, _) := instantiateMVarsCore (mctx := goalState.mctx) (e := expr)
   if expr.hasMVar then
+    -- Must not assert that the goal state is empty here. We could be in a branch goal.
+    --assert! ¬goalState.goals.isEmpty
     .none
   else
+    assert! goalState.goals.isEmpty
     .some expr
 
--- Diagnostics functions
-
-/-- Print the metavariables in a readable format -/
-protected def GoalState.print (goalState: GoalState) (options: Protocol.GoalPrint := {}): M Unit := do
-  let savedState := goalState.savedState
-  savedState.term.restore
-  let goals := savedState.tactic.goals
-  let mctx ← getMCtx
-  let root := goalState.root
-  -- Print the root
-  match mctx.decls.find? root with
-  | .some decl => printMVar ">" root decl
-  | .none => IO.println s!">{root.name}: ??"
-  goals.forM (fun mvarId => do
-    if mvarId != root then
-      match mctx.decls.find? mvarId with
-      | .some decl => printMVar "⊢" mvarId decl
-      | .none => IO.println s!"⊢{mvarId.name}: ??"
-  )
-  let goals := goals.toSSet
-  mctx.decls.forM (fun mvarId decl => do
-    if goals.contains mvarId || mvarId == root then
-      pure ()
-    -- Always print the root goal
-    else if mvarId == goalState.root then
-      printMVar (pref := ">") mvarId decl
-    -- Print the remainig ones that users don't see in Lean
-    else if options.printNonVisible then
-      let pref := if goalState.newMVars.contains mvarId then "~" else " "
-      printMVar pref mvarId decl
-    else
-      pure ()
-      --IO.println s!" {mvarId.name}{userNameToString decl.userName}"
-  )
-  where
-    printMVar (pref: String) (mvarId: MVarId) (decl: MetavarDecl): Elab.TermElabM Unit := do
-      if options.printContext then
-        decl.lctx.fvarIdToDecl.forM printFVar
-      let type_sexp := serialize_expression_ast (← instantiateMVars decl.type) (sanitize := false)
-      IO.println s!"{pref}{mvarId.name}{userNameToString decl.userName}: {← Meta.ppExpr decl.type} {type_sexp}"
-      if options.printValue then
-        if let Option.some value := (← getMCtx).eAssignment.find? mvarId then
-          IO.println s!"  = {← Meta.ppExpr value}"
-    printFVar (fvarId: FVarId) (decl: LocalDecl): Elab.TermElabM Unit := do
-      IO.println s!" | {fvarId.name}{userNameToString decl.userName}: {← Meta.ppExpr decl.type}"
-    userNameToString : Name → String
-      | .anonymous => ""
-      | other => s!"[{other}]"
-
 end Pantograph
diff --git a/Pantograph/Protocol.lean b/Pantograph/Protocol.lean
index a6bae29..1c05227 100644
--- a/Pantograph/Protocol.lean
+++ b/Pantograph/Protocol.lean
@@ -18,9 +18,10 @@ structure Options where
   printExprPretty: Bool    := true
   -- When enabled, print the raw AST of expressions
   printExprAST: Bool       := false
-  -- When enabled, the types and values of persistent variables in a proof goal
-  -- are not shown unless they are new to the proof step. Reduces overhead
-  proofVariableDelta: Bool := false
+  -- When enabled, the types and values of persistent variables in a goal
+  -- are not shown unless they are new to the proof step. Reduces overhead.
+  -- NOTE: that this assumes the type and assignment of variables can never change.
+  noRepeat: Bool := false
   -- See `pp.auxDecls`
   printAuxDecls: Bool      := false
   -- See `pp.implementationDetailHyps`
@@ -123,7 +124,7 @@ structure OptionsSet where
   printJsonPretty?: Option Bool
   printExprPretty?: Option Bool
   printExprAST?: Option Bool
-  proofVariableDelta?: Option Bool
+  noRepeat?: Option Bool
   printAuxDecls?: Option Bool
   printImplementationDetailHyps?: Option Bool
   deriving Lean.FromJson
diff --git a/Pantograph/Serial.lean b/Pantograph/Serial.lean
index 99f95ef..87552eb 100644
--- a/Pantograph/Serial.lean
+++ b/Pantograph/Serial.lean
@@ -4,6 +4,7 @@ All serialisation functions
 import Lean
 
 import Pantograph.Protocol
+import Pantograph.Goal
 
 namespace Pantograph
 open Lean
@@ -173,12 +174,12 @@ def serialize_expression_ast (expr: Expr) (sanitize: Bool := true): String :=
 def serialize_expression (options: Protocol.Options) (e: Expr): MetaM Protocol.Expression := do
   let pp := toString (← Meta.ppExpr e)
   let pp?: Option String := match options.printExprPretty with
-      | true => .some pp
-      | false => .none
+    | true => .some pp
+    | false => .none
   let sexp: String := serialize_expression_ast e
   let sexp?: Option String := match options.printExprAST with
-      | true => .some sexp
-      | false => .none
+    | true => .some sexp
+    | false => .none
   return {
     pp?,
     sexp?
@@ -239,7 +240,7 @@ def serialize_goal (options: Protocol.Options) (mvarDecl: MetavarDecl) (parentDe
       if skip then
         return acc
       else
-        let nameOnly := options.proofVariableDelta && (parentDecl?.map
+        let nameOnly := options.noRepeat && (parentDecl?.map
           (λ decl => decl.lctx.find? localDecl.fvarId |>.isSome) |>.getD false)
         let var ← match nameOnly with
           | true => ppVarNameOnly localDecl
@@ -254,6 +255,67 @@ def serialize_goal (options: Protocol.Options) (mvarDecl: MetavarDecl) (parentDe
   where
   of_name (n: Name) := name_to_ast n (sanitize := false)
 
+protected def GoalState.serializeGoals (state: GoalState) (parent: Option GoalState := .none) (options: Protocol.Options := {}): MetaM (Array Protocol.Goal):= do
+  let goals := state.goals.toArray
+  state.savedState.term.meta.restore
+  let parentDecl? := parent.bind (λ parentState =>
+    let parentGoal := parentState.goals.get! state.parentGoalId
+    parentState.mctx.findDecl? parentGoal)
+  goals.mapM fun goal => do
+    if options.noRepeat then
+      let key := if parentDecl?.isSome then "is some" else "is none"
+      IO.println s!"goal: {goal.name}, {key}"
+    match state.mctx.findDecl? goal with
+    | .some mvarDecl =>
+      let serializedGoal ← serialize_goal options mvarDecl (parentDecl? := parentDecl?)
+      pure serializedGoal
+    | .none => throwError s!"Metavariable does not exist in context {goal.name}"
 
+/-- Print the metavariables in a readable format -/
+protected def GoalState.print (goalState: GoalState) (options: Protocol.GoalPrint := {}): MetaM Unit := do
+  let savedState := goalState.savedState
+  savedState.term.meta.restore
+  let goals := savedState.tactic.goals
+  let mctx ← getMCtx
+  let root := goalState.root
+  -- Print the root
+  match mctx.decls.find? root with
+  | .some decl => printMVar ">" root decl
+  | .none => IO.println s!">{root.name}: ??"
+  goals.forM (fun mvarId => do
+    if mvarId != root then
+      match mctx.decls.find? mvarId with
+      | .some decl => printMVar "⊢" mvarId decl
+      | .none => IO.println s!"⊢{mvarId.name}: ??"
+  )
+  let goals := goals.toSSet
+  mctx.decls.forM (fun mvarId decl => do
+    if goals.contains mvarId || mvarId == root then
+      pure ()
+    -- Always print the root goal
+    else if mvarId == goalState.root then
+      printMVar (pref := ">") mvarId decl
+    -- Print the remainig ones that users don't see in Lean
+    else if options.printNonVisible then
+      let pref := if goalState.newMVars.contains mvarId then "~" else " "
+      printMVar pref mvarId decl
+    else
+      pure ()
+      --IO.println s!" {mvarId.name}{userNameToString decl.userName}"
+  )
+  where
+    printMVar (pref: String) (mvarId: MVarId) (decl: MetavarDecl): MetaM Unit := do
+      if options.printContext then
+        decl.lctx.fvarIdToDecl.forM printFVar
+      let type_sexp := serialize_expression_ast (← instantiateMVars decl.type) (sanitize := false)
+      IO.println s!"{pref}{mvarId.name}{userNameToString decl.userName}: {← Meta.ppExpr decl.type} {type_sexp}"
+      if options.printValue then
+        if let Option.some value := (← getMCtx).eAssignment.find? mvarId then
+          IO.println s!"  = {← Meta.ppExpr value}"
+    printFVar (fvarId: FVarId) (decl: LocalDecl): MetaM Unit := do
+      IO.println s!" | {fvarId.name}{userNameToString decl.userName}: {← Meta.ppExpr decl.type}"
+    userNameToString : Name → String
+      | .anonymous => ""
+      | other => s!"[{other}]"
 
 end Pantograph
diff --git a/Test/Proofs.lean b/Test/Proofs.lean
index a3d26fe..79f0f38 100644
--- a/Test/Proofs.lean
+++ b/Test/Proofs.lean
@@ -9,7 +9,7 @@ import Test.Common
 namespace Pantograph
 
 def TacticResult.toString : TacticResult → String
-  | .success _ goals => s!".success ({goals.size} goals)"
+  | .success state => s!".success ({state.goals.length} goals)"
   | .failure messages =>
     let messages := "\n".intercalate messages.toList
     s!".failure {messages}"
@@ -75,16 +75,6 @@ def buildGoal (nameType: List (String × String)) (target: String): Protocol.Goa
       isInaccessible? := .some false
     })).toArray
   }
--- Like `buildGoal` but allow certain variables to be elided.
-def buildGoalSelective (nameType: List (String × Option String)) (target: String): Protocol.Goal :=
-  {
-    target := { pp? := .some target},
-    vars := (nameType.map fun x => ({
-      userName := x.fst,
-      type? := x.snd.map (λ type => { pp? := type }),
-      isInaccessible? := x.snd.map (λ _ => false)
-    })).toArray
-  }
 def proofRunner (env: Lean.Environment) (tests: TestM Unit): IO LSpec.TestSeq := do
   let termElabM := tests.run LSpec.TestSeq.done |>.run {} -- with default options
 
@@ -120,12 +110,13 @@ def proof_nat_add_comm (manual: Bool): TestM Unit := do
       addTest $ assertUnreachable "Goal could not parse"
       return ()
 
-  let (state1, goal1) ← match ← state0.execute (goalId := 0) (tactic := "intro n m") with
-    | .success state #[goal] => pure (state, goal)
+  let state1 ← match ← state0.execute (goalId := 0) (tactic := "intro n m") with
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
-  addTest $ LSpec.check "intro n m" (goal1.devolatilize = buildGoal [("n", "Nat"), ("m", "Nat")] "n + m = m + n")
+  addTest $ LSpec.check "intro n m" ((← state1.serializeGoals (options := ← read)).map (·.devolatilize) =
+    #[buildGoal [("n", "Nat"), ("m", "Nat")] "n + m = m + n"])
 
   match ← state1.execute (goalId := 0) (tactic := "assumption") with
   | .failure #[message] =>
@@ -134,12 +125,49 @@ def proof_nat_add_comm (manual: Bool): TestM Unit := do
     addTest $ assertUnreachable $ other.toString
 
   let state2 ← match ← state1.execute (goalId := 0) (tactic := "rw [Nat.add_comm]") with
-    | .success state #[] => pure state
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
+  addTest $ LSpec.test "rw [Nat.add_comm]" state2.goals.isEmpty
 
   return ()
+def proof_delta_variable: TestM Unit := do
+  let options: Protocol.Options := { noRepeat := true }
+  let state? ← startProof <| .expr "∀ (a b: Nat), a + b = b + a"
+  addTest $ LSpec.check "Start goal" state?.isSome
+  let state0 ← match state? with
+    | .some state => pure state
+    | .none => do
+      addTest $ assertUnreachable "Goal could not parse"
+      return ()
+
+  let state1 ← match ← state0.execute (goalId := 0) (tactic := "intro n") with
+    | .success state => pure state
+    | other => do
+      addTest $ assertUnreachable $ other.toString
+      return ()
+  addTest $ LSpec.check "intro n" ((← state1.serializeGoals (parent := state0) options).map (·.devolatilize) =
+    #[buildGoalSelective [("n", .some "Nat")] "∀ (b : Nat), n + b = b + n"])
+  let state2 ← match ← state1.execute (goalId := 0) (tactic := "intro m") with
+    | .success state => pure state
+    | other => do
+      addTest $ assertUnreachable $ other.toString
+      return ()
+  addTest $ LSpec.check "intro m" ((← state2.serializeGoals (parent := state1) options).map (·.devolatilize) =
+    #[buildGoalSelective [("n", .none), ("m", .some "Nat")] "n + m = m + n"])
+  return ()
+  where
+-- Like `buildGoal` but allow certain variables to be elided.
+  buildGoalSelective (nameType: List (String × Option String)) (target: String): Protocol.Goal :=
+    {
+      target := { pp? := .some target},
+      vars := (nameType.map fun x => ({
+        userName := x.fst,
+        type? := x.snd.map (λ type => { pp? := type }),
+        isInaccessible? := x.snd.map (λ _ => false)
+      })).toArray
+    }
 
 example (w x y z : Nat) (p : Nat → Prop)
         (h : p (x * y + z * w * x)) : p (x * w * z + y * x) := by
@@ -153,23 +181,26 @@ def proof_arith: TestM Unit := do
       addTest $ assertUnreachable "Goal could not parse"
       return ()
 
-  let (state1, goal) ← match ← state0.execute (goalId := 0) (tactic := "intros") with
-    | .success state #[goal] => pure (state, goal)
+  let state1 ← match ← state0.execute (goalId := 0) (tactic := "intros") with
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
-  addTest $ LSpec.check "1 root" state1.rootExpr.isNone
-  let (state2, goal) ← match ← state1.execute (goalId := 0) (tactic := "simp [Nat.add_assoc, Nat.add_comm, Nat.add_left_comm, Nat.mul_comm, Nat.mul_assoc, Nat.mul_left_comm] at *") with
-    | .success state #[goal] => pure (state, goal)
+  addTest $ LSpec.check "intros" (state1.goals.length = 1)
+  addTest $ LSpec.test "1 root" state1.rootExpr.isNone
+  let state2 ← match ← state1.execute (goalId := 0) (tactic := "simp [Nat.add_assoc, Nat.add_comm, Nat.add_left_comm, Nat.mul_comm, Nat.mul_assoc, Nat.mul_left_comm] at *") with
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
+  addTest $ LSpec.check "simp ..." (state2.goals.length = 1)
   addTest $ LSpec.check "2 root" state2.rootExpr.isNone
   let state3 ← match ← state2.execute (goalId := 0) (tactic := "assumption") with
-    | .success state #[] => pure state
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
+  addTest $ LSpec.test "assumption" state3.goals.isEmpty
   addTest $ LSpec.check "3 root" state3.rootExpr.isSome
   return ()
 
@@ -196,59 +227,66 @@ def proof_or_comm: TestM Unit := do
       addTest $ assertUnreachable "Goal could not parse"
       return ()
 
-  let (state1, goal1) ← match ← state0.execute (goalId := 0) (tactic := "intro p q h") with
-    | .success state #[goal] => pure (state, goal)
+  let state1 ← match ← state0.execute (goalId := 0) (tactic := "intro p q h") with
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
-  addTest $ LSpec.check "p q h" (goal1.devolatilize = buildGoal [("p", "Prop"), ("q", "Prop"), ("h", "p ∨ q")] "q ∨ p")
-  let (state2, goal1, goal2) ← match ← state1.execute (goalId := 0) (tactic := "cases h") with
-    | .success state #[goal1, goal2] => pure (state, goal1, goal2)
+  addTest $ LSpec.check "intro n m" ((← state1.serializeGoals (options := ← read)).map (·.devolatilize) =
+    #[buildGoal [("p", "Prop"), ("q", "Prop"), ("h", "p ∨ q")] "q ∨ p"])
+  let state2 ← match ← state1.execute (goalId := 0) (tactic := "cases h") with
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
-  addTest $ LSpec.check "cases h/1" (goal1.devolatilize = branchGoal "inl" "p")
-  addTest $ LSpec.check "cases h/2" (goal2.devolatilize = branchGoal "inr" "q")
+  addTest $ LSpec.check "cases h" ((← state2.serializeGoals (options := ← read)).map (·.devolatilize) =
+    #[branchGoal "inl" "p", branchGoal "inr" "q"])
 
-  let (state3_1, _goal) ← match ← state2.execute (goalId := 0) (tactic := "apply Or.inr") with
-    | .success state #[goal] => pure (state, goal)
+  let state3_1 ← match ← state2.execute (goalId := 0) (tactic := "apply Or.inr") with
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
+  addTest $ LSpec.check "· apply Or.inr" (state3_1.goals.length = 1)
   let state4_1 ← match ← state3_1.execute (goalId := 0) (tactic := "assumption") with
-    | .success state #[] => pure state
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
-  let (state3_2, _goal) ← match ← state2.execute (goalId := 1) (tactic := "apply Or.inl") with
-    | .success state #[goal] => pure (state, goal)
+  addTest $ LSpec.check "· assumption" state4_1.goals.isEmpty
+  addTest $ LSpec.check "(4_1 root)" state4_1.rootExpr.isNone
+  let state3_2 ← match ← state2.execute (goalId := 1) (tactic := "apply Or.inl") with
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
+  addTest $ LSpec.check "· apply Or.inl" (state3_2.goals.length = 1)
   let state4_2 ← match ← state3_2.execute (goalId := 0) (tactic := "assumption") with
-    | .success state #[] => pure state
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
-
-  addTest $ LSpec.check "4_2 root" state4_2.rootExpr.isNone
+  addTest $ LSpec.check "· assumption" state4_2.goals.isEmpty
+  addTest $ LSpec.check "(4_2 root)" state4_2.rootExpr.isNone
   -- Ensure the proof can continue from `state4_2`.
   let state2b ← match state2.continue state4_2 with
     | .error msg => do
       addTest $ assertUnreachable $ msg
       return ()
     | .ok state => pure state
-  addTest $ LSpec.test "state2b" (state2b.goals == [state2.goals.get! 0])
-  let (state3_1, _goal) ← match ← state2b.execute (goalId := 0) (tactic := "apply Or.inr") with
-    | .success state #[goal] => pure (state, goal)
+  addTest $ LSpec.test "(resume)" (state2b.goals == [state2.goals.get! 0])
+  let state3_1 ← match ← state2b.execute (goalId := 0) (tactic := "apply Or.inr") with
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
+  addTest $ LSpec.check "· apply Or.inr" (state3_1.goals.length = 1)
   let state4_1 ← match ← state3_1.execute (goalId := 0) (tactic := "assumption") with
-    | .success state #[] => pure state
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
+  addTest $ LSpec.check "· assumption" state4_1.goals.isEmpty
   addTest $ LSpec.check "4_1 root" state4_1.rootExpr.isSome
 
   return ()
@@ -273,41 +311,30 @@ def proof_m_couple: TestM Unit := do
       addTest $ assertUnreachable "Goal could not parse"
       return ()
 
-  let (state1, goalL, goalR, goalM) ← match ← state0.execute (goalId := 0) (tactic := "apply Nat.le_trans") with
-    | .success state #[goalL, goalR, goalM] => pure (state, goalL, goalR, goalM)
+  let state1 ← match ← state0.execute (goalId := 0) (tactic := "apply Nat.le_trans") with
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
-  addTest $ LSpec.test "2 ≤ ?m" (goalL.target.pp? = .some "2 ≤ ?m")
-  addTest $ LSpec.test "?m ≤ 5" (goalR.target.pp? = .some "?m ≤ 5")
-  addTest $ LSpec.test "Nat" (goalM.target.pp? = .some "Nat")
+  addTest $ LSpec.check "apply Nat.le_trans" ((← state1.serializeGoals (options := ← read)).map (·.target.pp?) =
+    #[.some "2 ≤ ?m", .some "?m ≤ 5", .some "Nat"])
+  addTest $ LSpec.test "(1 root)" state1.rootExpr.isNone
   -- Set m to 3
   let state2 ← match ← state1.execute (goalId := 2) (tactic := "exact 3") with
-    | .success state #[] => pure state
+    | .success state => pure state
     | other => do
       addTest $ assertUnreachable $ other.toString
       return ()
+  addTest $ LSpec.test "(1b root)" state2.rootExpr.isNone
   let state1b ← match state1.continue state2 with
-    | .ok state => pure state
-    | .error error => do
-      addTest $ assertUnreachable $ error
+    | .error msg => do
+      addTest $ assertUnreachable $ msg
       return ()
-  state1b.print
---def proof_delta_variable: TestM Unit := withReader (fun _ => {proofVariableDelta := true}) do
---  let goal? ← startProof (.expr "∀ (a b: Nat), a + b = b + a")
---  addTest $ LSpec.check "Start goal" goal?.isSome
---  if let .some goal := goal? then
---    if let .success #[(goal, sGoal)] ← goal.execute "intro n" then
---      let sGoal1e: Protocol.Goal :=buildGoalSelective [("n", .some "Nat")] "∀ (b : Nat), n + b = b + n"
---      addTest $ LSpec.check "intro n" (sGoal = sGoal1e)
---
---      if let .success #[(_, sGoal)] ← goal.execute "intro m" then
---        let sGoal2e: Protocol.Goal :=buildGoalSelective [("n", .none), ("m", .some "Nat")] "n + m = m + n"
---        addTest $ LSpec.check "intro m" (sGoal = sGoal2e)
---      else
---        addTest $ assertUnreachable "intro m"
---    else
---      addTest $ assertUnreachable "intro n"
+    | .ok state => pure state
+  addTest $ LSpec.check "exact 3" ((← state1b.serializeGoals (options := ← read)).map (·.target.pp?) =
+    #[.some "2 ≤ 3", .some "3 ≤ 5"])
+  addTest $ LSpec.test "(2 root)" state1b.rootExpr.isNone
+  return ()
 
 /-- Tests the most basic form of proofs whose goals do not relate to each other -/
 def suite: IO LSpec.TestSeq := do
@@ -318,6 +345,7 @@ def suite: IO LSpec.TestSeq := do
   let tests := [
     ("Nat.add_comm", proof_nat_add_comm false),
     ("Nat.add_comm manual", proof_nat_add_comm true),
+    ("Nat.add_comm delta", proof_delta_variable),
     ("arithmetic", proof_arith),
     ("Or.comm", proof_or_comm),
     ("2 < 5", proof_m_couple)