Merge pull request 'Simplify goal bookkeeping mechanism' (#10) from tactic/book into dev

Reviewed-on: #10

Pantograph.lean

@@ -2,6 +2,7 @@ import Pantograph.Commands
 import Pantograph.Serial
 import Pantograph.Symbols
 import Pantograph.Tactic
+import Pantograph.SemihashMap
 
 namespace Pantograph
 
@@ -11,8 +12,7 @@ structure Context where
 /-- Stores state of the REPL -/
 structure State where
   options: Commands.Options := {}
-  --environments: Array Lean.Environment := #[]
-  proofTrees:   Array ProofTree := #[]
+  goalStates: SemihashMap GoalState := SemihashMap.empty
 
 -- State monad
 abbrev MainM := ReaderT Context (StateT State Lean.Elab.TermElabM)
@@ -30,14 +30,15 @@ def execute (command: Commands.Command): MainM Lean.Json := do
     | .error error => return Lean.toJson $ errorCommand s!"Unable to parse json: {error}"
   match command.cmd with
   | "reset"          => run reset
+  | "stat"           => run stat
   | "expr.echo"      => run expr_echo
   | "lib.catalog"    => run lib_catalog
   | "lib.inspect"    => run lib_inspect
   | "options.set"    => run options_set
   | "options.print"  => run options_print
-  | "proof.start"     => run proof_start
-  | "proof.tactic"    => run proof_tactic
-  | "proof.printTree" => run proof_print_tree
+  | "goal.start"     => run goal_start
+  | "goal.tactic"    => run goal_tactic
+  | "goal.delete"    => run goal_delete
   | cmd =>
     let error: Commands.InteractionError :=
       errorCommand s!"Unknown command {cmd}"
@@ -47,11 +48,15 @@ def execute (command: Commands.Command): MainM Lean.Json := do
   errorCommand := errorI "command"
   errorIndex := errorI "index"
   -- Command Functions
-  reset (_: Commands.Reset): MainM (CR Commands.ResetResult) := do
+  reset (_: Commands.Reset): MainM (CR Commands.StatResult) := do
     let state ← get
-    let nTrees := state.proofTrees.size
-    set { state with proofTrees := #[] }
-    return .ok { nTrees := nTrees }
+    let nGoals := state.goalStates.size
+    set { state with goalStates := SemihashMap.empty }
+    return .ok { nGoals }
+  stat (_: Commands.Stat): MainM (CR Commands.StatResult) := do
+    let state ← get
+    let nGoals := state.goalStates.size
+    return .ok { nGoals }
   lib_catalog (_: Commands.LibCatalog): MainM (CR Commands.LibCatalogResult) := do
     let env ← Lean.MonadEnv.getEnv
     let names := env.constants.fold (init := #[]) (λ acc name info =>
@@ -113,7 +118,7 @@ def execute (command: Commands.Command): MainM Lean.Json := do
     return .ok {  }
   options_print (_: Commands.OptionsPrint): MainM (CR Commands.OptionsPrintResult) := do
     return .ok (← get).options
-  proof_start (args: Commands.ProofStart): MainM (CR Commands.ProofStartResult) := do
+  goal_start (args: Commands.GoalStart): MainM (CR Commands.GoalStartResult) := do
     let state ← get
     let env ← Lean.MonadEnv.getEnv
     let expr?: Except _ Lean.Expr ← (match args.expr, args.copyFrom with
@@ -134,32 +139,36 @@ def execute (command: Commands.Command): MainM Lean.Json := do
     match expr? with
     | .error error => return .error error
     | .ok expr =>
-      let tree ← ProofTree.create expr
-      -- Put the new tree in the environment
-      let nextTreeId := state.proofTrees.size
-      set { state with proofTrees := state.proofTrees.push tree }
-      return .ok { treeId := nextTreeId }
-  proof_tactic (args: Commands.ProofTactic): MainM (CR Commands.ProofTacticResult) := do
+      let goalState ← GoalState.create expr
+      let (goalStates, goalId) := state.goalStates.insert goalState
+      set { state with goalStates }
+      return .ok { goalId }
+  goal_tactic (args: Commands.GoalTactic): MainM (CR Commands.GoalTacticResult) := do
     let state ← get
-    match state.proofTrees.get? args.treeId with
-    | .none => return .error $ errorIndex "Invalid tree index {args.treeId}"
-    | .some tree =>
-      let (result, nextTree) ← ProofTree.execute
-        (stateId := args.stateId)
-        (goalId := args.goalId.getD 0)
-        (tactic := args.tactic) |>.run state.options |>.run tree
+    match state.goalStates.get? args.goalId with
+    | .none => return .error $ errorIndex s!"Invalid goal index {args.goalId}"
+    | .some goalState =>
+      let result ← GoalState.execute goalState args.tactic |>.run state.options
       match result with
-      | .invalid message => return .error $  errorIndex message
-      | .success nextId? goals =>
-        set { state with proofTrees := state.proofTrees.set! args.treeId nextTree }
-        return .ok { nextId? := nextId?, goals? := .some goals }
+      | .success goals =>
+        if goals.isEmpty then
+          return .ok {}
+        else
+          -- Append all goals
+          let (goalStates, goalIds, sGoals) := Array.foldl (λ acc itr =>
+            let (map, indices, serializedGoals) := acc
+            let (goalState, sGoal) := itr
+            let (map, index) := map.insert goalState
+            (map, index :: indices, sGoal :: serializedGoals)
+            ) (state.goalStates, [], []) goals
+          set { state with goalStates }
+          return .ok { goals? := .some sGoals.reverse.toArray, goalIds? := .some goalIds.reverse.toArray }
       | .failure messages =>
         return .ok { tacticErrors? := .some messages }
-  proof_print_tree (args: Commands.ProofPrintTree): MainM (CR Commands.ProofPrintTreeResult) := do
+  goal_delete (args: Commands.GoalDelete): MainM (CR Commands.GoalDeleteResult) := do
     let state ← get
-    match state.proofTrees.get? args.treeId with
-    | .none => return .error $ errorIndex "Invalid tree index {args.treeId}"
-    | .some tree =>
-      return .ok { parents := tree.structure_array }
+    let goalStates := args.goalIds.foldl (λ map id => map.remove id) state.goalStates
+    set { state with goalStates }
+    return .ok {}
 
 end Pantograph

Pantograph/Commands.lean

@@ -80,8 +80,11 @@ structure InteractionError where
 
 structure Reset where
   deriving Lean.FromJson
-structure ResetResult where
-  nTrees: Nat
+structure Stat where
+  deriving Lean.FromJson
+structure StatResult where
+  -- Number of goals states
+  nGoals: Nat
   deriving Lean.ToJson
 
 -- Return the type of an expression
@@ -127,35 +130,34 @@ structure OptionsPrint where
   deriving Lean.FromJson
 abbrev OptionsPrintResult := Options
 
-structure ProofStart where
+structure GoalStart where
   -- Only one of the fields below may be populated.
   expr: Option String     -- Proof expression
   copyFrom: Option String -- Theorem name
   deriving Lean.FromJson
-structure ProofStartResult where
-  treeId: Nat := 0 -- Proof tree id
+structure GoalStartResult where
+  goalId: Nat := 0 -- Proof tree id
   deriving Lean.ToJson
-structure ProofTactic where
+structure GoalTactic where
   -- Identifiers for tree, state, and goal
-  treeId: Nat
-  stateId: Nat
-  goalId: Option Nat -- Defaults to 0
+  goalId: Nat
   tactic: String
   deriving Lean.FromJson
-structure ProofTacticResult where
+structure GoalTacticResult where
   -- Existence of this field shows success
   goals?: Option (Array Goal)          := .none
   -- Next proof state id, if successful
-  nextId?: Option Nat                  := .none
+  goalIds?: Option (Array Nat)          := .none
   -- Existence of this field shows failure
   tacticErrors?: Option (Array String) := .none
   deriving Lean.ToJson
-structure ProofPrintTree where
-  treeId: Nat
+
+-- Remove a bunch of goals.
+structure GoalDelete where
+  goalIds: List Nat
   deriving Lean.FromJson
-structure ProofPrintTreeResult where
-  -- "" if no parents, otherwise "parentId.goalId"
-  parents: Array String
+structure GoalDeleteResult where
   deriving Lean.ToJson
 
+
 end Pantograph.Commands

Pantograph/SemihashMap.lean

@@ -0,0 +1,149 @@
+
+namespace Pantograph.SemihashMap
+
+structure Imp (β: Type u) where
+  data: Array (Option β)
+
+  -- Number of elements currently in use
+  size: Nat
+
+  -- Next index that has never been touched
+  allocFront: Nat
+
+  -- Deallocated indices
+  deallocs: Array Nat
+
+  -- Number of valid entries in `deallocs` array
+  lastDealloc: Nat
+
+namespace Imp
+
+structure WF (m: Imp β): Prop where
+  capacity: m.data.size = m.deallocs.size
+  front_dealloc: ∀ i: Fin m.deallocs.size, (i < m.allocFront) → (m.deallocs.get i) < m.allocFront
+  front_data: ∀ i: Fin m.data.size, (i ≥ m.allocFront) → (m.data.get i).isNone
+
+def empty (capacity := 16): Imp β :=
+  {
+    data := mkArray capacity .none,
+    size := 0,
+    allocFront := 0,
+    deallocs := mkArray capacity 0,
+    lastDealloc := 0,
+  }
+
+private theorem list_get_replicate (x: α) (i: Fin (List.replicate n x).length):
+    List.get (List.replicate n x) i = x := by
+  sorry
+
+theorem empty_wf : WF (empty n: Imp β) := by
+  unfold empty
+  apply WF.mk
+  case capacity =>
+    conv =>
+      lhs
+      congr
+      simp
+    conv =>
+      rhs
+      congr
+      simp
+    simp
+  case front_dealloc =>
+    simp_all
+    intro i
+    intro a
+    contradiction
+  case front_data =>
+    simp_all
+    intro i
+    unfold Imp.data at i
+    simp at i
+    conv =>
+      lhs
+      unfold Array.get
+      unfold mkArray
+      simp [List.replicate]
+      rewrite [list_get_replicate]
+
+-- FIXME: Merge this with the well-formed versions below so proof and code can
+-- mesh seamlessly.
+@[inline] def insert (map: Imp β) (v: β): (Imp β × Nat) :=
+  match map.lastDealloc with
+  | 0 => -- Capacity is full, buffer expansion is required
+    if map.size == map.data.size then
+      let nextIndex := map.data.size
+      let extendCapacity := map.size
+      let result: Imp β := {
+        data := (map.data.append #[Option.some v]).append (mkArray extendCapacity .none),
+        size := map.size + 1,
+        allocFront := map.size + 1,
+        deallocs := mkArray (map.data.size + 1 + extendCapacity) 0,
+        lastDealloc := 0,
+      }
+      (result, nextIndex)
+    else
+      let nextIndex := map.size
+      let result: Imp β := {
+        map
+        with
+        data := map.data.set ⟨nextIndex, sorry⟩ (Option.some v),
+        size := map.size + 1,
+        allocFront := map.allocFront + 1,
+      }
+      (result, nextIndex)
+  | (.succ k) => -- Allocation list has space
+    let nextIndex := map.deallocs.get! k
+    let result: Imp β := {
+      map with
+      data := map.data.set ⟨nextIndex, sorry⟩ (Option.some v),
+      size := map.size + 1,
+      lastDealloc := map.lastDealloc - 1
+    }
+    (result, nextIndex)
+
+@[inline] def remove (map: Imp β) (index: Fin (map.size)): Imp β :=
+  have h: index.val < map.data.size := by sorry
+  match map.data.get ⟨index.val, h⟩ with
+  | .none => map
+  | .some _ =>
+    {
+      map with
+      data := map.data.set ⟨index, sorry⟩ .none,
+      size := map.size - 1,
+      deallocs := map.deallocs.set ⟨map.lastDealloc, sorry⟩ index,
+      lastDealloc := map.lastDealloc + 1,
+    }
+
+/-- Retrieval is efficient -/
+@[inline] def get? (map: Imp β) (index: Fin (map.size)): Option β :=
+  have h: index.val < map.data.size := by sorry
+  map.data.get ⟨index.val, h⟩
+@[inline] def capacity (map: Imp β): Nat := map.data.size
+
+end Imp
+
+
+/--
+This is like a hashmap but you cannot control the keys.
+-/
+def _root_.Pantograph.SemihashMap β := {m: Imp β // m.WF}
+
+@[inline] def empty (capacity := 16): SemihashMap β :=
+  ⟨ Imp.empty capacity, Imp.empty_wf ⟩
+@[inline] def insert (map: SemihashMap β) (v: β): (SemihashMap β × Nat) :=
+  let ⟨imp, pre⟩ := map
+  let ⟨result, id⟩ := imp.insert v
+  ( ⟨ result, sorry ⟩, id)
+@[inline] def remove (map: SemihashMap β) (index: Nat): SemihashMap β :=
+  let ⟨imp, pre⟩ := map
+  let result := imp.remove ⟨index, sorry⟩
+  ⟨ result, sorry ⟩
+@[inline] def get? (map: SemihashMap β) (index: Nat): Option β :=
+  let ⟨imp, _⟩ := map
+  imp.get? ⟨index, sorry⟩
+@[inline] def size (map: SemihashMap β): Nat :=
+  let ⟨imp, _⟩ := map
+  imp.size
+
+end Pantograph.SemihashMap

Pantograph/Tactic.lean

@@ -24,35 +24,22 @@ def Lean.MessageLog.getErrorMessages (log : MessageLog) : MessageLog :=
 namespace Pantograph
 open Lean
 
-structure ProofState where
-  goals : List MVarId
+structure GoalState where
+  mvarId: MVarId
   savedState : Elab.Tactic.SavedState
-  parent : Option Nat := none
-  parentGoalId : Nat  := 0
-structure ProofTree where
-  -- Set of proof states
-  states : Array ProofState := #[]
 
 abbrev M := Elab.TermElabM
 
-def ProofTree.create (expr: Expr): M ProofTree := do
+def GoalState.create (expr: Expr): M GoalState := do
   let expr ← instantiateMVars expr
   let goal := (← Meta.mkFreshExprMVar expr (kind := MetavarKind.synthetic))
   let savedStateMonad: Elab.Tactic.TacticM Elab.Tactic.SavedState := MonadBacktrack.saveState
   let savedState ← savedStateMonad { elaborator := .anonymous } |>.run' { goals := [goal.mvarId!]}
   return {
-    states := #[{
     savedState := savedState,
-      goals := [goal.mvarId!]
-    }]
+    mvarId := goal.mvarId!
   }
 
--- Print the tree structures in readable form
-def ProofTree.structure_array (tree: ProofTree): Array String :=
-  tree.states.map λ state => match state.parent with
-    | .none => ""
-    | .some parent => s!"{parent}.{state.parentGoalId}"
-
 def execute_tactic (state: Elab.Tactic.SavedState) (goal: MVarId) (tactic: String) :
     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
@@ -78,44 +65,38 @@ def execute_tactic (state: Elab.Tactic.SavedState) (goal: MVarId) (tactic: Strin
 
 /-- Response for executing a tactic -/
 inductive TacticResult where
-  -- Invalid id
-  | invalid (message: String): TacticResult
   -- Goes to next state
-  | success (nextId?: Option Nat) (goals: Array Commands.Goal)
+  | success (goals: Array (GoalState × Commands.Goal))
   -- Fails with messages
   | failure (messages: Array String)
 
+namespace TacticResult
+
+def is_success: TacticResult → Bool
+  | .success _ => true
+  | .failure _ => false
+
+end TacticResult
+
 /-- Execute tactic on given state -/
-def ProofTree.execute (stateId: Nat) (goalId: Nat) (tactic: String):
-    Commands.OptionsT StateRefT ProofTree M TacticResult := do
+def GoalState.execute (goal: GoalState) (tactic: String):
+    Commands.OptionsT M TacticResult := do
   let options ← read
-  let tree ← get
-  match tree.states.get? stateId with
-  | .none => return .invalid s!"Invalid state id {stateId}"
-  | .some state =>
-    match state.goals.get? goalId with
-    | .none => return .invalid s!"Invalid goal id {goalId}"
-    | .some goal =>
-      match (← execute_tactic (state := state.savedState) (goal := goal) (tactic := tactic)) with
+  match (← execute_tactic (state := goal.savedState) (goal := goal.mvarId) (tactic := tactic)) with
   | .error errors =>
     return .failure errors
   | .ok (nextState, nextGoals) =>
-        let nextId := tree.states.size
     if nextGoals.isEmpty then
-          return .success .none #[]
+      return .success #[]
     else
-          let proofState: ProofState := {
-            savedState := nextState,
-            goals := nextGoals,
-            parent := stateId,
-            parentGoalId := goalId
-          }
-          modify fun s => { s with states := s.states.push proofState }
-        let parentDecl? := (← MonadMCtx.getMCtx).findDecl? goal
-        let goals ← nextGoals.mapM fun mvarId => do
-          match (← MonadMCtx.getMCtx).findDecl? mvarId with
-          | .some mvarDecl => serialize_goal options mvarDecl (parentDecl? := parentDecl?)
-          | .none => throwError mvarId
-        return .success (.some nextId) goals.toArray
+      let nextGoals: List GoalState := nextGoals.map fun mvarId => { mvarId, savedState := nextState }
+      let parentDecl? := (← MonadMCtx.getMCtx).findDecl? goal.mvarId
+      let goals ← nextGoals.mapM fun nextGoal => do
+        match (← MonadMCtx.getMCtx).findDecl? nextGoal.mvarId with
+        | .some mvarDecl =>
+          let serializedGoal ← serialize_goal options mvarDecl (parentDecl? := parentDecl?)
+          return (nextGoal, serializedGoal)
+        | .none => throwError nextGoal.mvarId
+      return .success goals.toArray
 
 end Pantograph

Pantograph/Version.lean

@@ -1,5 +1,5 @@
 namespace Pantograph
 
-def version := "0.2.3"
+def version := "0.2.4"
 
 end Pantograph

README.md

@@ -50,15 +50,16 @@ Example with `mathlib4` (~90k symbols, may stack overflow, see troubleshooting)
 $ lake env build/bin/Pantograph Mathlib.Analysis.Seminorm
 lib.catalog
 ```
-Example proving a theorem: (alternatively use `proof.start {"copyFrom": "Nat.add_comm"}`) to prime the proof
+Example proving a theorem: (alternatively use `goal.start {"copyFrom": "Nat.add_comm"}`) to prime the proof
 ```
 $ env build/bin/Pantograph Init
-proof.start {"expr": "∀ (n m : Nat), n + m = m + n"}
-proof.tactic {"treeId": 0, "stateId": 0, "goalId": 0, "tactic": "intro n m"}
-proof.tactic {"treeId": 0, "stateId": 1, "goalId": 0, "tactic": "assumption"}
-proof.printTree {"treeId": 0}
-proof.tactic {"treeId": 0, "stateId": 1, "goalId": 0, "tactic": "rw [Nat.add_comm]"}
-proof.printTree {"treeId": 0}
+goal.start {"expr": "∀ (n m : Nat), n + m = m + n"}
+goal.tactic {"goalId": 0, "tactic": "intro n m"}
+goal.tactic {"goalId": 1, "tactic": "assumption"}
+goal.delete {"goalIds": [0]}
+stat {}
+goal.tactic {"goalId": 1, "tactic": "rw [Nat.add_comm]"}
+stat
 ```
 where the application of `assumption` should lead to a failure.
 
@@ -74,9 +75,10 @@ See `Pantograph/Commands.lean` for a description of the parameters and return va
 - `options.set { key: value, ... }`: Set one or more options (not Lean options; those
   have to be set via command line arguments.), for options, see `Pantograph/Commands.lean`
 - `options.print`: Display the current set of options
-- `proof.start {["name": <name>], ["expr": <expr>], ["copyFrom": <symbol>]}`: Start a new proof state from a given expression or symbol
-- `proof.tactic {"treeId": <id>, "stateId": <id>, "goalId": <id>, "tactic": string}`: Execute a tactic on a given proof state
-- `proof.printTree {"treeId": <id>}`: Print the topological structure of a proof tree
+- `goal.start {["name": <name>], ["expr": <expr>], ["copyFrom": <symbol>]}`: Start a new goal from a given expression or symbol
+- `goal.tactic {"goalId": <id>, "tactic": <tactic>}`: Execute a tactic string on a given goal
+- `goal.remove {"goalIds": [<id>]}"`: Remove a bunch of stored goals.
+- `stat`: Display resource usage
 
 ## Errors
 

Test/Proofs.lean

@@ -10,87 +10,46 @@ inductive Start where
   | copy (name: String) -- Start from some name in the environment
   | expr (expr: String) -- Start from some expression
 
-abbrev TestM := ReaderT Commands.Options StateRefT ProofTree M
-
-def start_proof (start: Start): M (LSpec.TestSeq × Option ProofTree) := do
-  let env ← Lean.MonadEnv.getEnv
-  let mut testSeq := LSpec.TestSeq.done
-  match start with
-  | .copy name =>
-    let cInfo? := str_to_name name |> env.find?
-    testSeq := testSeq ++ LSpec.check s!"Symbol exists {name}" cInfo?.isSome
-    match cInfo? with
-    | .some cInfo =>
-      let state ← ProofTree.create
-        (expr := cInfo.type)
-      return (testSeq, Option.some state)
-    | .none =>
-      return (testSeq, Option.none)
-  | .expr expr =>
-    let syn? := syntax_from_str env expr
-    testSeq := testSeq ++ LSpec.check s!"Parsing {expr}" (syn?.isOk)
-    match syn? with
-    | .error error =>
-      IO.println error
-      return (testSeq, Option.none)
-    | .ok syn =>
-      let expr? ← syntax_to_expr syn
-      testSeq := testSeq ++ LSpec.check s!"Elaborating" expr?.isOk
-      match expr? with
-      | .error error =>
-        IO.println error
-        return (testSeq, Option.none)
-      | .ok expr =>
-        let state ← ProofTree.create
-          (expr := expr)
-        return (testSeq, Option.some state)
+abbrev TestM := StateRefT LSpec.TestSeq (ReaderT Commands.Options M)
 
 deriving instance DecidableEq, Repr for Commands.Expression
 deriving instance DecidableEq, Repr for Commands.Variable
 deriving instance DecidableEq, Repr for Commands.Goal
-deriving instance DecidableEq, Repr for TacticResult
 
-/-- Check the output of each proof step -/
-def proof_step (stateId: Nat) (goalId: Nat) (tactic: String)
-    (expected: TacticResult) : TestM LSpec.TestSeq := do
-  let options ← read
-  let result: TacticResult ← ProofTree.execute stateId goalId tactic |>.run options
-  match expected, result with
-  | .success (.some i) #[], .success (.some _) goals =>
-    -- If the goals are omitted but the next state is specified, we imply that
-    -- the tactic succeeded.
-    let expected := .success (.some i) goals
-    return LSpec.test s!"{stateId}.{goalId} {tactic}"   (result = expected)
-  | _, _ =>
-    return LSpec.test s!"{stateId}.{goalId} {tactic}"   (result = expected)
+def add_test (test: LSpec.TestSeq): TestM Unit := do
+  set $ (← get) ++ test
 
-/-- Check that the tree structure is correct -/
-def proof_inspect (expected: Array String) : TestM LSpec.TestSeq := do
-  let result := (← get).structure_array
-  return LSpec.test s!"tree structure" (result = expected)
+def start_proof (start: Start): TestM (Option GoalState) := do
+  let env ← Lean.MonadEnv.getEnv
+  match start with
+  | .copy name =>
+    let cInfo? := str_to_name name |> env.find?
+    add_test $ LSpec.check s!"Symbol exists {name}" cInfo?.isSome
+    match cInfo? with
+    | .some cInfo =>
+      let goal ← GoalState.create (expr := cInfo.type)
+      return Option.some goal
+    | .none =>
+      return Option.none
+  | .expr expr =>
+    let syn? := syntax_from_str env expr
+    add_test $ LSpec.check s!"Parsing {expr}" (syn?.isOk)
+    match syn? with
+    | .error error =>
+      IO.println error
+      return Option.none
+    | .ok syn =>
+      let expr? ← syntax_to_expr syn
+      add_test $ LSpec.check s!"Elaborating" expr?.isOk
+      match expr? with
+      | .error error =>
+        IO.println error
+        return Option.none
+      | .ok expr =>
+        let goal ← GoalState.create (expr := expr)
+        return Option.some goal
 
-def proof_runner (env: Lean.Environment) (options: Commands.Options) (start: Start) (steps: List (TestM LSpec.TestSeq)): IO LSpec.TestSeq := do
-  let termElabM := do
-    let (testSeq, state?) ← start_proof start
-    match state? with
-    | .none => return testSeq
-    | .some state => steps.foldlM (fun tests m => do pure $ tests ++ (← m)) testSeq |>.run options |>.run' state
-
-  let coreContext: Lean.Core.Context := {
-    currNamespace := str_to_name "Aniva",
-    openDecls := [],     -- No 'open' directives needed
-    fileName := "<Pantograph>",
-    fileMap := { source := "", positions := #[0], lines := #[1] }
-  }
-  let metaM := termElabM.run' (ctx := {
-    declName? := some "_pantograph",
-    errToSorry := false
-  })
-  let coreM := metaM.run'
-  match ← (coreM.run' coreContext { env := env }).toBaseIO with
-  | .error exception =>
-    return LSpec.test "Exception" (s!"internal exception #{← exception.toMessageData.toString}" = "")
-  | .ok a            => return a
+def assert_unreachable (message: String): LSpec.TestSeq := LSpec.check message false
 
 def build_goal (nameType: List (String × String)) (target: String): Commands.Goal :=
   {
@@ -101,30 +60,61 @@ def build_goal (nameType: List (String × String)) (target: String): Commands.Go
       isInaccessible? := .some false
     })).toArray
   }
+-- Like `build_goal` but allow certain variables to be elided.
+def build_goal_selective (nameType: List (String × Option String)) (target: String): Commands.Goal :=
+  {
+    target := { pp? := .some target},
+    vars := (nameType.map fun x => ({
+      name := x.fst,
+      type? := x.snd.map (λ type => { pp? := type }),
+      isInaccessible? := x.snd.map (λ _ => false)
+    })).toArray
+  }
 
+
+-- Individual test cases
 example: ∀ (a b: Nat), a + b = b + a := by
   intro n m
   rw [Nat.add_comm]
-def proof_nat_add_comm (env: Lean.Environment): IO LSpec.TestSeq := do
-  let goal1: Commands.Goal := build_goal [("n", "Nat"), ("m", "Nat")] "n + m = m + n"
-  proof_runner env {} (.copy "Nat.add_comm") [
-    proof_step 0 0 "intro n m"
-      (.success (.some 1) #[goal1]),
-    proof_step 1 0 "assumption"
-      (.failure #[s!"tactic 'assumption' failed\nn m : Nat\n⊢ n + m = m + n"]),
-    proof_step 1 0 "rw [Nat.add_comm]"
-      (.success .none #[])
-  ]
-def proof_nat_add_comm_manual (env: Lean.Environment): IO LSpec.TestSeq := do
-  let goal1: Commands.Goal := build_goal [("n", "Nat"), ("m", "Nat")] "n + m = m + n"
-  proof_runner env {} (.expr "∀ (a b: Nat), a + b = b + a") [
-    proof_step 0 0 "intro n m"
-      (.success (.some 1) #[goal1]),
-    proof_step 1 0 "assumption"
-      (.failure #[s!"tactic 'assumption' failed\nn m : Nat\n⊢ n + m = m + n"]),
-    proof_step 1 0 "rw [Nat.add_comm]"
-      (.success .none #[])
-  ]
+def proof_nat_add_comm: TestM Unit := do
+  let goal? ← start_proof (.copy "Nat.add_comm")
+  add_test $ LSpec.check "Start goal" goal?.isSome
+  if let .some goal := goal? then
+    if let .success #[(goal, sGoal)] ← goal.execute "intro n m" then
+      let sGoal1e: Commands.Goal := build_goal [("n", "Nat"), ("m", "Nat")] "n + m = m + n"
+      add_test $ LSpec.check "intro n m" (sGoal = sGoal1e)
+
+      if let .failure #[message] ← goal.execute "assumption" then
+        add_test $ LSpec.check "assumption" (message = "tactic 'assumption' failed\nn m : Nat\n⊢ n + m = m + n")
+      else
+        add_test $ assert_unreachable "assumption"
+
+      if let .success #[] ← goal.execute "rw [Nat.add_comm]" then
+        return ()
+      else
+        add_test $ assert_unreachable "rw [Nat.add_comm]"
+    else
+      add_test $ assert_unreachable "intro n m"
+def proof_nat_add_comm_manual: TestM Unit := do
+  let goal? ← start_proof (.expr "∀ (a b: Nat), a + b = b + a")
+  add_test $ LSpec.check "Start goal" goal?.isSome
+  if let .some goal := goal? then
+    if let .success #[(goal, sGoal)] ← goal.execute "intro n m" then
+      let sGoal1e: Commands.Goal := build_goal [("n", "Nat"), ("m", "Nat")] "n + m = m + n"
+      add_test $ LSpec.check "intro n m" (sGoal = sGoal1e)
+
+      if let .failure #[message] ← goal.execute "assumption" then
+        add_test $ LSpec.check "assumption" (message = "tactic 'assumption' failed\nn m : Nat\n⊢ n + m = m + n")
+      else
+        add_test $ assert_unreachable "assumption"
+
+      if let .success #[] ← goal.execute "rw [Nat.add_comm]" then
+        return ()
+      else
+        add_test $ assert_unreachable "rw [Nat.add_comm]"
+    else
+      add_test $ assert_unreachable "intro n m"
+
 
 -- Two ways to write the same theorem
 example: ∀ (p q: Prop), p ∨ q → q ∨ p := by
@@ -141,7 +131,7 @@ example: ∀ (p q: Prop), p ∨ q → q ∨ p := by
     assumption
   . apply Or.inl
     assumption
-def proof_or_comm (env: Lean.Environment): IO LSpec.TestSeq := do
+def proof_or_comm: TestM Unit := do
   let typeProp: Commands.Expression := { pp? := .some "Prop" }
   let branchGoal (caseName name: String): Commands.Goal := {
     caseName? := .some caseName,
@@ -152,69 +142,111 @@ def proof_or_comm (env: Lean.Environment): IO LSpec.TestSeq := do
       { name := "h✝", type? := .some { pp? := .some name }, isInaccessible? := .some true }
     ]
   }
-  proof_runner env {} (.expr "∀ (p q: Prop), p ∨ q → q ∨ p") [
-    proof_step 0 0 "intro p q h"
-      (.success (.some 1) #[build_goal [("p", "Prop"), ("q", "Prop"), ("h", "p ∨ q")] "q ∨ p"]),
-    proof_step 1 0 "cases h"
-      (.success (.some 2) #[branchGoal "inl" "p", branchGoal "inr" "q"]),
-    proof_inspect #["", "0.0", "1.0"],
-    proof_step 2 0 "apply Or.inr"
-      (.success (.some 3) #[]),
-    proof_inspect #["", "0.0", "1.0", "2.0"],
-    proof_step 3 0 "assumption"
-      (.success .none #[]),
-    proof_step 2 1 "apply Or.inl"
-      (.success (.some 4) #[]),
-    proof_step 4 0 "assumption"
-      (.success .none #[]),
-    proof_inspect #["", "0.0", "1.0", "2.0", "2.1"]
-  ]
+  let goal? ← start_proof (.expr "∀ (p q: Prop), p ∨ q → q ∨ p")
+  add_test $ LSpec.check "Start goal" goal?.isSome
+  if let .some goal := goal? then
+    if let .success #[(goal, sGoal)] ← goal.execute "intro p q h" then
+      let sGoal1e := build_goal [("p", "Prop"), ("q", "Prop"), ("h", "p ∨ q")] "q ∨ p"
+      add_test $ LSpec.check "intro p q h" (sGoal = sGoal1e)
+
+      if let .success #[(goal1, sGoal1), (goal2, sGoal2)] ← goal.execute "cases h" then
+        add_test $ LSpec.check "cases h/1" (sGoal1 = branchGoal "inl" "p")
+        if let .success #[(goal, _)] ← goal1.execute "apply Or.inr" then
+          if let .success #[] ← goal.execute "assumption" then
+            return ()
+          else
+            add_test $ assert_unreachable "assumption"
+        else
+          add_test $ assert_unreachable "apply Or.inr"
+
+
+        add_test $ LSpec.check "cases h/2" (sGoal2 = branchGoal "inr" "q")
+        if let .success #[(goal, _)] ← goal2.execute "apply Or.inl" then
+          if let .success #[] ← goal.execute "assumption" then
+            return ()
+          else
+            add_test $ assert_unreachable "assumption"
+        else
+          add_test $ assert_unreachable "apply Or.inl"
+
+      else
+        add_test $ assert_unreachable "cases h"
+    else
+      add_test $ assert_unreachable "intro p q h"
 
 example (w x y z : Nat) (p : Nat → Prop)
         (h : p (x * y + z * w * x)) : p (x * w * z + y * x) := by
   simp [Nat.add_assoc, Nat.add_comm, Nat.add_left_comm, Nat.mul_comm, Nat.mul_assoc, Nat.mul_left_comm] at *
   assumption
-def proof_arith_1 (env: Lean.Environment): IO LSpec.TestSeq := do
-  proof_runner env {} (.expr "∀ (w x y z : Nat) (p : Nat → Prop) (h : p (x * y + z * w * x)), p (x * w * z + y * x)") [
-    proof_step 0 0 "intros"
-      (.success (.some 1) #[]),
-    proof_step 1 0 "simp [Nat.add_assoc, Nat.add_comm, Nat.add_left_comm, Nat.mul_comm, Nat.mul_assoc, Nat.mul_left_comm] at *"
-      (.success (.some 2) #[]),
-    proof_step 2 0 "assumption"
-      (.success .none #[])
-  ]
+def proof_arith_1: TestM Unit := do
+  let goal? ← start_proof (.expr "∀ (w x y z : Nat) (p : Nat → Prop) (h : p (x * y + z * w * x)), p (x * w * z + y * x)")
+  add_test $ LSpec.check "Start goal" goal?.isSome
+  if let .some goal := goal? then
+    if let .success #[(goal, _)] ← goal.execute "intros" then
+      if let .success #[(goal, _)] ← goal.execute "simp [Nat.add_assoc, Nat.add_comm, Nat.add_left_comm, Nat.mul_comm, Nat.mul_assoc, Nat.mul_left_comm] at *" then
+        if let .success #[] ← goal.execute "assumption" then
+          return ()
+        else
+          add_test $ assert_unreachable "assumption"
+      else
+        add_test $ assert_unreachable "simp ..."
+    else
+      add_test $ assert_unreachable "intros"
 
-def build_goal_selective (nameType: List (String × Option String)) (target: String): Commands.Goal :=
-  {
-    target := { pp? := .some target},
-    vars := (nameType.map fun x => ({
-      name := x.fst,
-      type? := x.snd.map (λ type => { pp? := type }),
-      isInaccessible? := x.snd.map (λ _ => false)
-    })).toArray
+def proof_delta_variable: TestM Unit := withReader (fun _ => {proofVariableDelta := true}) do
+  let goal? ← start_proof (.expr "∀ (a b: Nat), a + b = b + a")
+  add_test $ LSpec.check "Start goal" goal?.isSome
+  if let .some goal := goal? then
+    if let .success #[(goal, sGoal)] ← goal.execute "intro n" then
+      let sGoal1e: Commands.Goal := build_goal_selective [("n", .some "Nat")] "∀ (b : Nat), n + b = b + n"
+      add_test $ LSpec.check "intro n" (sGoal = sGoal1e)
+
+      if let .success #[(_, sGoal)] ← goal.execute "intro m" then
+        let sGoal2e: Commands.Goal := build_goal_selective [("n", .none), ("m", .some "Nat")] "n + m = m + n"
+        add_test $ LSpec.check "intro m" (sGoal = sGoal2e)
+      else
+        add_test $ assert_unreachable "intro m"
+    else
+      add_test $ assert_unreachable "intro n"
+
+def proof_runner (env: Lean.Environment) (tests: TestM Unit): IO LSpec.TestSeq := do
+  let termElabM := tests.run LSpec.TestSeq.done |>.run {} -- with default options
+
+  let coreContext: Lean.Core.Context := {
+    currNamespace := str_to_name "Aniva",
+    openDecls := [],     -- No 'open' directives needed
+    fileName := "<Pantograph>",
+    fileMap := { source := "", positions := #[0], lines := #[1] }
   }
-def proof_delta_variable (env: Lean.Environment): IO LSpec.TestSeq := do
-  let goal1: Commands.Goal := build_goal_selective [("n", .some "Nat")] "∀ (b : Nat), n + b = b + n"
-  let goal2: Commands.Goal := build_goal_selective [("n", .none), ("m", .some "Nat")] "n + m = m + n"
-  proof_runner env { proofVariableDelta := true } (.expr "∀ (a b: Nat), a + b = b + a") [
-    proof_step 0 0 "intro n"
-      (.success (.some 1) #[goal1]),
-    proof_step 1 0 "intro m"
-      (.success (.some 2) #[goal2])
-  ]
+  let metaM := termElabM.run' (ctx := {
+    declName? := some "_pantograph",
+    errToSorry := false
+  })
+  let coreM := metaM.run'
+  match ← (coreM.run' coreContext { env := env }).toBaseIO with
+  | .error exception =>
+    return LSpec.test "Exception" (s!"internal exception #{← exception.toMessageData.toString}" = "")
+  | .ok (_, a) =>
+    return a
 
 def test_proofs : IO LSpec.TestSeq := do
   let env: Lean.Environment ← Lean.importModules
     (imports := ["Init"].map (λ str => { module := str_to_name str, runtimeOnly := false }))
     (opts := {})
     (trustLevel := 1)
+  let tests := [
+    ("Nat.add_comm", proof_nat_add_comm),
+    ("nat.add_comm manual", proof_nat_add_comm_manual),
+    ("Or.comm", proof_or_comm),
+    ("arithmetic 1", proof_arith_1),
+    ("delta variable", proof_delta_variable)
+  ]
+  let tests ← tests.foldlM (fun acc tests => do
+    let (name, tests) := tests
+    let tests ← proof_runner env tests
+    return acc ++ (LSpec.group name tests)) LSpec.TestSeq.done
 
-  return LSpec.group "Proofs" $
-    (LSpec.group "Nat.add_comm" $ (← proof_nat_add_comm env)) ++
-    (LSpec.group "Nat.add_comm manual" $ (← proof_nat_add_comm_manual env)) ++
-    (LSpec.group "Or.comm" $ (← proof_or_comm env)) ++
-    (LSpec.group "Arithmetic 1" $ (← proof_arith_1 env)) ++
-    (LSpec.group "Delta variable" $ (← proof_delta_variable env))
+  return LSpec.group "Proofs" tests
 
 end Pantograph.Test