fix: snapshot subtree was not restored on reuse (leanprover#4643)

This could lead to nested error messages and info trees vanishing on
edits (strictly) below them

Fixes the second issue in leanprover#4623

src/Lean/Elab/MutualDef.lean

@@ -167,7 +167,7 @@ private def elabHeaders (views : Array DefView)
         else
           reuseBody := false
 
-      let mut (newHeader, newState) ← withRestoreOrSaveFull reusableResult? do
+      let mut (newHeader, newState) ← withRestoreOrSaveFull reusableResult? none do
         withRef view.headerRef do
         addDeclarationRanges declName view.ref  -- NOTE: this should be the full `ref`
         applyAttributesAt declName view.modifiers.attrs .beforeElaboration
@@ -337,14 +337,9 @@ private def elabFunValues (headers : Array DefViewElabHeader) : TermElabM (Array
         -- elaboration
         if let some old := old.val.get then
           snap.new.resolve <| some old
-          -- also make sure to reuse tactic snapshots if present so that body reuse does not lead to
-          -- missed tactic reuse on further changes
-          if let some tacSnap := header.tacSnap? then
-            if let some oldTacSnap := tacSnap.old? then
-              tacSnap.new.resolve oldTacSnap.val.get
           reusableResult? := some (old.value, old.state)
 
-    let (val, state) ← withRestoreOrSaveFull reusableResult? do
+    let (val, state) ← withRestoreOrSaveFull reusableResult? header.tacSnap? do
       withDeclName header.declName <| withLevelNames header.levelNames do
       let valStx ← liftMacroM <| declValToTerm header.value
       forallBoundedTelescope header.type header.numParams fun xs type => do

src/Lean/Elab/Tactic/Basic.lean

@@ -96,14 +96,15 @@ def SavedState.restore (b : SavedState) (restoreInfo := false) : TacticM Unit :=
   b.term.restore restoreInfo
   set b.tactic
 
-@[specialize, inherit_doc Core.withRestoreOrSaveFull]
-def withRestoreOrSaveFull (reusableResult? : Option (α × SavedState)) (act : TacticM α) :
+@[specialize, inherit_doc Term.withRestoreOrSaveFull]
+def withRestoreOrSaveFull (reusableResult? : Option (α × SavedState))
+    (tacSnap? : Option (Language.SnapshotBundle Tactic.TacticParsedSnapshot)) (act : TacticM α) :
     TacticM (α × SavedState) := do
   if let some (_, state) := reusableResult? then
     set state.tactic
   let reusableResult? := reusableResult?.map (fun (val, state) => (val, state.term))
   let (a, term) ← controlAt TermElabM fun runInBase => do
-    Term.withRestoreOrSaveFull reusableResult? <| runInBase act
+    Term.withRestoreOrSaveFull reusableResult? tacSnap? <| runInBase act
   return (a, { term, tactic := (← get) })
 
 protected def getCurrMacroScope : TacticM MacroScope := do pure (← readThe Core.Context).currMacroScope

src/Lean/Elab/Tactic/BuiltinTactic.lean

@@ -90,11 +90,10 @@ where
               {
                 range? := stxs |>.getRange?
                 task := next.result }]
-            let (_, state) ← withRestoreOrSaveFull reusableResult? do
-              -- set up nested reuse; `evalTactic` will check for `isIncrementalElab`
-              withTheReader Term.Context ({ · with
-                  tacSnap? := some { old? := oldInner?, new := inner } }) do
-                evalTactic tac
+            let (_, state) ← withRestoreOrSaveFull reusableResult?
+                -- set up nested reuse; `evalTactic` will check for `isIncrementalElab`
+                (tacSnap? := some { old? := oldInner?, new := inner }) do
+              evalTactic tac
             finished.resolve { state? := state }
 
         withTheReader Term.Context ({ · with tacSnap? := some {

src/Lean/Elab/Term.lean

@@ -327,14 +327,33 @@ def SavedState.restore (s : SavedState) (restoreInfo : Bool := false) : TermElab
   unless restoreInfo do
     setInfoState infoState
 
-@[specialize, inherit_doc Core.withRestoreOrSaveFull]
-def withRestoreOrSaveFull (reusableResult? : Option (α × SavedState)) (act : TermElabM α) :
+/--
+Like `Meta.withRestoreOrSaveFull` for `TermElabM`, but also takes a `tacSnap?` that
+* when running `act`, is set as `Context.tacSnap?`
+* otherwise (i.e. on restore) is used to update the new snapshot promise to the old task's
+  value.
+This extra restore step is necessary because while `reusableResult?` can be used to replay any
+effects on `State`, `Context.tacSnap?` is not part of it but changed via an `IO` side effect, so
+it needs to be replayed separately.
+
+We use an explicit parameter instead of accessing `Context.tacSnap?` directly because this prevents
+`withRestoreOrSaveFull` and `withReader` from being used in the wrong order.
+-/
+@[specialize]
+def withRestoreOrSaveFull (reusableResult? : Option (α × SavedState))
+    (tacSnap? : Option (Language.SnapshotBundle Tactic.TacticParsedSnapshot)) (act : TermElabM α) :
     TermElabM (α × SavedState) := do
   if let some (_, state) := reusableResult? then
     set state.elab
+    if let some snap := tacSnap? then
+      let some old := snap.old?
+        | throwError "withRestoreOrSaveFull: expected old snapshot in `tacSnap?`"
+      snap.new.resolve old.val.get
+
   let reusableResult? := reusableResult?.map (fun (val, state) => (val, state.meta))
-  let (a, meta) ← controlAt MetaM fun runInBase => do
-    Meta.withRestoreOrSaveFull reusableResult? <| runInBase act
+  let (a, meta) ← withReader ({ · with tacSnap? }) do
+    controlAt MetaM fun runInBase => do
+      Meta.withRestoreOrSaveFull reusableResult? <| runInBase act
   return (a, { meta, «elab» := (← get) })
 
 instance : MonadBacktrack SavedState TermElabM where

tests/lean/interactive/incrementalTactic.lean

@@ -92,3 +92,16 @@ def tacInTermInTac2 : True := by
   | zero => done
   | succ => sorry
 --^ collectDiagnostics
+
+/-!
+Error messages (or rather, whole snapshot subtrees) may vanish if we do not properly restore them on
+reuse (which `Term.withRestoreOrSaveFull` now makes sure of).
+-/
+-- RESET
+example : True := by
+  · · done
+    · done
+  · skip
+      --^ sync
+      --^ insert: " "
+      --^ collectDiagnostics

tests/lean/interactive/incrementalTactic.lean.expected.out

@@ -85,3 +85,27 @@ s
    "fullRange":
    {"start": {"line": 3, "character": 12},
     "end": {"line": 3, "character": 16}}}]}
+{"version": 2,
+ "uri": "file:///incrementalTactic.lean",
+ "diagnostics":
+ [{"source": "Lean 4",
+   "severity": 1,
+   "range":
+   {"start": {"line": 2, "character": 6}, "end": {"line": 2, "character": 10}},
+   "message": "unsolved goals\n⊢ True",
+   "fullRange":
+   {"start": {"line": 2, "character": 6}, "end": {"line": 2, "character": 10}}},
+  {"source": "Lean 4",
+   "severity": 1,
+   "range":
+   {"start": {"line": 3, "character": 4}, "end": {"line": 3, "character": 10}},
+   "message": "no goals to be solved",
+   "fullRange":
+   {"start": {"line": 3, "character": 4}, "end": {"line": 3, "character": 10}}},
+  {"source": "Lean 4",
+   "severity": 1,
+   "range":
+   {"start": {"line": 4, "character": 2}, "end": {"line": 4, "character": 8}},
+   "message": "no goals to be solved",
+   "fullRange":
+   {"start": {"line": 4, "character": 2}, "end": {"line": 4, "character": 8}}}]}