Use TermElabM as the main monad stack instead of IO

Main.lean

@@ -10,28 +10,14 @@ namespace Pantograph
 
 
 structure Context where
-  coreContext: Lean.Core.Context
 
 /-- Stores state of the REPL -/
 structure State where
-  environments: Array Lean.Environment := #[]
+  --environments: Array Lean.Environment := #[]
   proofTrees:   Array Meta.ProofTree := #[]
 
 -- State monad
-abbrev Subroutine := ReaderT Context (StateT State IO)
-
-def State.getEnv (state: State) (id: Nat): Except String Lean.Environment :=
-  match state.environments.get? id with
-  | .some env => return env
-  | .none => throw s!"Invalid environment id {id}"
-
-
--- Utilities
-def option_expect (o: Option α) (error: String): Except String α :=
-  match o with
-  | .some value => return value
-  | .none => throw error
-
+abbrev Subroutine := ReaderT Context (StateT State Lean.Elab.TermElabM)
 
 open Commands
 
@@ -52,20 +38,12 @@ def parse_command (s: String): Except String Command := do
 
 
 
-unsafe def execute (command: Command): Subroutine Lean.Json := do
+def execute (command: Command): Subroutine Lean.Json := do
   match command.cmd with
-  | "create" =>
-    match Lean.fromJson? command.payload with
-    | .ok args => create args
-    | .error x => return errorJson x
   | "catalog" =>
     match Lean.fromJson? command.payload with
     | .ok args => catalog args
     | .error x => return errorJson x
-  | "clear" =>
-    -- Delete all the environments
-    let ret ← clear
-    return Lean.toJson ret
   | "inspect" =>
     match Lean.fromJson? command.payload with
     | .ok args => inspect args
@@ -88,93 +66,49 @@ unsafe def execute (command: Command): Subroutine Lean.Json := do
   where
   errorJson (s: String) := Lean.toJson ({ error := "json", desc := s }: InteractionError)
   errorIndex (s: String) := Lean.toJson ({ error := "index", desc := s }: InteractionError)
-  create (args: Create): Subroutine Lean.Json := do
-    let state ← get
-    let envId := state.environments.size
-    let env ← Lean.importModules
-      (imports := args.imports.map (λ str => { module := str_to_name str, runtimeOnly := false }))
-      (opts := {})
-      (trustLevel := 1)
-    modify fun s => { environments := s.environments.push env }
-    let num_filtered_symbols := env.constants.fold (init := 0) (λ acc name info =>
-      acc + if is_symbol_unsafe_or_internal name info then 0 else 1)
-    return Lean.toJson ({
-      envId := envId,
-      symbols := env.constants.size,
-      filtered_symbols := num_filtered_symbols }: CreateResult)
-  catalog (args: Catalog): Subroutine Lean.Json := do
-    let state ← get
-    match state.getEnv args.envId with
-    | .error error => return Lean.toJson <| errorIndex error
-    | .ok env =>
-      let names := env.constants.fold (init := []) (λ es name info =>
-        match to_filtered_symbol name info with
-        | .some x => x::es
-        | .none => es)
-      return Lean.toJson <| ({ symbols := names }: CatalogResult)
-  clear: Subroutine Lean.Json := do
-    let state ← get
-    let nEnv := state.environments.size
-    for env in state.environments do
-      env.freeRegions
-    set ({ }: State)
-    return Lean.toJson ({ nEnv := nEnv }: ClearResult)
+  catalog (_: Catalog): Subroutine Lean.Json := do
+    let env ← Lean.MonadEnv.getEnv
+    let names := env.constants.fold (init := []) (λ es name info =>
+      match to_filtered_symbol name info with
+      | .some x => x::es
+      | .none => es)
+    return Lean.toJson <| ({ symbols := names }: CatalogResult)
   inspect (args: Inspect): Subroutine Lean.Json := do
