refactor: `conv` and `calc` into tactic fragments

Pantograph/Goal.lean

@@ -20,14 +20,10 @@ structure GoalState where
   root: MVarId
 
   -- Parent state metavariable source
-  parentMVar?: Option MVarId
+  parentMVar?: Option MVarId := .none
 
-  -- Existence of this field shows that we are currently in `conv` mode.
+  -- If this is not `.none` there is a partial tactic being executed
-  -- (convRhs, goal, dormant)
+  fragment? : Option TacticFragment := .none
-  convMVar?: Option (MVarId × MVarId × List MVarId) := .none
-  -- Previous RHS for calc, so we don't have to repeat it every time
-  -- WARNING: If using `state with` outside of `calc`, this must be set to `.none`
-  calcPrevRhs?: Option (MVarId × Expr) := .none
 
 @[export pantograph_goal_state_create_m]
 protected def GoalState.create (expr: Expr): Elab.TermElabM GoalState := do
@@ -53,9 +49,6 @@ protected def GoalState.createFromMVars (goals: List MVarId) (root: MVarId): Met
     savedState,
     parentMVar? := .none,
   }
-@[export pantograph_goal_state_is_conv]
-protected def GoalState.isConv (state: GoalState): Bool :=
-  state.convMVar?.isSome
 protected def GoalState.goals (state: GoalState): List MVarId :=
   state.savedState.tactic.goals
 @[export pantograph_goal_state_goals]
@@ -101,15 +94,17 @@ private def GoalState.restoreTacticM (state: GoalState) (goal: MVarId): Elab.Tac
   Elab.Tactic.setGoals [goal]
 
 @[export pantograph_goal_state_focus]
