Save core state in proofs

Main.lean

@@ -200,10 +200,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 parents := tree.states.map λ state => match state.parent with
-        | .none => ""
-        | .some parent => s!"{parent}.{state.parentGoalId}"
-      return Lean.toJson ({parents := parents}: ProofPrintTreeResult)
+      return Lean.toJson ({parents := tree.structure_array}: ProofPrintTreeResult)
 
 
 end Pantograph

Pantograph/Meta.lean

@@ -27,11 +27,22 @@ structure ProofState where
   parentGoalId : Nat  := 0
 structure ProofTree where
   -- All parameters needed to run a `TermElabM` monad
-  env: Lean.Environment
   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 := #[]
 
@@ -39,24 +50,34 @@ abbrev ProofM := StateRefT ProofTree IO
 
 def createProofTree (name: Lean.Name) (env: Lean.Environment) (coreContext: Lean.Core.Context): ProofTree :=
   {
-    env := env,
     name := name,
     coreContext := coreContext,
     elabContext := {
       declName? := some (name ++ "_pantograph"),
       errToSorry := false
     }
+    coreState := {
+      env := env
     }
+  }
+
+-- Tree structures
+
+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'
-  let coreState : Lean.Core.State := { env := context.env }
-  coreM.toIO context.coreContext coreState
+  coreM.toIO context.coreContext context.coreState
 def ProofM.runTermElabM' (termElabM: Lean.Elab.TermElabM α): ProofM α := do
-  return Prod.fst <| ← ProofM.runTermElabM termElabM
+  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
@@ -103,11 +124,11 @@ def execute_tactic (env: Lean.Environment) (state: Lean.Elab.Tactic.SavedState)
     (fileName := "<stdin>") with
   | Except.error err => return .error #[err]
   | Except.ok stx    => do
-    state.term.restore
     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
@@ -146,7 +167,7 @@ def ProofM.execute (stateId: Nat) (goalId: Nat) (tactic: String): ProofM TacticR
   match context.states.get? stateId with
   | .none => return .invalid s!"Invalid state id {stateId}"
   | .some state =>
-    match (← ProofM.runTermElabM' <| execute_tactic (env := context.env) (state := state.savedState) (goalId := goalId) (tactic := tactic)) with
+    match (← ProofM.runTermElabM' <| execute_tactic (env := context.coreState.env) (state := state.savedState) (goalId := goalId) (tactic := tactic)) with
     | .error errors =>
       return .failure errors
     | .ok nextState =>

Test/Proofs.lean

@@ -65,6 +65,9 @@ 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
@@ -96,6 +99,7 @@ def proof_nat_add_comm_manual: IO LSpec.TestSeq := do
       (.success .none #[])
   ]
 
+-- Two ways to write the same theorem
 example: ∀ (p q: Prop), p ∨ q → q ∨ p := by
   intro p q h
   cases h
@@ -103,20 +107,30 @@ example: ∀ (p q: Prop), p ∨ q → q ∨ p := by
   assumption
   apply Or.inl
   assumption
+example: ∀ (p q: Prop), p ∨ q → q ∨ p := by
+  intro p q h
+  cases h
+  . apply Or.inr
+    assumption
+  . apply Or.inl
+    assumption
 def proof_or_comm: IO LSpec.TestSeq := do
   proof_runner (.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 #[]),
     proof_step 2 1 "apply Or.inl"
       (.success (.some 4) #[]),
     proof_step 4 0 "assumption"
-      (.success .none #[])
+      (.success .none #[]),
+    proof_inspect #["", "0.0", "1.0", "2.0", "1.1"]
   ]
 
 def test_proofs : IO LSpec.TestSeq := do