-    let context ← read
-    let state ← get
-    match state.getEnv args.envId with
-    | .error error => return Lean.toJson <| errorIndex error
-    | .ok env =>
-      let name := str_to_name args.name
-      let info? := env.find? name
-      match info? with
-      | none => return Lean.toJson <| errorIndex s!"Symbol not found {args.name}"
-      | some info =>
-        let format ← Serial.expr_to_str
-          (env := env)
-          (coreContext := context.coreContext)
-          (expr := info.toConstantVal.type)
-        let module? := env.getModuleIdxFor? name >>=
-          (λ idx => env.allImportedModuleNames.get? idx.toNat) |>.map toString
-        return Lean.toJson ({
-          type := format,
-          module? := module?
-        }: InspectResult)
+    let env ← Lean.MonadEnv.getEnv
+    let name := str_to_name args.name
+    let info? := env.find? name
+    match info? with
+    | none => return Lean.toJson <| errorIndex s!"Symbol not found {args.name}"
+    | some info =>
+      let format ← Lean.Meta.ppExpr info.toConstantVal.type
+      let module? := env.getModuleIdxFor? name >>=
+        (λ idx => env.allImportedModuleNames.get? idx.toNat) |>.map toString
+      return Lean.toJson ({
+        type := toString format,
+        module? := module?
+      }: InspectResult)
   proof_start (args: ProofStart): Subroutine Lean.Json := do
-    let context ← read
     let state ← get
