Pantograph.lean

@@ -2,7 +2,6 @@ import Pantograph.Commands
 import Pantograph.Serial
 import Pantograph.Symbols
 import Pantograph.Tactic
-import Pantograph.SemihashMap
 
 namespace Pantograph
 
@@ -12,7 +11,8 @@ structure Context where
 /-- Stores state of the REPL -/
 structure State where
   options: Commands.Options := {}
-  goalStates: SemihashMap GoalState := SemihashMap.empty
+  --environments: Array Lean.Environment := #[]
+  proofTrees:   Array ProofTree := #[]
 
 -- State monad
 abbrev MainM := ReaderT Context (StateT State Lean.Elab.TermElabM)
@@ -49,9 +49,9 @@ def execute (command: Commands.Command): MainM Lean.Json := do
   -- Command Functions
   reset (_: Commands.Reset): MainM (CR Commands.ResetResult) := do
     let state ← get
-    let nStates := state.goalStates.size
+    let nTrees := state.proofTrees.size
-    set { state with goalStates := SemihashMap.empty }
+    set { state with proofTrees := #[] }
-    return .ok { nStates := nStates }
+    return .ok { nTrees := nTrees }
   lib_catalog (_: Commands.LibCatalog): MainM (CR Commands.LibCatalogResult) := do
     let env ← Lean.MonadEnv.getEnv
     let names := env.constants.fold (init := #[]) (λ acc name info =>
@@ -134,34 +134,32 @@ def execute (command: Commands.Command): MainM Lean.Json := do
     match expr? with
     | .error error => return .error error
     | .ok expr =>
-      let goalState ← GoalState.create expr
+      let tree ← ProofTree.create (str_to_name <| args.name.getD "Untitled") expr
-      let (goalStates, goalId) := state.goalStates.insert goalState
+      -- Put the new tree in the environment
-      set { state with goalStates }
+      let nextTreeId := state.proofTrees.size
-      return .ok { goalId }
+      set { state with proofTrees := state.proofTrees.push tree }
+      return .ok { treeId := nextTreeId }
   proof_tactic (args: Commands.ProofTactic): MainM (CR Commands.ProofTacticResult) := do
     let state ← get
-    match state.goalStates.get? args.goalId with
+    match state.proofTrees.get? args.treeId with
-    | .none => return .error $ errorIndex "Invalid goal index {args.goalId}"
+    | .none => return .error $ errorIndex "Invalid tree index {args.treeId}"
-    | .some goalState =>
+    | .some tree =>
-      let result ← GoalState.execute goalState args.tactic |>.run state.options
+      let (result, nextTree) ← ProofTree.execute
+        (stateId := args.stateId)
+        (goalId := args.goalId.getD 0)
+        (tactic := args.tactic) |>.run state.options |>.run tree
       match result with
-      | .success goals =>
+      | .invalid message => return .error $  errorIndex message
-        if goals.isEmpty then
+      | .success nextId? goals =>
-          return .ok {}
+        set { state with proofTrees := state.proofTrees.set! args.treeId nextTree }
-        else
+        return .ok { nextId? := nextId?, goals? := .some goals }
-          -- 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 (_: Commands.ProofPrintTree): MainM (CR Commands.ProofPrintTreeResult) := do
+  proof_print_tree (args: Commands.ProofPrintTree): MainM (CR Commands.ProofPrintTreeResult) := do
     let state ← get
-    return .ok { nGoals := state.goalStates.size }
+    match state.proofTrees.get? args.treeId with
+    | .none => return .error $ errorIndex "Invalid tree index {args.treeId}"
+    | .some tree =>
+      return .ok { parents := tree.structure_array }
 
 end Pantograph

Pantograph/Commands.lean

@@ -81,7 +81,7 @@ structure InteractionError where
 structure Reset where
   deriving Lean.FromJson
 structure ResetResult where
-  nStates: Nat
+  nTrees: Nat
   deriving Lean.ToJson
 
 -- Return the type of an expression
@@ -128,32 +128,35 @@ structure OptionsPrint where
 abbrev OptionsPrintResult := Options
 
 structure ProofStart where
+  name: Option String     -- Identifier of the proof
   -- 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
-  goalId: Nat := 0 -- Proof tree id
+  treeId: Nat := 0 -- Proof tree id
   deriving Lean.ToJson
 structure ProofTactic where
   -- Identifiers for tree, state, and goal
-  goalId: Nat
+  treeId: Nat
+  stateId: Nat
+  goalId: Option Nat
   tactic: String
   deriving Lean.FromJson
 structure ProofTacticResult where
   -- Existence of this field shows success
   goals?: Option (Array Goal)          := .none
   -- Next proof state id, if successful
-  goalIds?: Option (Array Nat)          := .none
+  nextId?: Option Nat                  := .none
   -- Existence of this field shows failure
   tacticErrors?: Option (Array String) := .none
   deriving Lean.ToJson
-
 structure ProofPrintTree where
+  treeId: Nat
   deriving Lean.FromJson
 structure ProofPrintTreeResult where
-  -- Total number of goals
+  -- "" if no parents, otherwise "parentId.goalId"
-  nGoals: Nat
+  parents: Array String
   deriving Lean.ToJson
 
 end Pantograph.Commands

Pantograph/SemihashMap.lean

@@ -1,89 +0,0 @@
-
-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
-
-
-@[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: Nat): Imp β :=
-  match map.data.getD index .none 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: Nat): Option β :=
-  map.data.getD index .none
-@[inline] def capacity (map: Imp β): Nat := map.data.size
-
-end Imp
-
-def empty (capacity := 16): Imp β :=
-  {
-    data := mkArray capacity .none,
-    size := 0,
-    allocFront := 0,
-    deallocs := mkArray capacity 0,
-    lastDealloc := 0,
-  }
-
-/--
-This is like a hashmap but you cannot control the keys.
--/
-def _root_.Pantograph.SemihashMap β := Imp β
-
-end Pantograph.SemihashMap

Pantograph/Tactic.lean

@@ -24,22 +24,39 @@ def Lean.MessageLog.getErrorMessages (log : MessageLog) : MessageLog :=
 namespace Pantograph
 open Lean
 
-structure GoalState where
+structure ProofState where
-  mvarId: MVarId
+  goals : List MVarId
   savedState : Elab.Tactic.SavedState
+  parent : Option Nat := none
+  parentGoalId : Nat  := 0
+structure ProofTree where
+  -- All parameters needed to run a `TermElabM` monad
+  name: Name
+
+  -- Set of proof states
+  states : Array ProofState := #[]
 
 abbrev M := Elab.TermElabM
 
-def GoalState.create (expr: Expr): M GoalState := do
+def ProofTree.create (name: Name) (expr: Expr): M ProofTree := 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 {
+    name := name,
+    states := #[{
       savedState := savedState,
-    mvarId := goal.mvarId!
+      goals := [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
@@ -65,38 +82,44 @@ 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 (goals: Array (GoalState × Commands.Goal))
+  | success (nextId?: Option Nat) (goals: Array 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 GoalState.execute (goal: GoalState) (tactic: String):
+def ProofTree.execute (stateId: Nat) (goalId: Nat) (tactic: String):
-    Commands.OptionsT M TacticResult := do
+    Commands.OptionsT StateRefT ProofTree M TacticResult := do
   let options ← read
-  match (← execute_tactic (state := goal.savedState) (goal := goal.mvarId) (tactic := tactic)) with
+  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
       | .error errors =>
         return .failure errors
       | .ok (nextState, nextGoals) =>
+        let nextId := tree.states.size
         if nextGoals.isEmpty then
-      return .success #[]
+          return .success .none #[]
         else
-      let nextGoals: List GoalState := nextGoals.map fun mvarId => { mvarId, savedState := nextState }
+          let proofState: ProofState := {
-      let parentDecl? := (← MonadMCtx.getMCtx).findDecl? goal.mvarId
+            savedState := nextState,
-      let goals ← nextGoals.mapM fun nextGoal => do
+            goals := nextGoals,
-        match (← MonadMCtx.getMCtx).findDecl? nextGoal.mvarId with
+            parent := stateId,
-        | .some mvarDecl =>
+            parentGoalId := goalId
-          let serializedGoal ← serialize_goal options mvarDecl (parentDecl? := parentDecl?)
+          }
-          return (nextGoal, serializedGoal)
+          modify fun s => { s with states := s.states.push proofState }
-        | .none => throwError nextGoal.mvarId
+        let parentDecl? := (← MonadMCtx.getMCtx).findDecl? goal
-      return .success goals.toArray
+        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
 
 end Pantograph

README.md

@@ -54,11 +54,11 @@ Example proving a theorem: (alternatively use `proof.start {"copyFrom": "Nat.add
 ```
 $ env build/bin/Pantograph Init
 proof.start {"expr": "∀ (n m : Nat), n + m = m + n"}
-proof.tactic {"goalId": 0, "tactic": "intro n m"}
+proof.tactic {"treeId": 0, "stateId": 0, "goalId": 0, "tactic": "intro n m"}
-proof.tactic {"goalId": 1, "tactic": "assumption"}
+proof.tactic {"treeId": 0, "stateId": 1, "goalId": 0, "tactic": "assumption"}
-proof.printTree {}
+proof.printTree {"treeId": 0}
-proof.tactic {"goalId": 1, "tactic": "rw [Nat.add_comm]"}
+proof.tactic {"treeId": 0, "stateId": 1, "goalId": 0, "tactic": "rw [Nat.add_comm]"}
-proof.printTree
+proof.printTree {"treeId": 0}
 ```
 where the application of `assumption` should lead to a failure.
 
@@ -75,8 +75,8 @@ See `Pantograph/Commands.lean` for a description of the parameters and return va
   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 {"goalId": <id>, "tactic": <tactic>}`: Execute a tactic string on a given proof state
+- `proof.tactic {"treeId": <id>, "stateId": <id>, "goalId": <id>, "tactic": string}`: Execute a tactic on a given proof state
-- `proof.printTree`: Print the number of goals
+- `proof.printTree {"treeId": <id>}`: Print the topological structure of a proof tree
 
 ## Errors
 

Test/Proofs.lean

@@ -10,46 +10,89 @@ inductive Start where
   | copy (name: String) -- Start from some name in the environment
   | expr (expr: String) -- Start from some expression
 
-abbrev TestM := StateRefT LSpec.TestSeq (ReaderT Commands.Options M)
+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
+        (name := str_to_name "TestExample")
+        (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
+          (name := str_to_name "TestExample")
+          (expr := expr)
+        return (testSeq, Option.some state)
 
 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
 
-def add_test (test: LSpec.TestSeq): TestM Unit := do
+/-- Check the output of each proof step -/
-  set $ (← get) ++ test
+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 start_proof (start: Start): TestM (Option GoalState) := do
+/-- Check that the tree structure is correct -/
-  let env ← Lean.MonadEnv.getEnv
+def proof_inspect (expected: Array String) : TestM LSpec.TestSeq := do
-  match start with
+  let result := (← get).structure_array
-  | .copy name =>
+  return LSpec.test s!"tree structure" (result = expected)
-    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 assert_unreachable (message: String): LSpec.TestSeq := LSpec.check message false
+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 build_goal (nameType: List (String × String)) (target: String): Commands.Goal :=
   {
@@ -60,61 +103,30 @@ 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: TestM Unit := do
+def proof_nat_add_comm (env: Lean.Environment): IO LSpec.TestSeq := do
-  let goal? ← start_proof (.copy "Nat.add_comm")
+  let goal1: Commands.Goal := build_goal [("n", "Nat"), ("m", "Nat")] "n + m = m + n"
-  add_test $ LSpec.check "Start goal" goal?.isSome
+  proof_runner env {} (.copy "Nat.add_comm") [
-  if let .some goal := goal? then
+    proof_step 0 0 "intro n m"
-    if let .success #[(goal, sGoal)] ← goal.execute "intro n m" then
+      (.success (.some 1) #[goal1]),
-      let sGoal1e: Commands.Goal := build_goal [("n", "Nat"), ("m", "Nat")] "n + m = m + n"
+    proof_step 1 0 "assumption"
-      add_test $ LSpec.check "intro n m" (sGoal = sGoal1e)
+      (.failure #[s!"tactic 'assumption' failed\nn m : Nat\n⊢ n + m = m + n"]),
-
+    proof_step 1 0 "rw [Nat.add_comm]"
-      if let .failure #[message] ← goal.execute "assumption" then
+      (.success .none #[])
-        add_test $ LSpec.check "assumption" (message = "tactic 'assumption' failed\nn m : Nat\n⊢ n + m = m + n")
+  ]
-      else
+def proof_nat_add_comm_manual (env: Lean.Environment): IO LSpec.TestSeq := do
-        add_test $ assert_unreachable "assumption"
+  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") [
-      if let .success #[] ← goal.execute "rw [Nat.add_comm]" then
+    proof_step 0 0 "intro n m"
-        return ()
+      (.success (.some 1) #[goal1]),
-      else
+    proof_step 1 0 "assumption"
-        add_test $ assert_unreachable "rw [Nat.add_comm]"
+      (.failure #[s!"tactic 'assumption' failed\nn m : Nat\n⊢ n + m = m + n"]),
-    else
+    proof_step 1 0 "rw [Nat.add_comm]"
-      add_test $ assert_unreachable "intro n m"
+      (.success .none #[])
-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
@@ -131,7 +143,7 @@ example: ∀ (p q: Prop), p ∨ q → q ∨ p := by
     assumption
   . apply Or.inl
     assumption
-def proof_or_comm: TestM Unit := do
+def proof_or_comm (env: Lean.Environment): IO LSpec.TestSeq := do
   let typeProp: Commands.Expression := { pp? := .some "Prop" }
   let branchGoal (caseName name: String): Commands.Goal := {
     caseName? := .some caseName,
@@ -142,111 +154,69 @@ def proof_or_comm: TestM Unit := do
       { name := "h✝", type? := .some { pp? := .some name }, isInaccessible? := .some true }
     ]
   }
-  let goal? ← start_proof (.expr "∀ (p q: Prop), p ∨ q → q ∨ p")
+  proof_runner env {} (.expr "∀ (p q: Prop), p ∨ q → q ∨ p") [
-  add_test $ LSpec.check "Start goal" goal?.isSome
+    proof_step 0 0 "intro p q h"
-  if let .some goal := goal? then
+      (.success (.some 1) #[build_goal [("p", "Prop"), ("q", "Prop"), ("h", "p ∨ q")] "q ∨ p"]),
-    if let .success #[(goal, sGoal)] ← goal.execute "intro p q h" then
+    proof_step 1 0 "cases h"
-      let sGoal1e := build_goal [("p", "Prop"), ("q", "Prop"), ("h", "p ∨ q")] "q ∨ p"
+      (.success (.some 2) #[branchGoal "inl" "p", branchGoal "inr" "q"]),
-      add_test $ LSpec.check "intro p q h" (sGoal = sGoal1e)
+    proof_inspect #["", "0.0", "1.0"],
-
+    proof_step 2 0 "apply Or.inr"
-      if let .success #[(goal1, sGoal1), (goal2, sGoal2)] ← goal.execute "cases h" then
+      (.success (.some 3) #[]),
-        add_test $ LSpec.check "cases h/1" (sGoal1 = branchGoal "inl" "p")
+    proof_inspect #["", "0.0", "1.0", "2.0"],
-        if let .success #[(goal, _)] ← goal1.execute "apply Or.inr" then
+    proof_step 3 0 "assumption"
-          if let .success #[] ← goal.execute "assumption" then
+      (.success .none #[]),
-            return ()
+    proof_step 2 1 "apply Or.inl"
-          else
+      (.success (.some 4) #[]),
-            add_test $ assert_unreachable "assumption"
+    proof_step 4 0 "assumption"
-        else
+      (.success .none #[]),
-          add_test $ assert_unreachable "apply Or.inr"
+    proof_inspect #["", "0.0", "1.0", "2.0", "2.1"]
-
+  ]
-
-        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: TestM Unit := do
+def proof_arith_1 (env: Lean.Environment): IO LSpec.TestSeq := 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)")
+  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)") [
-  add_test $ LSpec.check "Start goal" goal?.isSome
+    proof_step 0 0 "intros"
-  if let .some goal := goal? then
+      (.success (.some 1) #[]),
-    if let .success #[(goal, _)] ← goal.execute "intros" then
+    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 *"
-      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
+      (.success (.some 2) #[]),
-        if let .success #[] ← goal.execute "assumption" then
+    proof_step 2 0 "assumption"
-          return ()
+      (.success .none #[])
-        else
+  ]
-          add_test $ assert_unreachable "assumption"
-      else
-        add_test $ assert_unreachable "simp ..."
-    else
-      add_test $ assert_unreachable "intros"
 
-def proof_delta_variable: TestM Unit := withReader (fun _ => {proofVariableDelta := true}) do
+def build_goal_selective (nameType: List (String × Option String)) (target: String): Commands.Goal :=
-  let goal? ← start_proof (.expr "∀ (a b: Nat), a + b = b + a")
+  {
-  add_test $ LSpec.check "Start goal" goal?.isSome
+    target := { pp? := .some target},
-  if let .some goal := goal? then
+    vars := (nameType.map fun x => ({
-    if let .success #[(goal, sGoal)] ← goal.execute "intro n" then
+      name := x.fst,
-      let sGoal1e: Commands.Goal := build_goal_selective [("n", .some "Nat")] "∀ (b : Nat), n + b = b + n"
+      type? := x.snd.map (λ type => { pp? := type }),
-      add_test $ LSpec.check "intro n" (sGoal = sGoal1e)
+      isInaccessible? := x.snd.map (λ _ => false)
-
+    })).toArray
-      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] }
   }
-  let metaM := termElabM.run' (ctx := {
+def proof_delta_variable (env: Lean.Environment): IO LSpec.TestSeq := do
-    declName? := some "_pantograph",
+  let goal1: Commands.Goal := build_goal_selective [("n", .some "Nat")] "∀ (b : Nat), n + b = b + n"
-    errToSorry := false
+  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") [
-  let coreM := metaM.run'
+    proof_step 0 0 "intro n"
-  match ← (coreM.run' coreContext { env := env }).toBaseIO with
+      (.success (.some 1) #[goal1]),
-  | .error exception =>
+    proof_step 1 0 "intro m"
-    return LSpec.test "Exception" (s!"internal exception #{← exception.toMessageData.toString}" = "")
+      (.success (.some 2) #[goal2])
-  | .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" tests
+  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))
 
 end Pantograph.Test