fix: Consolidate TermElabM blocks

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?,

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

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