-    let ret?: Except Lean.Json Meta.ProofTree ← ExceptT.run <| (do
-      let env ← match state.getEnv args.envId with
-        | .error error => throw <| Lean.toJson <| errorIndex error
-        | .ok env => pure env
-      let tree := Meta.createProofTree
-        (name := args.name.getD "Untitled")
-        (env := env)
-        (coreContext := context.coreContext)
-      let expr: Lean.Expr ← match args.expr, args.copyFrom with
-        | .some expr, .none =>
-          let syn ← match Serial.syntax_from_str env expr with
-            | .error str => throw <| Lean.toJson ({ error := "parsing", desc := str }: InteractionError)
-            | .ok syn => pure syn
-          let expr: Lean.Expr ← match (← Meta.ProofM.syntax_to_expr syn |>.run' tree) with
-            | .error str => throw <| Lean.toJson ({ error := "elab", desc := str }: InteractionError)
-            | .ok expr => pure expr
-          pure expr
-        | .none, .some copyFrom =>
-          match Serial.expr_from_const env (name := str_to_name copyFrom) with
-          | .error str =>
-            IO.println "Symbol not found"
-            throw <| errorIndex str
-          | .ok expr => pure expr
-        | .none, .none =>
-          throw <| Lean.toJson ({ error := "arguments", desc := "At least one of {expr, copyFrom} must be supplied" }: InteractionError)
-        | _, _ => throw <| Lean.toJson ({ error := "arguments", desc := "Cannot populate both of {expr, copyFrom}" }: InteractionError)
-      let (_, tree) := ← (Meta.ProofM.start expr |>.run tree)
-      return tree
-    )
-    match ret? with
+    let env ← Lean.MonadEnv.getEnv
+    let expr?: Except Lean.Json Lean.Expr ← (match args.expr, args.copyFrom with
+      | .some expr, .none =>
+        (match Serial.syntax_from_str env expr with
+        | .error str => return .error <| Lean.toJson ({ error := "parsing", desc := str }: InteractionError)
+        | .ok syn => do
+          (match (← Meta.syntax_to_expr syn) with
+          | .error str => return .error <| Lean.toJson ({ error := "elab", desc := str }: InteractionError)
+          | .ok expr => return .ok expr))
+      | .none, .some copyFrom =>
+        (match env.find? <| str_to_name copyFrom with
+        | .none => return .error <| errorIndex s!"Symbol not found: {copyFrom}"
+        | .some cInfo => return .ok cInfo.type)
+      | .none, .none =>
+        return .error <| Lean.toJson ({ error := "arguments", desc := "At least one of {expr, copyFrom} must be supplied" }: InteractionError)
+      | _, _ => return .error <| Lean.toJson ({ error := "arguments", desc := "Cannot populate both of {expr, copyFrom}" }: InteractionError))
+    match expr? with
     | .error error => return error
-    | .ok tree =>
+    | .ok expr =>
+      let tree ← Meta.ProofTree.create (str_to_name <| args.name.getD "Untitled") expr
       -- Put the new tree in the environment
       let nextTreeId := state.proofTrees.size
       set { state with proofTrees := state.proofTrees.push tree }
@@ -184,7 +118,7 @@ unsafe def execute (command: Command): Subroutine Lean.Json := do
     match state.proofTrees.get? args.treeId with
     | .none => return Lean.toJson <| errorIndex "Invalid tree index {args.treeId}"
     | .some tree =>
-      let (result, nextTree) ← Meta.ProofM.execute
+      let (result, nextTree) ← Meta.ProofTree.execute
         (stateId := args.stateId)
         (goalId := args.goalId.getD 0)
         (tactic := args.tactic) |>.run tree
@@ -221,19 +155,30 @@ unsafe def loop : Subroutine Unit := do
     IO.println <| toString <| ret
   loop
 
-unsafe def main : IO Unit := do
+unsafe def main (args: List String): IO Unit := do
   Lean.enableInitializersExecution
   Lean.initSearchPath (← Lean.findSysroot)
+
+  let env ← Lean.importModules
+    (imports := args.map (λ str => { module := str_to_name str, runtimeOnly := false }))
+    (opts := {})
+    (trustLevel := 1)
   let context: Context := {
-    coreContext := {
-      currNamespace := str_to_name "Aniva",
-      openDecls := [],     -- No 'open' directives needed
-      fileName := "<Pantograph>",
-      fileMap := { source := "", positions := #[0], lines := #[1] }
-    }
+  }
+  let coreContext: Lean.Core.Context := {
+    currNamespace := str_to_name "Aniva",
+    openDecls := [],     -- No 'open' directives needed
+    fileName := "<Pantograph>",
+    fileMap := { source := "", positions := #[0], lines := #[1] }
   }
   try
-    loop.run context |>.run' {}
+    let termElabM := loop.run context |>.run' {}
+    let metaM := termElabM.run' (ctx := {
+      declName? := some "_pantograph",
+      errToSorry := false
+    })
+    let coreM := metaM.run'
+    discard <| coreM.toIO coreContext { env := env }
   catch ex =>
     IO.println "Uncaught IO exception"
     IO.println ex.toString

Pantograph/Commands.lean

@@ -16,32 +16,16 @@ structure InteractionError where
 
 -- Individual command and return types
 
--- Create a new environment using the given imports
-structure Create where
-  imports : List String  := []
-  deriving Lean.FromJson
-structure CreateResult where
-  envId: Nat
-  symbols: Nat
-  filtered_symbols: Nat
-  deriving Lean.ToJson
 
 -- Print all symbols in environment
 structure Catalog where
-  envId: Nat
   deriving Lean.FromJson
 structure CatalogResult where
   symbols: List String
   deriving Lean.ToJson
 
--- Reset the state of REPL
-structure ClearResult where
-  nEnv: Nat   -- Number of environments reset
-  deriving Lean.ToJson
-
 -- Print the type of a symbol
 structure Inspect where
-  envId: Nat -- Environment id
   name: String
   deriving Lean.FromJson
 structure InspectResult where
@@ -50,7 +34,6 @@ structure InspectResult where
   deriving Lean.ToJson
 
 structure ProofStart where
-  envId: Nat -- Environment id
   name: Option String     -- Identifier of the proof
   -- Only one of the fields below may be populated.
   expr: Option String     -- Proof expression

Pantograph/Meta.lean

@@ -22,133 +22,74 @@ def Lean.MessageLog.getErrorMessages (log : Lean.MessageLog) : Lean.MessageLog :
 namespace Pantograph.Meta
 
 structure ProofState where
+  goals : List Lean.MVarId
   savedState : Lean.Elab.Tactic.SavedState
   parent : Option Nat := none
   parentGoalId : Nat  := 0
 structure ProofTree where
   -- All parameters needed to run a `TermElabM` monad
   name: Lean.Name
-  coreContext : Lean.Core.Context
-  elabContext : Lean.Elab.Term.Context
-
-  /-
-  This state must be saved so it preserves existing variable assignments. See
-
-    https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Resume.20proof.20in.20IO.20monad/near/360429763
-
-  It is unknown what will happen to this in the case of backtracking. Since we
-  never delete any proof states, it should be fine to store this here for now. A
-  test case `Or.comm` illustrates branching which will fail if the core state is
-  replaced every time.
-  -/
-  coreState : Lean.Core.State
 
   -- Set of proof states
   states : Array ProofState := #[]
 
-abbrev ProofM := StateRefT ProofTree IO
+abbrev M := Lean.Elab.TermElabM
 
-def createProofTree (name: Lean.Name) (env: Lean.Environment) (coreContext: Lean.Core.Context): ProofTree :=
-  {
+def ProofTree.create (name: Lean.Name) (expr: Lean.Expr): M ProofTree := do
+  let expr ← Lean.instantiateMVars expr
+  let goal := (← Lean.Meta.mkFreshExprMVar expr (kind := Lean.MetavarKind.synthetic))
+  let savedStateMonad: Lean.Elab.Tactic.TacticM Lean.Elab.Tactic.SavedState := Lean.MonadBacktrack.saveState
+  let savedState ← savedStateMonad { elaborator := .anonymous } |>.run' { goals := [goal.mvarId!]}
+  return {
     name := name,
-    coreContext := coreContext,
-    elabContext := {
-      declName? := some (name ++ "_pantograph"),
-      errToSorry := false
-    }
-    coreState := {
-      env := env
-    }
+    states := #[{
+      savedState := savedState,
+      goals := [goal.mvarId!]
+    }]
   }
 
