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

2025-06-23 20:57:53 -07:00 · 2025-06-23 20:57:53 -07:00 · 58ae791da3
parent 05a8e3b13c
commit 58ae791da3
6 changed files with 204 additions and 179 deletions
--- a/Pantograph/Goal.lean
+++ b/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.
-  -- (convRhs, goal, dormant)
-  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
+  -- If this is not `.none` there is a partial tactic being executed
+  fragment? : Option TacticFragment := .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
-  let goal ← state.savedState.tactic.goals[goalId]?
+protected def GoalState.focus (state : GoalState) (goal : MVarId): Option GoalState := do
+  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)
-  : Elab.TermElabM GoalState := do
+protected def GoalState.step' { α } (state: GoalState) (goal: MVarId) (tacticM: Elab.Tactic.TacticM α) (guardMVarErrors : Bool := false)
+  : 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 -/
-protected def GoalState.tryTacticM
-    (state: GoalState) (goal: MVarId) (tacticM: Elab.Tactic.TacticM Unit)
-    (guardMVarErrors : Bool := false)
-    : Elab.TermElabM TacticResult := do
-  assert! ¬ (← goal.isAssigned)
-  let prevMessageLength := state.coreState.messages.toList.length
+def withCapturingError (elabM : Elab.Term.TermElabM GoalState) : Elab.TermElabM TacticResult := do
+  -- FIXME: Maybe message log should be empty
+  let prevMessageLength := (← Core.getMessageLog).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
-    state.restoreTacticM goal
-
-    -- See Lean.Elab.Tactic.Conv.convTarget
-    let convMVar ← Elab.Tactic.withMainContext do
-      let (rhs, newGoal) ← Elab.Tactic.Conv.mkConvGoalFor (← Elab.Tactic.getMainTarget)
-      Elab.Tactic.replaceMainGoal [newGoal.mvarId!]
-      pure rhs.mvarId!
-    return (← MonadBacktrack.saveState, convMVar)
-  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]
+  withCapturingError do
+    let (fragment, state') ← state.step' goal Fragment.enterConv
+    let fragment? := .some {
+      goal := state'.goals[0]!,
+      fragment,
+    }
+    return {
+      state' with
+      fragment?
+    }

 /-- 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
-    | .some mvar => pure mvar
-    | .none => return .invalidAction "Not in conv state"
-  let tacticM :  Elab.Tactic.TacticM Elab.Tactic.SavedState:= do
-    -- Vide `Lean.Elab.Tactic.Conv.convert`
-    state.savedState.restore
-
-    -- 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 [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]
+  let .some { goal, fragment } := state.fragment? |
+    return .invalidAction "Not in conv state"
+  unless fragment matches .conv .. do
+    return .invalidAction "Not in conv state"
+  withCapturingError do
+    let state' ← state.step goal (fragment.exit goal)
+    return {
+      state' with
+      fragment? := .none,
+    }

 protected def GoalState.calcPrevRhsOf? (state: GoalState) (goal: MVarId): Option Expr := do
-  let (mvarId, rhs) ← state.calcPrevRhs?
-  if mvarId == goal then
-    .some rhs
-  else
-    .none
+  match state.fragment? with
+  | .some { goal := goal', fragment := .calc prevRhs? } =>
+    if goal == goal' then prevRhs? else .none
+  | .some _ => unreachable!
+  | .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
-  if state.convMVar?.isSome then
-    return .invalidAction "Cannot initiate `calc` while in `conv` state"
-  let `(term|$pred) ← match Parser.runParserCategory
-    (env := state.env)
-    (catName := `term)
-    (input := pred)
-    (fileName := ← getFileName) with
-    | .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]
+  let prevRhs? := state.calcPrevRhsOf? goal
+  withCapturingError do
+    let (fragment?, state') ← state.step' goal do
+      let fragment := Fragment.calc prevRhs?
+      fragment.step goal pred
+    return {
+      state' with
+      fragment?,
+    }

 end Pantograph
--- a/Pantograph/Serial.lean
+++ b/Pantograph/Serial.lean
@ -104,8 +104,7 @@ def goalStatePickle (goalState : GoalState) (path : System.FilePath) : IO Unit :
    }
    root,
    parentMVar?,
-    convMVar?,
-    calcPrevRhs?,
+    fragment?,
  } := goalState
  Pantograph.pickle path (
    env.constants.map₂,
@ -117,8 +116,7 @@ def goalStatePickle (goalState : GoalState) (path : System.FilePath) : IO Unit :

    root,
    parentMVar?,
-    convMVar?,
-    calcPrevRhs?,
+    fragment?,
  )

@[export pantograph_goal_state_unpickle_m]
@ -134,8 +132,7 @@ def goalStateUnpickle (path : System.FilePath) (env : Environment)

    root,
    parentMVar?,
-    convMVar?,
-    calcPrevRhs?,
+    fragment?,
  ), 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 (MVarId × Expr)
+    Option TacticFragment
  ) 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?,
-    calcPrevRhs?,
+    fragment?,
  }
  return (goalState, region)

--- a/Pantograph/Tactic.lean
+++ b/Pantograph/Tactic.lean
@ -1,2 +1,3 @@
 import Pantograph.Tactic.Assign
+import Pantograph.Tactic.Fragment
 import Pantograph.Tactic.Prograde
--- a/Pantograph/Tactic/Fragment.lean
+++ b/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
--- a/Repl.lean
+++ b/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
-          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
+        | true, .some false => pure nextGoalState
        | false, _ => pure nextGoalState
      let nextStateId ← newGoalState nextGoalState
      let parentExpr := nextGoalState.parentExpr?.get!
--- a/Test/Metavar.lean
+++ b/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 := [