Add SemihashMap structure for goal bookkeeping

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 := #[]
+  goalStates: SemihashMap GoalState := SemihashMap.empty
-  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 nTrees := state.proofTrees.size
+    let nStates := state.goalStates.size
-    set { state with proofTrees := #[] }
+    set { state with goalStates := SemihashMap.empty }
-    return .ok { nTrees := nTrees }
+    return .ok { nStates := nStates }
   lib_catalog (_: Commands.LibCatalog): MainM (CR Commands.LibCatalogResult) := do
     let env ← Lean.MonadEnv.getEnv
     let names := env.constants.fold (init := #[]) (λ acc name info =>
@@ -134,32 +134,34 @@ def execute (command: Commands.Command): MainM Lean.Json := do
     match expr? with
     | .error error => return .error error
     | .ok expr =>
-      let tree ← ProofTree.create expr
+      let goalState ← GoalState.create expr
-      -- Put the new tree in the environment
+      let (goalStates, goalId) := state.goalStates.insert goalState
-      let nextTreeId := state.proofTrees.size
+      set { state with goalStates }
-      set { state with proofTrees := state.proofTrees.push tree }
+      return .ok { goalId }
-      return .ok { treeId := nextTreeId }
   proof_tactic (args: Commands.ProofTactic): MainM (CR Commands.ProofTacticResult) := do
     let state ← get
-    match state.proofTrees.get? args.treeId with
+    match state.goalStates.get? args.goalId with
-    | .none => return .error $ errorIndex "Invalid tree index {args.treeId}"
+    | .none => return .error $ errorIndex "Invalid goal index {args.goalId}"
-    | .some tree =>
+    | .some goalState =>
-      let (result, nextTree) ← ProofTree.execute
+      let result ← GoalState.execute goalState args.tactic |>.run state.options
-        (stateId := args.stateId)
-        (goalId := args.goalId.getD 0)
-        (tactic := args.tactic) |>.run state.options |>.run tree
       match result with
-      | .invalid message => return .error $  errorIndex message
+      | .success goals =>
-      | .success nextId? goals =>
+        if goals.isEmpty then
-        set { state with proofTrees := state.proofTrees.set! args.treeId nextTree }
+          return .ok {}
-        return .ok { nextId? := nextId?, goals? := .some goals }
+        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
+  proof_print_tree (_: Commands.ProofPrintTree): MainM (CR Commands.ProofPrintTreeResult) := do
     let state ← get
-    match state.proofTrees.get? args.treeId with
+    return .ok { nGoals := state.goalStates.size }
-    | .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
-  nTrees: Nat
+  nStates: Nat
   deriving Lean.ToJson
 
 -- Return the type of an expression
@@ -133,29 +133,27 @@ structure ProofStart where
   copyFrom: Option String -- Theorem name
   deriving Lean.FromJson
 structure ProofStartResult where
-  treeId: Nat := 0 -- Proof tree id
+  goalId: Nat := 0 -- Proof tree id
   deriving Lean.ToJson
 structure ProofTactic where
   -- Identifiers for tree, state, and goal
-  treeId: Nat
+  goalId: Nat
-  stateId: Nat
-  goalId: Option Nat -- Defaults to 0
   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
-  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
   deriving Lean.FromJson
 structure ProofPrintTreeResult where
-  -- "" if no parents, otherwise "parentId.goalId"
+  -- Total number of goals
-  parents: Array String
+  nGoals: Nat
   deriving Lean.ToJson
 
 end Pantograph.Commands

Pantograph/SemihashMap.lean

@@ -0,0 +1,89 @@
+
+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,35 +24,22 @@ def Lean.MessageLog.getErrorMessages (log : MessageLog) : MessageLog :=
 namespace Pantograph
 open Lean
 
-structure ProofState where
+structure GoalState where
-  goals : List MVarId
+  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):
+def GoalState.execute (goal: GoalState) (tactic: String):
-    Commands.OptionsT StateRefT ProofTree M TacticResult := do
+    Commands.OptionsT M TacticResult := do
   let options ← read
-  let tree ← get
+  match (← execute_tactic (state := goal.savedState) (goal := goal.mvarId) (tactic := tactic)) with
-  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 .none #[]
+      return .success #[]
     else
-          let proofState: ProofState := {
+      let nextGoals: List GoalState := nextGoals.map fun mvarId => { mvarId, savedState := nextState }
-            savedState := nextState,
+      let parentDecl? := (← MonadMCtx.getMCtx).findDecl? goal.mvarId
-            goals := nextGoals,
+      let goals ← nextGoals.mapM fun nextGoal => do
-            parent := stateId,
+        match (← MonadMCtx.getMCtx).findDecl? nextGoal.mvarId with
-            parentGoalId := goalId
+        | .some mvarDecl =>
-          }
+          let serializedGoal ← serialize_goal options mvarDecl (parentDecl? := parentDecl?)
-          modify fun s => { s with states := s.states.push proofState }
+          return (nextGoal, serializedGoal)
-        let parentDecl? := (← MonadMCtx.getMCtx).findDecl? goal
+        | .none => throwError nextGoal.mvarId
-        let goals ← nextGoals.mapM fun mvarId => do
+      return .success goals.toArray
-          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 {"treeId": 0, "stateId": 0, "goalId": 0, "tactic": "intro n m"}
+proof.tactic {"goalId": 0, "tactic": "intro n m"}
-proof.tactic {"treeId": 0, "stateId": 1, "goalId": 0, "tactic": "assumption"}
+proof.tactic {"goalId": 1, "tactic": "assumption"}
-proof.printTree {"treeId": 0}
+proof.printTree {}
-proof.tactic {"treeId": 0, "stateId": 1, "goalId": 0, "tactic": "rw [Nat.add_comm]"}
+proof.tactic {"goalId": 1, "tactic": "rw [Nat.add_comm]"}
-proof.printTree {"treeId": 0}
+proof.printTree
 ```
 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 {"treeId": <id>, "stateId": <id>, "goalId": <id>, "tactic": string}`: Execute a tactic on a given proof state
+- `proof.tactic {"goalId": <id>, "tactic": <tactic>}`: Execute a tactic string on a given proof state
-- `proof.printTree {"treeId": <id>}`: Print the topological structure of a proof tree
+- `proof.printTree`: Print the number of goals
 
 ## 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
+abbrev TestM := StateRefT LSpec.TestSeq (ReaderT Commands.Options 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)
 
 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 add_test (test: LSpec.TestSeq): TestM Unit := do
-def proof_step (stateId: Nat) (goalId: Nat) (tactic: String)
+  set $ (← get) ++ test
-    (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)
 
-/-- Check that the tree structure is correct -/
+def start_proof (start: Start): TestM (Option GoalState) := do
-def proof_inspect (expected: Array String) : TestM LSpec.TestSeq := do
+  let env ← Lean.MonadEnv.getEnv
-  let result := (← get).structure_array
+  match start with
-  return LSpec.test s!"tree structure" (result = expected)
+  | .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
+def assert_unreachable (message: String): LSpec.TestSeq := LSpec.check message false
-  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 :=
   {
@@ -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
+def proof_nat_add_comm: TestM Unit := do
-  let goal1: Commands.Goal := build_goal [("n", "Nat"), ("m", "Nat")] "n + m = m + n"
+  let goal? ← start_proof (.copy "Nat.add_comm")
-  proof_runner env {} (.copy "Nat.add_comm") [
+  add_test $ LSpec.check "Start goal" goal?.isSome
-    proof_step 0 0 "intro n m"
+  if let .some goal := goal? then
-      (.success (.some 1) #[goal1]),
+    if let .success #[(goal, sGoal)] ← goal.execute "intro n m" then
-    proof_step 1 0 "assumption"
+      let sGoal1e: Commands.Goal := build_goal [("n", "Nat"), ("m", "Nat")] "n + m = m + n"
-      (.failure #[s!"tactic 'assumption' failed\nn m : Nat\n⊢ n + m = m + n"]),
+      add_test $ LSpec.check "intro n m" (sGoal = sGoal1e)
-    proof_step 1 0 "rw [Nat.add_comm]"
+
-      (.success .none #[])
+      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")
-def proof_nat_add_comm_manual (env: Lean.Environment): IO LSpec.TestSeq := do
+      else
-  let goal1: Commands.Goal := build_goal [("n", "Nat"), ("m", "Nat")] "n + m = m + n"
+        add_test $ assert_unreachable "assumption"
-  proof_runner env {} (.expr "∀ (a b: Nat), a + b = b + a") [
+
-    proof_step 0 0 "intro n m"
+      if let .success #[] ← goal.execute "rw [Nat.add_comm]" then
-      (.success (.some 1) #[goal1]),
+        return ()
-    proof_step 1 0 "assumption"
+      else
-      (.failure #[s!"tactic 'assumption' failed\nn m : Nat\n⊢ n + m = m + n"]),
+        add_test $ assert_unreachable "rw [Nat.add_comm]"
-    proof_step 1 0 "rw [Nat.add_comm]"
+    else
-      (.success .none #[])
+      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") [
+  let goal? ← start_proof (.expr "∀ (p q: Prop), p ∨ q → q ∨ p")
-    proof_step 0 0 "intro p q h"
+  add_test $ LSpec.check "Start goal" goal?.isSome
-      (.success (.some 1) #[build_goal [("p", "Prop"), ("q", "Prop"), ("h", "p ∨ q")] "q ∨ p"]),
+  if let .some goal := goal? then
-    proof_step 1 0 "cases h"
+    if let .success #[(goal, sGoal)] ← goal.execute "intro p q h" then
-      (.success (.some 2) #[branchGoal "inl" "p", branchGoal "inr" "q"]),
+      let sGoal1e := build_goal [("p", "Prop"), ("q", "Prop"), ("h", "p ∨ q")] "q ∨ p"
-    proof_inspect #["", "0.0", "1.0"],
+      add_test $ LSpec.check "intro p q h" (sGoal = sGoal1e)
-    proof_step 2 0 "apply Or.inr"
+
-      (.success (.some 3) #[]),
+      if let .success #[(goal1, sGoal1), (goal2, sGoal2)] ← goal.execute "cases h" then
-    proof_inspect #["", "0.0", "1.0", "2.0"],
+        add_test $ LSpec.check "cases h/1" (sGoal1 = branchGoal "inl" "p")
-    proof_step 3 0 "assumption"
+        if let .success #[(goal, _)] ← goal1.execute "apply Or.inr" then
-      (.success .none #[]),
+          if let .success #[] ← goal.execute "assumption" then
-    proof_step 2 1 "apply Or.inl"
+            return ()
-      (.success (.some 4) #[]),
+          else
-    proof_step 4 0 "assumption"
+            add_test $ assert_unreachable "assumption"
-      (.success .none #[]),
+        else
-    proof_inspect #["", "0.0", "1.0", "2.0", "2.1"]
+          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
+def proof_arith_1: TestM Unit := 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)") [
+  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_step 0 0 "intros"
+  add_test $ LSpec.check "Start goal" goal?.isSome
-      (.success (.some 1) #[]),
+  if let .some goal := goal? 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 "intros" then
-      (.success (.some 2) #[]),
+      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
-    proof_step 2 0 "assumption"
+        if let .success #[] ← goal.execute "assumption" then
-      (.success .none #[])
+          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 :=
+def proof_delta_variable: TestM Unit := withReader (fun _ => {proofVariableDelta := true}) do
-  {
+  let goal? ← start_proof (.expr "∀ (a b: Nat), a + b = b + a")
-    target := { pp? := .some target},
+  add_test $ LSpec.check "Start goal" goal?.isSome
-    vars := (nameType.map fun x => ({
+  if let .some goal := goal? then
-      name := x.fst,
+    if let .success #[(goal, sGoal)] ← goal.execute "intro n" then
-      type? := x.snd.map (λ type => { pp? := type }),
+      let sGoal1e: Commands.Goal := build_goal_selective [("n", .some "Nat")] "∀ (b : Nat), n + b = b + n"
-      isInaccessible? := x.snd.map (λ _ => false)
+      add_test $ LSpec.check "intro n" (sGoal = sGoal1e)
-    })).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] }
   }
-def proof_delta_variable (env: Lean.Environment): IO LSpec.TestSeq := do
+  let metaM := termElabM.run' (ctx := {
-  let goal1: Commands.Goal := build_goal_selective [("n", .some "Nat")] "∀ (b : Nat), n + b = b + n"
+    declName? := some "_pantograph",
-  let goal2: Commands.Goal := build_goal_selective [("n", .none), ("m", .some "Nat")] "n + m = m + n"
+    errToSorry := false
-  proof_runner env { proofVariableDelta := true } (.expr "∀ (a b: Nat), a + b = b + a") [
+  })
-    proof_step 0 0 "intro n"
+  let coreM := metaM.run'
-      (.success (.some 1) #[goal1]),
+  match ← (coreM.run' coreContext { env := env }).toBaseIO with
-    proof_step 1 0 "intro m"
+  | .error exception =>
-      (.success (.some 2) #[goal2])
+    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" $
+  return LSpec.group "Proofs" tests
-    (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