--- Tree structures
-
+-- 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}"
 
--- Executing a `TermElabM`
-def ProofM.runTermElabM (termElabM: Lean.Elab.TermElabM α): ProofM (α × Lean.Core.State) := do
-  let context ← get
-  let metaM : Lean.MetaM α := termElabM.run' (ctx := context.elabContext)
-  let coreM : Lean.CoreM α := metaM.run'
-  coreM.toIO context.coreContext context.coreState
-def ProofM.runTermElabM' (termElabM: Lean.Elab.TermElabM α): ProofM α := do
-  let (ret, coreState) ← ProofM.runTermElabM termElabM
-  set { ← get with coreState := coreState }
-  return ret
-
 -- Parsing syntax under the environment
-def ProofM.syntax_to_expr (syn: Lean.Syntax): ProofM (Except String Lean.Expr) := do
-  let termElabM : Lean.Elab.TermElabM (Except String Lean.Expr) :=
+def syntax_to_expr (syn: Lean.Syntax): Lean.Elab.TermElabM (Except String Lean.Expr) := do
+  try
+    let expr ← Lean.Elab.Term.elabType syn
+    -- Immediately synthesise all metavariables if we need to leave the elaboration context.
+    -- See https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Unknown.20universe.20metavariable/near/360130070
+    --Lean.Elab.Term.synthesizeSyntheticMVarsNoPostponing
+    let expr ← Lean.instantiateMVars expr
+    return .ok expr
+  catch ex => return .error (← ex.toMessageData.toString)
+
+def execute_tactic (state: Lean.Elab.Tactic.SavedState) (goal: Lean.MVarId) (tactic: String) :
+    M (Except (Array String) (Lean.Elab.Tactic.SavedState × List Lean.MVarId)):= do
+  let tacticM (stx: Lean.Syntax):  Lean.Elab.Tactic.TacticM (Except (Array String) (Lean.Elab.Tactic.SavedState × List Lean.MVarId)) := do
+    state.restore
+    Lean.Elab.Tactic.setGoals [goal]
     try
