From 1fb189a38fd78e27e0a069db5829e3b6d8b22b50 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Fri, 8 Dec 2023 17:31:25 -0800
Subject: [PATCH] fix: Consolidate TermElabM blocks

---
 Pantograph.lean        | 12 ++++++------
 Pantograph/Goal.lean   |  6 ++++++
 Pantograph/Serial.lean |  4 ++--
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/Pantograph.lean b/Pantograph.lean
index ab6a7b1..a5a181b 100644
--- a/Pantograph.lean
+++ b/Pantograph.lean
@@ -134,24 +134,23 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
   goal_start (args: Protocol.GoalStart): MainM (CR Protocol.GoalStartResult) := do
     let state ← get
     let env ← Lean.MonadEnv.getEnv
-    let expr?: Except _ Lean.Expr ← (match args.expr, args.copyFrom with
+    let expr?: Except _ GoalState ← runTermElabM (match args.expr, args.copyFrom with
       | .some expr, .none =>
         (match syntax_from_str env expr with
         | .error str => return .error <| errorI "parsing" str
         | .ok syn => do
-          (match ← runTermElabM <| syntax_to_expr syn with
+          (match ← syntax_to_expr syn with
           | .error str => return .error <| errorI "elab" str
-          | .ok expr => return .ok expr))
+          | .ok expr => return .ok (← GoalState.create expr)))
       | .none, .some copyFrom =>
         (match env.find? <| copyFrom.toName with
         | .none => return .error <| errorIndex s!"Symbol not found: {copyFrom}"
-        | .some cInfo => return .ok cInfo.type)
+        | .some cInfo => return .ok (← GoalState.create cInfo.type))
       | _, _ =>
         return .error <| errorI "arguments" "Exactly one of {expr, copyFrom} must be supplied")
     match expr? with
     | .error error => return .error error
-    | .ok expr =>
-      let goalState ← runTermElabM <| GoalState.create expr
+    | .ok goalState =>
       let stateId := state.nextId
       set { state with
         goalStates := state.goalStates.insert stateId goalState,
@@ -225,6 +224,7 @@ def execute (command: Protocol.Command): MainM Lean.Json := do
     match state.goalStates.find? args.stateId with
     | .none => return .error $ errorIndex s!"Invalid state index {args.stateId}"
     | .some goalState => do
+      goalState.restoreMetaM
       let root? ← goalState.rootExpr?.mapM (λ expr => serialize_expression state.options expr)
       return .ok {
         root?,
diff --git a/Pantograph/Goal.lean b/Pantograph/Goal.lean
index 3645087..1589408 100644
--- a/Pantograph/Goal.lean
+++ b/Pantograph/Goal.lean
@@ -51,6 +51,10 @@ protected def GoalState.env (state: GoalState): Environment :=
   state.savedState.term.meta.core.env
 private def GoalState.mvars (state: GoalState): SSet MVarId :=
   state.mctx.decls.foldl (init := .empty) fun acc k _ => acc.insert k
+private def GoalState.restoreElabM (state: GoalState): Elab.TermElabM Unit :=
+  state.savedState.term.restore
+def GoalState.restoreMetaM (state: GoalState): MetaM Unit :=
+  state.savedState.term.meta.restore
 
 /-- Inner function for executing tactic on goal state -/
 def executeTactic (state: Elab.Tactic.SavedState) (goal: MVarId) (tactic: Syntax) :
@@ -84,6 +88,7 @@ inductive TacticResult where
 /-- Execute tactic on given state -/
 protected def GoalState.execute (state: GoalState) (goalId: Nat) (tactic: String):
     M TacticResult := do
+  state.restoreElabM
   let goal ← match state.savedState.tactic.goals.get? goalId with
     | .some goal => pure $ goal
     | .none => return .indexError goalId
@@ -118,6 +123,7 @@ protected def GoalState.execute (state: GoalState) (goalId: Nat) (tactic: String
     }
 
 protected def GoalState.tryAssign (state: GoalState) (goalId: Nat) (expr: String): M TacticResult := do
+  state.restoreElabM
   let goal ← match state.savedState.tactic.goals.get? goalId with
     | .some goal => pure goal
     | .none => return .indexError goalId
diff --git a/Pantograph/Serial.lean b/Pantograph/Serial.lean
index 42fc4e2..3d0d945 100644
--- a/Pantograph/Serial.lean
+++ b/Pantograph/Serial.lean
@@ -242,8 +242,8 @@ def serialize_goal (options: Protocol.Options) (goal: MVarId) (mvarDecl: Metavar
   of_name (n: Name) := name_to_ast n (sanitize := false)
 
 protected def GoalState.serializeGoals (state: GoalState) (parent: Option GoalState := .none) (options: Protocol.Options := {}): MetaM (Array Protocol.Goal):= do
+  state.restoreMetaM
   let goals := state.goals.toArray
-  state.savedState.term.meta.restore
   let parentDecl? := parent.bind (λ parentState =>
     let parentGoal := parentState.goals.get! state.parentGoalId
     parentState.mctx.findDecl? parentGoal)
@@ -256,8 +256,8 @@ protected def GoalState.serializeGoals (state: GoalState) (parent: Option GoalSt
 
 /-- Print the metavariables in a readable format -/
 protected def GoalState.print (goalState: GoalState) (options: Protocol.GoalDiag := {}): MetaM Unit := do
+  goalState.restoreMetaM
   let savedState := goalState.savedState
-  savedState.term.meta.restore
   let goals := savedState.tactic.goals
   let mctx ← getMCtx
   let root := goalState.root