-protected def GoalState.focus (state: GoalState) (goalId: Nat): Option GoalState := do
+protected def GoalState.focus (state : GoalState) (goal : MVarId): Option GoalState := do
-  let goal ← state.savedState.tactic.goals[goalId]?
+  let fragment? := match state.fragment? with
+    | .some { goal := goal', .. } => if goal' == goal then state.fragment? else .none
+    | .none => .none
   return {
     state with
     savedState := {
       state.savedState with
       tactic := { goals := [goal] },
     },
-    calcPrevRhs? := .none,
+    fragment?,
   }
 
 /-- Immediately bring all parent goals back into scope. Used in automatic mode -/
@@ -336,7 +331,7 @@ protected def GoalState.replay (dst : GoalState) (src src' : GoalState) : CoreM
         throwError "Conflicting assignment of expr metavariable (d != d) {mvarId.name}"
 
     Meta.saveState
-  -- FIXME: Handle calc goals
+  -- FIXME: Handle fragments
   let goals :=dst.savedState.tactic.goals ++
     src'.savedState.tactic.goals.map (⟨mapId ·.name⟩)
   return {
@@ -384,12 +379,12 @@ Set `guardMVarErrors` to true to capture mvar errors. Lean will not
 automatically collect mvars from text tactics (vide
 `test_tactic_failure_synthesize_placeholder`)
 -/
-protected def GoalState.step (state: GoalState) (goal: MVarId) (tacticM: Elab.Tactic.TacticM Unit) (guardMVarErrors : Bool := false)
+protected def GoalState.step' { α } (state: GoalState) (goal: MVarId) (tacticM: Elab.Tactic.TacticM α) (guardMVarErrors : Bool := false)
-  : Elab.TermElabM GoalState := do
+  : Elab.TermElabM (α × GoalState) := do
   unless (← getMCtx).decls.contains goal do
     throwError s!"Goal is not in context: {goal.name}"
   goal.checkNotAssigned `GoalState.step
-  let (_, { goals }) ← tacticM { elaborator := .anonymous } |>.run { goals := [goal] }
+  let (a, { goals }) ← tacticM { elaborator := .anonymous } |>.run { goals := [goal] }
   let nextElabState ← MonadBacktrack.saveState
   --Elab.Term.synthesizeSyntheticMVarsNoPostponing
 
@@ -397,12 +392,15 @@ protected def GoalState.step (state: GoalState) (goal: MVarId) (tacticM: Elab.Ta
       pure $ mergeMVarLists goals (← collectAllErroredMVars goal)
     else
       pure goals
-  return {
+  let state' := {
     state with
     savedState := { term := nextElabState, tactic := { goals }, },
     parentMVar? := .some goal,
-    calcPrevRhs? := .none,
   }
+  return (a, state')
+protected def GoalState.step (state: GoalState) (goal: MVarId) (tacticM: Elab.Tactic.TacticM Unit) (guardMVarErrors : Bool := false)
+  : Elab.TermElabM GoalState :=
+  Prod.snd <$> GoalState.step' state goal tacticM guardMVarErrors
 
 /-- Response for executing a tactic -/
 inductive TacticResult where
@@ -422,22 +420,18 @@ private def dumpMessageLog (prevMessageLength : Nat) : CoreM (Bool × Array Stri
   Core.resetMessageLog
   return (hasErrors, newMessages.toArray)
 
-/-- Executes a `TacticM` monad on this `GoalState`, collecting the errors as necessary -/
+def withCapturingError (elabM : Elab.Term.TermElabM GoalState) : Elab.TermElabM TacticResult := do
-protected def GoalState.tryTacticM
+  -- FIXME: Maybe message log should be empty
-    (state: GoalState) (goal: MVarId) (tacticM: Elab.Tactic.TacticM Unit)
+  let prevMessageLength := (← Core.getMessageLog).toList.length
-    (guardMVarErrors : Bool := false)
-    : Elab.TermElabM TacticResult := do
-  assert! ¬ (← goal.isAssigned)
-  let prevMessageLength := state.coreState.messages.toList.length
   try
-    let nextState ← state.step goal tacticM guardMVarErrors
+    let state ← elabM
 
     -- Check if error messages have been generated in the core.
     let (hasError, newMessages) ← dumpMessageLog prevMessageLength
     if hasError then
       return .failure newMessages
     else
-      return .success nextState newMessages
+      return .success state newMessages
   catch exception =>
     match exception with
     | .internal _ =>
@@ -445,19 +439,34 @@ protected def GoalState.tryTacticM
       return .failure messages
     | _ => return .failure #[← exception.toMessageData.toString]
 
+/-- Executes a `TacticM` monad on this `GoalState`, collecting the errors as necessary -/
+protected def GoalState.tryTacticM
+    (state: GoalState) (goal: MVarId) (tacticM: Elab.Tactic.TacticM Unit)
+    (guardMVarErrors : Bool := false)
+    : Elab.TermElabM TacticResult := do
+  withCapturingError do
+    state.step goal tacticM guardMVarErrors
+
 /-- Execute a string tactic on given state. Restores TermElabM -/
 @[export pantograph_goal_state_try_tactic_m]
 protected def GoalState.tryTactic (state: GoalState) (goal: MVarId) (tactic: String):
     Elab.TermElabM TacticResult := do
   state.restoreElabM
+  if let .some { goal := goal', fragment } := state.fragment? then
+    if goal == goal' then
+      return ← withCapturingError do
+        let (fragment?', state') ← state.step' goal (fragment.step goal tactic)
+        return { state' with fragment? := fragment?' }
   let tactic ← match Parser.runParserCategory
     (env := ← MonadEnv.getEnv)
-    (catName := if state.isConv then `conv else `tactic)
+    (catName := `tactic)
     (input := tactic)
     (fileName := ← getFileName) with
     | .ok stx => pure $ stx
     | .error error => return .parseError error
-  state.tryTacticM goal (Elab.Tactic.evalTactic tactic) true
+  let tacticM := Elab.Tactic.evalTactic tactic
+  withCapturingError do
+    state.step goal tacticM (guardMVarErrors := true)
 
 protected def GoalState.tryAssign (state: GoalState) (goal: MVarId) (expr: String):
       Elab.TermElabM TacticResult := do
@@ -488,145 +497,53 @@ protected def GoalState.tryLet (state: GoalState) (goal: MVarId) (binderName: St
 /-- Enter conv tactic mode -/
 protected def GoalState.conv (state: GoalState) (goal: MVarId):
       Elab.TermElabM TacticResult := do
-  if state.convMVar?.isSome then
+  if state.fragment? matches .some { fragment := .conv .., .. } then
     return .invalidAction "Already in conv state"
   goal.checkNotAssigned `GoalState.conv
-  let tacticM :  Elab.Tactic.TacticM (Elab.Tactic.SavedState × MVarId) := do
+  withCapturingError do
-    state.restoreTacticM goal
+    let (fragment, state') ← state.step' goal Fragment.enterConv
-
+    let fragment? := .some {
-    -- See Lean.Elab.Tactic.Conv.convTarget
+      goal := state'.goals[0]!,
-    let convMVar ← Elab.Tactic.withMainContext do
+      fragment,
-      let (rhs, newGoal) ← Elab.Tactic.Conv.mkConvGoalFor (← Elab.Tactic.getMainTarget)
+    }
-      Elab.Tactic.replaceMainGoal [newGoal.mvarId!]
+    return {
-      pure rhs.mvarId!
+      state' with
-    return (← MonadBacktrack.saveState, convMVar)
+      fragment?
-  try
+    }
-    let (nextSavedState, convRhs) ← tacticM { elaborator := .anonymous } |>.run' state.savedState.tactic
-    -- Other goals are now dormant
-    let otherGoals := state.goals.filter $ λ g => g != goal
-    return .success {
-      root := state.root,
-      savedState := nextSavedState
-      parentMVar? := .some goal,
-      convMVar? := .some (convRhs, goal, otherGoals),
-      calcPrevRhs? := .none
-    } #[]
-  catch exception =>
-    return .failure #[← exception.toMessageData.toString]
 
 /-- Exit from `conv` mode. Resumes all goals before the mode starts and applys the conv -/
 @[export pantograph_goal_state_conv_exit_m]
 protected def GoalState.convExit (state: GoalState):
       Elab.TermElabM TacticResult := do
-  let (convRhs, convGoal, _) ← match state.convMVar? with
+  let .some { goal, fragment } := state.fragment? |
-    | .some mvar => pure mvar
+    return .invalidAction "Not in conv state"
-    | .none => return .invalidAction "Not in conv state"
+  unless fragment matches .conv .. do
-  let tacticM :  Elab.Tactic.TacticM Elab.Tactic.SavedState:= do
+    return .invalidAction "Not in conv state"
-    -- Vide `Lean.Elab.Tactic.Conv.convert`
+  withCapturingError do
-    state.savedState.restore
+    let state' ← state.step goal (fragment.exit goal)
-
+    return {
-    -- Close all existing goals with `refl`
+      state' with
-    for mvarId in (← Elab.Tactic.getGoals) do
+      fragment? := .none,
-      liftM <| mvarId.refl <|> mvarId.inferInstance <|> pure ()
+    }
-    Elab.Tactic.pruneSolvedGoals
-    unless (← Elab.Tactic.getGoals).isEmpty do
-      throwError "convert tactic failed, there are unsolved goals\n{Elab.goalsToMessageData (← Elab.Tactic.getGoals)}"
-
-    Elab.Tactic.setGoals [convGoal]
-
-    let targetNew ← instantiateMVars (.mvar convRhs)
-    let proof ← instantiateMVars (.mvar convGoal)
-
-    Elab.Tactic.liftMetaTactic1 fun mvarId => mvarId.replaceTargetEq targetNew proof
-    MonadBacktrack.saveState
-  try
-    let nextSavedState ← tacticM { elaborator := .anonymous } |>.run' state.savedState.tactic
-    return .success {
-      root := state.root,
-      savedState := nextSavedState
-      parentMVar? := .some convGoal,
-      convMVar? := .none
-      calcPrevRhs? := .none
-    } #[]
-  catch exception =>
-    return .failure #[← exception.toMessageData.toString]
 
 protected def GoalState.calcPrevRhsOf? (state: GoalState) (goal: MVarId): Option Expr := do
-  let (mvarId, rhs) ← state.calcPrevRhs?
+  match state.fragment? with
-  if mvarId == goal then
+  | .some { goal := goal', fragment := .calc prevRhs? } =>
-    .some rhs
+    if goal == goal' then prevRhs? else .none
-  else
+  | .some _ => unreachable!
-    .none
+  | .none => .none
+
 @[export pantograph_goal_state_try_calc_m]
 protected def GoalState.tryCalc (state: GoalState) (goal: MVarId) (pred: String):
       Elab.TermElabM TacticResult := do
-  state.restoreElabM
+  let prevRhs? := state.calcPrevRhsOf? goal
-  if state.convMVar?.isSome then
+  withCapturingError do
-    return .invalidAction "Cannot initiate `calc` while in `conv` state"
+    let (fragment?, state') ← state.step' goal do
-  let `(term|$pred) ← match Parser.runParserCategory
+      let fragment := Fragment.calc prevRhs?
-    (env := state.env)
+      fragment.step goal pred
-    (catName := `term)
+    return {
-    (input := pred)
+      state' with
-    (fileName := ← getFileName) with
+      fragment?,
-    | .ok syn => pure syn
+    }
-    | .error error => return .parseError error
-  goal.checkNotAssigned `GoalState.tryCalc
-  let calcPrevRhs? := state.calcPrevRhsOf? goal
-  let decl ← goal.getDecl
-  let target ← instantiateMVars decl.type
-  let tag := decl.userName
-  try
-    goal.withContext do
-
-    let mut step ← Elab.Term.elabType <| ← do
-      if let some prevRhs := calcPrevRhs? then
-        Elab.Term.annotateFirstHoleWithType pred (← Meta.inferType prevRhs)
-      else
-        pure pred
-
-    let some (_, lhs, rhs) ← Elab.Term.getCalcRelation? step |
-      throwErrorAt pred "invalid 'calc' step, relation expected{indentExpr step}"
-    if let some prevRhs := calcPrevRhs? then
-      unless ← Meta.isDefEqGuarded lhs prevRhs do
-        throwErrorAt pred "invalid 'calc' step, left-hand-side is{indentD m!"{lhs} : {← Meta.inferType lhs}"}\nprevious right-hand-side is{indentD m!"{prevRhs} : {← Meta.inferType prevRhs}"}"
-
-    -- Creates a mvar to represent the proof that the calc tactic solves the
-    -- current branch
-    -- In the Lean `calc` tactic this is gobbled up by
-    -- `withCollectingNewGoalsFrom`
-    let mut proof ← Meta.mkFreshExprMVarAt (← getLCtx) (← Meta.getLocalInstances) step
-      (userName := tag ++ `calc)
-    let mvarBranch := proof.mvarId!
-
-    let mut proofType ← Meta.inferType proof
-    let mut remainder? := Option.none
-
-    -- The calc tactic either solves the main goal or leaves another relation.
-    -- Replace the main goal, and save the new goal if necessary
-    unless ← Meta.isDefEq proofType target do
-      let rec throwFailed :=
-        throwError "'calc' tactic failed, has type{indentExpr proofType}\nbut it is expected to have type{indentExpr target}"
-      let some (_, _, rhs) ← Elab.Term.getCalcRelation? proofType | throwFailed
-      let some (r, _, rhs') ← Elab.Term.getCalcRelation? target | throwFailed
-      let lastStep := mkApp2 r rhs rhs'
-      let lastStepGoal ← Meta.mkFreshExprSyntheticOpaqueMVar lastStep tag
-      (proof, proofType) ← Elab.Term.mkCalcTrans proof proofType lastStepGoal lastStep
-      unless ← Meta.isDefEq proofType target do throwFailed
-      remainder? := .some lastStepGoal.mvarId!
-    goal.assign proof
-
-    let goals := [ mvarBranch ] ++ remainder?.toList
-    let calcPrevRhs? := remainder?.map $ λ g => (g, rhs)
-    return .success {
-      root := state.root,
-      savedState := {
-        term := ← MonadBacktrack.saveState,
-        tactic := { goals },
-      },
-      parentMVar? := .some goal,
-      calcPrevRhs?
-    } #[]
-  catch exception =>
-    return .failure #[← exception.toMessageData.toString]
 
 end Pantograph

Pantograph/Serial.lean

@@ -104,8 +104,7 @@ def goalStatePickle (goalState : GoalState) (path : System.FilePath) : IO Unit :
     }
     root,
     parentMVar?,
-    convMVar?,
+    fragment?,
-    calcPrevRhs?,
   } := goalState
   Pantograph.pickle path (
     env.constants.map₂,
@@ -117,8 +116,7 @@ def goalStatePickle (goalState : GoalState) (path : System.FilePath) : IO Unit :
 
     root,
     parentMVar?,
-    convMVar?,
+    fragment?,
-    calcPrevRhs?,
   )
 
 @[export pantograph_goal_state_unpickle_m]
@@ -134,8 +132,7 @@ def goalStateUnpickle (path : System.FilePath) (env : Environment)
 
     root,
     parentMVar?,
-    convMVar?,
+    fragment?,
-    calcPrevRhs?,
   ), region) ← Pantograph.unpickle (
     PHashMap Name ConstantInfo ×
 
@@ -146,8 +143,7 @@ def goalStateUnpickle (path : System.FilePath) (env : Environment)
 
     MVarId ×
     Option MVarId ×
-    Option (MVarId × MVarId × List MVarId) ×
+    Option TacticFragment
-    Option (MVarId × Expr)
   ) path
   let env ← env.replay (Std.HashMap.ofList map₂.toList)
   let goalState := {
@@ -167,8 +163,7 @@ def goalStateUnpickle (path : System.FilePath) (env : Environment)
     },
     root,
     parentMVar?,
-    convMVar?,
+    fragment?,
-    calcPrevRhs?,
   }
   return (goalState, region)
 

Pantograph/Tactic.lean

@@ -1,2 +1,3 @@
 import Pantograph.Tactic.Assign
+import Pantograph.Tactic.Fragment
 import Pantograph.Tactic.Prograde

Pantograph/Tactic/Fragment.lean

@@ -0,0 +1,118 @@
+/- Fragmented tactics are the tactics which can give incremental feedback and
+whose integrity as a block is crucial to its operation. e.g. `calc` or `conv`.
+Here, a unified system handles all fragments.
+
+Inside a tactic fragment, the parser category may be different. An incomplete
+fragmented tactic may not be elaboratable..
+-/
+import Lean.Meta
+import Lean.Elab
+
+open Lean
+
+namespace Pantograph
+
+inductive Fragment where
+  | calc (prevRhs? : Option Expr)
+  | conv (rhs : MVarId) (dormant : List MVarId)
+
+protected def Fragment.enterCalc : Elab.Tactic.TacticM Fragment := do
+  return .calc .none
+protected def Fragment.enterConv : Elab.Tactic.TacticM Fragment := do
+  let goal ← Elab.Tactic.getMainGoal
+  let convGoal ← goal.withContext do
+    let (rhs, newGoal) ← Elab.Tactic.Conv.mkConvGoalFor (← Elab.Tactic.getMainTarget)
+    Elab.Tactic.replaceMainGoal [newGoal.mvarId!]
+    pure rhs.mvarId!
+  let otherGoals := (← Elab.Tactic.getGoals).filter (· != goal)
+  return .conv convGoal otherGoals
+
+protected def Fragment.exit (fragment : Fragment) (goal: MVarId) : Elab.Tactic.TacticM Unit :=
+  match fragment with
+  | .calc _prevRhs? => Elab.Tactic.setGoals [goal]
+  | .conv rhs otherGoals => do
+    -- Close all existing goals with `refl`
+    for mvarId in (← Elab.Tactic.getGoals) do
+      liftM <| mvarId.refl <|> mvarId.inferInstance <|> pure ()
+    Elab.Tactic.pruneSolvedGoals
+    unless (← Elab.Tactic.getGoals).isEmpty do
+      throwError "convert tactic failed, there are unsolved goals\n{Elab.goalsToMessageData (← Elab.Tactic.getGoals)}"
+
+    Elab.Tactic.setGoals $ [goal] ++ otherGoals
+    let targetNew ← instantiateMVars (.mvar rhs)
+    let proof ← instantiateMVars (.mvar goal)
+
+    Elab.Tactic.liftMetaTactic1 (·.replaceTargetEq targetNew proof)
+
+structure TacticFragment where
+  -- The goal which the fragment acts on
+  goal : MVarId
+  fragment : Fragment
+
+protected def Fragment.step (fragment : Fragment) (goal : MVarId) (s : String)
+  : Elab.Tactic.TacticM (Option TacticFragment) := goal.withContext do
+  match fragment with
+  | .calc prevRhs? => do
+    let .ok stx := Parser.runParserCategory
+      (env := ← getEnv)
+      (catName := `term)
+      (input := s)
+      (fileName := ← getFileName) | throwError s!"Failed to parse calc element {s}"
+    let `(term|$pred) := stx
+    let decl ← goal.getDecl
+    let target ← instantiateMVars decl.type
+    let tag := decl.userName
+
+    let mut step ← Elab.Term.elabType <| ← do
+      if let some prevRhs := prevRhs? then
+        Elab.Term.annotateFirstHoleWithType pred (← Meta.inferType prevRhs)
+      else
+        pure pred
+
+    let some (_, lhs, rhs) ← Elab.Term.getCalcRelation? step |
+      throwErrorAt pred "invalid 'calc' step, relation expected{indentExpr step}"
+    if let some prevRhs := prevRhs? then
+      unless ← Meta.isDefEqGuarded lhs prevRhs do
+        throwErrorAt pred "invalid 'calc' step, left-hand-side is{indentD m!"{lhs} : {← Meta.inferType lhs}"}\nprevious right-hand-side is{indentD m!"{prevRhs} : {← Meta.inferType prevRhs}"}"
+
+    -- Creates a mvar to represent the proof that the calc tactic solves the
+    -- current branch
+    -- In the Lean `calc` tactic this is gobbled up by
+    -- `withCollectingNewGoalsFrom`
+    let mut proof ← Meta.mkFreshExprMVarAt (← getLCtx) (← Meta.getLocalInstances) step
+      (userName := tag ++ `calc)
+    let mvarBranch := proof.mvarId!
+
+    let mut proofType ← Meta.inferType proof
+    let mut remainder? := Option.none
+
+    -- The calc tactic either solves the main goal or leaves another relation.
+    -- Replace the main goal, and save the new goal if necessary
+    unless ← Meta.isDefEq proofType target do
+      let rec throwFailed :=
+        throwError "'calc' tactic failed, has type{indentExpr proofType}\nbut it is expected to have type{indentExpr target}"
+      let some (_, _, rhs) ← Elab.Term.getCalcRelation? proofType | throwFailed
+      let some (r, _, rhs') ← Elab.Term.getCalcRelation? target | throwFailed
+      let lastStep := mkApp2 r rhs rhs'
+      let lastStepGoal ← Meta.mkFreshExprSyntheticOpaqueMVar lastStep tag
+      (proof, proofType) ← Elab.Term.mkCalcTrans proof proofType lastStepGoal lastStep
+      unless ← Meta.isDefEq proofType target do throwFailed
+      remainder? := .some lastStepGoal.mvarId!
+    goal.assign proof
+
+    let goals := [ mvarBranch ] ++ remainder?.toList
+    Elab.Tactic.setGoals goals
+    match remainder? with
+    | .some goal => return .some { goal, fragment := .calc (.some rhs) }
+    | .none => return .none
+  | fragment@(.conv _ _) => do
+    let .ok tactic := Parser.runParserCategory
+      (env := ← MonadEnv.getEnv)
+      (catName := `conv)
+      (input := s)
+      (fileName := ← getFileName) | throwError "Could not parse `conv tactic {s}"
+    Elab.Tactic.evalTactic tactic
+    let goal ← Elab.Tactic.getMainGoal
+    return .some { goal, fragment }
+
+end Pantograph

Repl.lean

@@ -336,12 +336,7 @@ def execute (command: Protocol.Command): MainM Json := do
         | true, .none => do
           pure $ nextGoalState.immediateResume goalState
         | true, .some true => pure nextGoalState
-        | true, .some false => do
+        | true, .some false => pure nextGoalState
-          let .some (_, _, dormantGoals) := goalState.convMVar? |
-            Protocol.throw $ errorIO "If conv exit succeeded this should not fail"
-          let .ok result := nextGoalState.resume (nextGoalState.goals ++ dormantGoals) |
-            Protocol.throw $ errorIO "Resuming known goals"
-          pure result
         | false, _ => pure nextGoalState
       let nextStateId ← newGoalState nextGoalState
       let parentExpr := nextGoalState.parentExpr?.get!

Test/Metavar.lean

@@ -259,11 +259,11 @@ def test_partial_continuation: TestM Unit := do
 
   -- Continuation should fail if the state does not exist:
   match state0.resume coupled_goals with
-  | .error error => addTest $ LSpec.check "(continuation failure message)" (error = "Goals [_uniq.44, _uniq.45, _uniq.42, _uniq.51] are not in scope")
+  | .error error => checkEq "(continuation failure message)" error "Goals [_uniq.44, _uniq.45, _uniq.42, _uniq.51] are not in scope"
   | .ok _ => fail "(continuation should fail)"
   -- Continuation should fail if some goals have not been solved
   match state2.continue state1 with
-  | .error error => addTest $ LSpec.check "(continuation failure message)" (error = "Target state has unresolved goals")
+  | .error error => checkEq "(continuation failure message)" error "Target state has unresolved goals"
   | .ok _ => fail "(continuation should fail)"
   return ()
 
@@ -274,13 +274,12 @@ def test_branch_unification : TestM Unit := do
   let .success state _  ← state.tacticOn' 0 (← `(tactic|apply And.intro)) | fail "apply And.intro failed to run"
   let .success state1 _  ← state.tacticOn' 0 (← `(tactic|exact h)) | fail "exact h failed to run"
   let .success state2 _  ← state.tacticOn' 1 (← `(tactic|apply Or.inl)) | fail "apply Or.inl failed to run"
-  assert! state2.goals.length == 1
+  checkEq "(state2 goals)" state2.goals.length 1
   let state' ← state2.replay state state1
-  assert! state'.goals.length == 1
+  checkEq "(state' goals)" state'.goals.length 1
   let .success stateT _ ← state'.tacticOn' 0 (← `(tactic|exact h)) | fail "exact h failed to run"
   let .some root := stateT.rootExpr? | fail "Root expression must exist"
   checkEq "(root)" (toString $ ← Meta.ppExpr root) "fun p q h => ⟨h, Or.inl h⟩"
-  return ()
 
 def suite (env: Environment): List (String × IO LSpec.TestSeq) :=
   let tests := [