-      --let expr ← Lean.Elab.Term.elabTerm syn
-      --  (expectedType? := none)
-      --  (catchExPostpone := false)
-      --  (implicitLambda := true)
-      let expr ← Lean.Elab.Term.elabType syn
-
-      -- Immediately synthesise all metavariables since we need to leave the elaboration context.
-      -- See https://leanprover.zulipchat.com/#narrow/stream/270676-lean4/topic/Unknown.20universe.20metavariable/near/360130070
-      Lean.Elab.Term.synthesizeSyntheticMVarsNoPostponing
-      let expr ← Lean.instantiateMVars expr
-
-      return .ok expr
-
-    catch ex => return .error (← ex.toMessageData.toString)
-  ProofM.runTermElabM' <| termElabM
-
-def start_tactic_state (expr: Lean.Expr): Lean.Elab.TermElabM Lean.Elab.Tactic.SavedState := do
-    let mvar ← Lean.Meta.mkFreshExprMVar (some expr) (kind := Lean.MetavarKind.synthetic)
-    let termState : Lean.Elab.Term.SavedState ← Lean.Elab.Term.saveState
-    let tacticState : Lean.Elab.Tactic.SavedState := { term := termState, tactic := { goals := [mvar.mvarId!] }}
-    return tacticState
-/-- Create the initial proof state of the proof tree -/
-def ProofM.start (expr: Lean.Expr): ProofM Unit := do
-  let state: ProofState := {
-    savedState := (← ProofM.runTermElabM' <| start_tactic_state expr),
-    parent := none
-  }
-  let tree ← get
-  set { tree with states := #[state] }
-
-
-def execute_tactic (env: Lean.Environment) (state: Lean.Elab.Tactic.SavedState) (goalId: Nat) (tactic: String) :
-    Lean.Elab.TermElabM (Except (Array String) Lean.Elab.Tactic.SavedState):= do
-  -- Factor this one out to allow for direct syntactic communication
+      Lean.Elab.Tactic.evalTactic stx
+      if (← getThe Lean.Core.State).messages.hasErrors then
+        let messages := (← getThe Lean.Core.State).messages.getErrorMessages |>.toList.toArray
+        let errors ← (messages.map Lean.Message.data).mapM fun md => md.toString
+        return .error errors
+      else
+        return .ok (← Lean.MonadBacktrack.saveState, ← Lean.Elab.Tactic.getUnsolvedGoals)
+    catch exception =>
+      return .error #[← exception.toMessageData.toString]
   match Lean.Parser.runParserCategory
-    (env := env)
+    (env := ← Lean.MonadEnv.getEnv)
     (catName := `tactic)
     (input := tactic)
     (fileName := "<stdin>") with
   | Except.error err => return .error #[err]
-  | Except.ok stx    => do
-    let tac : Lean.Elab.Tactic.TacticM Unit := set state.tactic *> Lean.Elab.Tactic.evalTactic stx
-    match state.tactic.goals.get? goalId with
-    | .none => return .error #[s!"Invalid goalId {goalId}"]
-    | .some mvarId =>
-      state.term.restore
-      try
-        Lean.Elab.Term.synthesizeSyntheticMVarsNoPostponing
-        let unsolvedGoals ← Lean.Elab.Tactic.run mvarId tac
-        if (← getThe Lean.Core.State).messages.hasErrors then
-          let messages := (← getThe Lean.Core.State).messages.getErrorMessages |>.toList.toArray
-          let errors ← (messages.map Lean.Message.data).mapM fun md => md.toString
-          return .error errors
-        else
-          unsolvedGoals.forM Lean.instantiateMVarDeclMVars
-          let nextState : Lean.Elab.Tactic.SavedState := {
-            term := (← Lean.Elab.Term.saveState),
-            tactic := { goals := unsolvedGoals }
-          }
-          return .ok nextState
-      catch ex =>
-        return .error #[← ex.toMessageData.toString]
+  | Except.ok stx    => tacticM stx { elaborator := .anonymous } |>.run' state.tactic
 
-def extract_goals (state: Lean.Elab.Tactic.SavedState): Lean.Elab.TermElabM (Array String) := do
-  state.term.restore
-  let goals ← state.tactic.goals.mapM fun g => do pure $ toString (← Lean.Meta.ppGoal g)
+def goals_to_string (goals: List Lean.MVarId): M (Array String) := do
+  let goals ← goals.mapM fun g => do pure $ toString (← Lean.Meta.ppGoal g)
   pure goals.toArray
 
 
@@ -162,31 +103,29 @@ inductive TacticResult where
   | failure (messages: Array String)
 
 /-- Execute tactic on given state -/
-def ProofM.execute (stateId: Nat) (goalId: Nat) (tactic: String): ProofM TacticResult := do
-  let context ← get
-  match context.states.get? stateId with
+def ProofTree.execute (stateId: Nat) (goalId: Nat) (tactic: String): StateRefT ProofTree M TacticResult := do
+  let tree ← get
+  match tree.states.get? stateId with
   | .none => return .invalid s!"Invalid state id {stateId}"
   | .some state =>
-    match (← ProofM.runTermElabM' <| execute_tactic (env := context.coreState.env) (state := state.savedState) (goalId := goalId) (tactic := tactic)) with
-    | .error errors =>
-      return .failure errors
-    | .ok nextState =>
-      let nextId := context.states.size
-      -- Return goals
-      let goals ← ProofM.runTermElabM' <| extract_goals nextState
-
-      if goals.size = 0 then
-        return .success .none #[]
-      else
-        -- Create next proof state node
-        let proofState: ProofState := {
-          savedState := nextState,
-          parent := stateId,
-          parentGoalId := goalId
-        }
-        modify fun s => { s with states := s.states.push proofState }
-
-        return .success (.some nextId) goals
-
+    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 #[]
+        else
+          let proofState: ProofState := {
+            savedState := nextState,
+            goals := nextGoals,
+            parent := stateId,
+            parentGoalId := goalId
+          }
+          modify fun s => { s with states := s.states.push proofState }
+        return .success (.some nextId) (← goals_to_string nextGoals)
 
 end Pantograph.Meta

Pantograph/Serial.lean

@@ -20,12 +20,4 @@ def syntax_from_str (env: Lean.Environment) (s: String): Except String Lean.Synt
     (input := s)
     (fileName := "<stdin>")
 
-def expr_to_str (env: Lean.Environment) (coreContext: Lean.Core.Context) (expr: Lean.Expr): IO String := do
-  let metaM := Lean.Meta.ppExpr expr
-  let coreM : Lean.CoreM Lean.Format := metaM.run'
-  let coreState : Lean.Core.State := { env := env }
-  let (format, _) ← coreM.toIO coreContext coreState
-  return format.pretty
-
-
 end Pantograph.Serial

README.md

@@ -31,24 +31,23 @@ command { ... }
 The list of available commands can be found in `Pantograph/Commands.lean`. An
 empty command aborts the REPL.
 
+The `Pantograph` executable must be run with a list of modules to import.
+
 Example: (~5k symbols)
 ```
-$ lake env build/bin/Pantograph
-create {"imports": ["Init"]}
-catalog {"envId": 0}
-inspect {"envId": 0, "name": "Nat.le_add_left"}
+$ lake env build/bin/Pantograph "Init"
+catalog
+inspect {"name": "Nat.le_add_left"}
 ```
 Example with `mathlib` (~90k symbols)
 ```
-$ lake env build/bin/Pantograph
-create {"imports": ["Mathlib.Analysis.Seminorm"]}
-catalog {"envId": 0}
+$ lake env build/bin/Pantograph "Mathlib.Analysis.Seminorm"
+catalog
 ```
-Example proving a theorem: (alternatively use `proof.start {"id": 0, "name": "aa", "copyFrom": "Nat.add_comm", "expr": ""}`) to prime the proof
+Example proving a theorem: (alternatively use `proof.start {"copyFrom": "Nat.add_comm"}`) to prime the proof
 ```
-$ lake env build/bin/Pantograph
-create {"imports": ["Init"]}
-proof.start {"envId": 0, "expr": "∀ (n m : Nat), n + m = m + n"}
+$ lake 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}
@@ -63,8 +62,6 @@ If lean encounters stack overflow problems when printing catalog, execute this b
 ```sh
 ulimit -s unlimited
 ```
-Due to a current bug in Lean (which existed in Lean 3), [the default value of structures are not honoured when parsing Json's](https://github.com/leanprover/lean4/issues/2225).
-
 
 ## Testing
 

Test/Proofs.lean

@@ -9,28 +9,21 @@ inductive Start where
   | copy (name: String) -- Start from some name in the environment
   | expr (expr: String) -- Start from some expression
 
-def start_proof (start: Start): IO (LSpec.TestSeq × Option Meta.ProofTree) := do
-  let imports := ["Init"]
-  let env: Lean.Environment ← Lean.importModules
-    (imports := imports.map (λ str => { module := str_to_name str, runtimeOnly := false }))
-    (opts := {})
-    (trustLevel := 1)
-  let state := Meta.createProofTree
-    (name := str_to_name "TestExample") env
-    (coreContext := {
-      currNamespace := str_to_name "Aniva",
-      openDecls := [],     -- No 'open' directives needed
-      fileName := "<Pantograph>",
-      fileMap := { source := "", positions := #[0], lines := #[1] }
-    })
+abbrev M := Meta.M
+abbrev TestM := StateRefT Meta.ProofTree M
+
+def start_proof (start: Start): M (LSpec.TestSeq × Option Meta.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) ← Meta.ProofM.start cinfo.type |>.run state
+    | .some cInfo =>
+      let state ← Meta.ProofTree.create
+        (name := str_to_name "TestExample")
+        (expr := cInfo.type)
       return (testSeq, Option.some state)
     | .none =>
       return (testSeq, Option.none)
@@ -42,21 +35,23 @@ def start_proof (start: Start): IO (LSpec.TestSeq × Option Meta.ProofTree) := d
       IO.println error
       return (testSeq, Option.none)
     | .ok syn =>
-      let expr? := (← Meta.ProofM.syntax_to_expr syn |>.run' state)
+      let expr? ← Meta.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) ← Meta.ProofM.start expr |>.run state
+        let state ← Meta.ProofTree.create
+          (name := str_to_name "TestExample")
+          (expr := expr)
         return (testSeq, Option.some state)
 
 deriving instance DecidableEq, Repr for Meta.TacticResult
 
 def proof_step (stateId: Nat) (goalId: Nat) (tactic: String)
-    (expected: Meta.TacticResult) : Meta.ProofM LSpec.TestSeq := do
-  let result: Meta.TacticResult ← Meta.ProofM.execute stateId goalId tactic
+    (expected: Meta.TacticResult) : TestM LSpec.TestSeq := do
+  let result: Meta.TacticResult ← Meta.ProofTree.execute stateId goalId tactic
   match expected, result with
   | .success (.some i) #[], .success (.some _) goals =>
     -- If the goals are omitted but the next state is specified, we imply that
@@ -65,22 +60,41 @@ def proof_step (stateId: Nat) (goalId: Nat) (tactic: String)
     return LSpec.test s!"{stateId}.{goalId} {tactic}"   (result = expected)
   | _, _ =>
     return LSpec.test s!"{stateId}.{goalId} {tactic}"   (result = expected)
-def proof_inspect (expected: Array String) : Meta.ProofM LSpec.TestSeq := do
-  let result := (← get).structure_array
-  return LSpec.test s!"Tree structure" (result = expected)
 
-def proof_runner (start: Start) (steps: List (Meta.ProofM LSpec.TestSeq)): IO LSpec.TestSeq := 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' state
+def proof_inspect (expected: Array String) : TestM LSpec.TestSeq := do
+  let result := (← get).structure_array
+  return LSpec.test s!"tree structure" (result = expected)
+
+def proof_runner (env: Lean.Environment) (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' 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
+
 
 example: ∀ (a b: Nat), a + b = b + a := by
   intro n m
   rw [Nat.add_comm]
-def proof_nat_add_comm: IO LSpec.TestSeq :=
+def proof_nat_add_comm (env: Lean.Environment): IO LSpec.TestSeq :=
   let goal1 := "n m : Nat\n⊢ n + m = m + n"
-  proof_runner (.copy "Nat.add_comm") [
+  proof_runner env (.copy "Nat.add_comm") [
     proof_step 0 0 "intro n m"
       (.success (.some 1) #[goal1]),
     proof_step 1 0 "assumption"
@@ -88,9 +102,9 @@ def proof_nat_add_comm: IO LSpec.TestSeq :=
     proof_step 1 0 "rw [Nat.add_comm]"
       (.success .none #[])
   ]
-def proof_nat_add_comm_manual: IO LSpec.TestSeq := do
+def proof_nat_add_comm_manual (env: Lean.Environment): IO LSpec.TestSeq := do
   let goal1 := "n m : Nat\n⊢ n + m = m + n"
-  proof_runner (.expr "∀ (a b: Nat), a + b = b + a") [
+  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"
@@ -114,15 +128,15 @@ example: ∀ (p q: Prop), p ∨ q → q ∨ p := by
     assumption
   . apply Or.inl
     assumption
-def proof_or_comm: IO LSpec.TestSeq := do
-  proof_runner (.expr "∀ (p q: Prop), p ∨ q → q ∨ p") [
+def proof_or_comm (env: Lean.Environment): IO LSpec.TestSeq := do
+  proof_runner env (.expr "∀ (p q: Prop), p ∨ q → q ∨ p") [
     proof_step 0 0 "intro p q h"
       (.success (.some 1) #["p q : Prop\nh : p ∨ q\n⊢ q ∨ p"]),
     proof_step 1 0 "cases h"
       (.success (.some 2) #[]),
     proof_inspect #["", "0.0", "1.0"],
     proof_step 2 0 "apply Or.inr"
-      (.success (.some 3) #[""]),
+      (.success (.some 3) #[]),
     proof_inspect #["", "0.0", "1.0", "2.0"],
     proof_step 3 0 "assumption"
       (.success .none #[]),
@@ -130,14 +144,34 @@ def proof_or_comm: IO LSpec.TestSeq := do
       (.success (.some 4) #[]),
     proof_step 4 0 "assumption"
       (.success .none #[]),
-    proof_inspect #["", "0.0", "1.0", "2.0", "1.1"]
+    proof_inspect #["", "0.0", "1.0", "2.0", "2.1"]
+  ]
+
+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 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)
+
   return LSpec.group "Proofs" $
-    (LSpec.group "Nat.add_comm" $ (← proof_nat_add_comm)) ++
-    (LSpec.group "Nat.add_comm manual" $ (← proof_nat_add_comm_manual)) ++
-    (LSpec.group "Or.comm" $ (← proof_or_comm))
+    (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))
 
 end Pantograph.Test