fix: incorrect info tree reuse (leanprover#4340)

The `save` happened in a slightly different context from the restore,
which a refinement of the `saveOrRestoreFull` signature now makes
impossible.

Fixes leanprover#4328

src/Lean/CoreM.lean

@@ -219,31 +219,31 @@ def saveState : CoreM SavedState := do
   return { toState := s, passedHearbeats := 0 }
 
 /--
-Incremental reuse primitive: if `reusableResult?` is `none`, runs `cont` with an action `save` that
-on execution returns the saved monadic state at this point including the heartbeats used by `cont`
-so far. If `reusableResult?` on the other hand is `some (a, state)`, restores full `state` including
-heartbeats used and returns `a`.
+Incremental reuse primitive: if `reusableResult?` is `none`, runs `act` and returns its result
+together with the saved monadic state after `act` including the heartbeats used by it. If
+`reusableResult?` on the other hand is `some (a, state)`, restores full `state` including heartbeats
+used and returns `(a, state)`.
 
 The intention is for steps that support incremental reuse to initially pass `none` as
-`reusableResult?` and call `save` as late as possible in `cont`. In a further run, if reuse is
-possible, `reusableResult?` should be set to the previous state and result, ensuring that the state
+`reusableResult?` and store the result and state in a snapshot. In a further run, if reuse is
+possible, `reusableResult?` should be set to the previous result and state, ensuring that the state
 after running `withRestoreOrSaveFull` is identical in both runs. Note however that necessarily this
-is only an approximation in the case of heartbeats as heartbeats used by `withRestoreOrSaveFull`, by
-the remainder of `cont` after calling `save`, as well as by reuse-handling code such as the one
-supplying `reusableResult?` are not accounted for.
+is only an approximation in the case of heartbeats as heartbeats used by `withRestoreOrSaveFull`
+itself after calling `act` as well as by reuse-handling code such as the one supplying
+`reusableResult?` are not accounted for.
 -/
 @[specialize] def withRestoreOrSaveFull (reusableResult? : Option (α × SavedState))
-    (cont : (save : CoreM SavedState) → CoreM α) : CoreM α := do
+    (act : CoreM α) : CoreM (α × SavedState) := do
   if let some (val, state) := reusableResult? then
     set state.toState
     IO.addHeartbeats state.passedHearbeats.toUInt64
-    return val
+    return (val, state)
 
   let startHeartbeats ← IO.getNumHeartbeats
-  cont (do
-    let s ← get
-    let stopHeartbeats ← IO.getNumHeartbeats
-    return { toState := s, passedHearbeats := stopHeartbeats - startHeartbeats })
+  let a ← act
+  let s ← get
+  let stopHeartbeats ← IO.getNumHeartbeats
+  return (a, { toState := s, passedHearbeats := stopHeartbeats - startHeartbeats })
 
 /-- Restore backtrackable parts of the state. -/
 def SavedState.restore (b : SavedState) : CoreM Unit :=

src/Lean/Elab/MutualDef.lean

@@ -141,17 +141,12 @@ private def elabHeaders (views : Array DefView)
     for view in views, ⟨shortDeclName, declName, levelNames⟩ in expandedDeclIds,
         tacPromise in tacPromises, bodyPromise in bodyPromises do
       let mut reusableResult? := none
+      let mut oldBodySnap? := none
+      let mut oldTacSnap? := none
       if let some snap := view.headerSnap? then
         -- by the `DefView.headerSnap?` invariant, safe to reuse results at this point, so let's
         -- wait for them!
         if let some old := snap.old?.bind (·.val.get) then
-          let (tacStx?, newTacTask?) ← mkTacTask view.value tacPromise
-          snap.new.resolve <| some { old with
-            tacStx?
-            tacSnap? := newTacTask?
-            bodyStx := view.value
-            bodySnap := mkBodyTask view.value bodyPromise
-          }
           -- Transition from `DefView.snap?` to `DefViewElabHeader.tacSnap?` invariant: if all
           -- headers and all previous bodies could be reused, then the state at the *start* of the
           -- top-level tactic block (if any) is unchanged
@@ -161,26 +156,18 @@ private def elabHeaders (views : Array DefView)
           -- we can reuse the result
           reuseBody := reuseBody &&
             view.value.structRangeEqWithTraceReuse (← getOptions) old.bodyStx
-          let header := { old.view, view with
-            -- We should only forward the promise if we are actually waiting on the corresponding
-            -- task; otherwise, diagnostics assigned to it will be lost
-            tacSnap? := guard newTacTask?.isSome *> some {
-              old? := do
-                guard reuseTac
-                some ⟨(← old.tacStx?), (← old.tacSnap?)⟩
-              new := tacPromise
-            }
-            bodySnap? := some {
-              -- no syntax guard to store, we already did the necessary checks
-              old? := guard reuseBody *> pure ⟨.missing, old.bodySnap⟩
-              new := bodyPromise
-            }
-          }
-          reusableResult? := some (header, old.state)
+          -- no syntax guard to store, we already did the necessary checks
+          oldBodySnap? := guard reuseBody *> pure ⟨.missing, old.bodySnap⟩
+          oldTacSnap? := do
+              guard reuseTac
+              some ⟨(← old.tacStx?), (← old.tacSnap?)⟩
+          let newHeader : DefViewElabHeader := { view, old.view with
+            bodySnap? := none, tacSnap? := none }  -- filled below
+          reusableResult? := some (newHeader, old.state)
         else
           reuseBody := false
 
-      let header ← withRestoreOrSaveFull reusableResult? fun save => do
+      let mut (newHeader, newState) ← withRestoreOrSaveFull reusableResult? do
         withRef view.headerRef do
         addDeclarationRanges declName view.ref  -- NOTE: this should be the full `ref`
         applyAttributesAt declName view.modifiers.attrs .beforeElaboration
@@ -219,29 +206,29 @@ private def elabHeaders (views : Array DefView)
               declName, shortDeclName, type, levelNames, binderIds
               numParams := xs.size
             }
-            let mut newHeader : DefViewElabHeader := { view, newHeader with
+            let newHeader : DefViewElabHeader := { view, newHeader with
               bodySnap? := none, tacSnap? := none }
-            if let some snap := view.headerSnap? then
-              let (tacStx?, newTacTask?) ← mkTacTask view.value tacPromise
-              snap.new.resolve <| some {
-                diagnostics :=
-                  (← Language.Snapshot.Diagnostics.ofMessageLog (← Core.getAndEmptyMessageLog))
-                view := newHeader.toDefViewElabHeaderData
-                state := (← save)
-                tacStx?
-                tacSnap? := newTacTask?
-                bodyStx := view.value
-                bodySnap := mkBodyTask view.value bodyPromise
-              }
-              newHeader := { newHeader with
-                -- We should only forward the promise if we are actually waiting on the
-                -- corresponding task; otherwise, diagnostics assigned to it will be lost
-                tacSnap? := guard newTacTask?.isSome *> some { old? := none, new := tacPromise }
-                bodySnap? := some { old? := none, new := bodyPromise }
-              }
             check headers newHeader
             return newHeader
-      headers := headers.push header
+      if let some snap := view.headerSnap? then
+        let (tacStx?, newTacTask?) ← mkTacTask view.value tacPromise
+        snap.new.resolve <| some {
+          diagnostics :=
+            (← Language.Snapshot.Diagnostics.ofMessageLog (← Core.getAndEmptyMessageLog))
+          view := newHeader.toDefViewElabHeaderData
+          state := newState
+          tacStx?
+          tacSnap? := newTacTask?
+          bodyStx := view.value
+          bodySnap := mkBodyTask view.value bodyPromise
+        }
+        newHeader := { newHeader with
+          -- We should only forward the promise if we are actually waiting on the
+          -- corresponding task; otherwise, diagnostics assigned to it will be lost
+          tacSnap? := guard newTacTask?.isSome *> some { old? := oldTacSnap?, new := tacPromise }
+          bodySnap? := some { old? := oldBodySnap?, new := bodyPromise }
+        }
+      headers := headers.push newHeader
     return headers
 where
   getBodyTerm? (stx : Syntax) : Option Syntax :=
@@ -357,7 +344,7 @@ private def elabFunValues (headers : Array DefViewElabHeader) : TermElabM (Array
               tacSnap.new.resolve oldTacSnap.val.get
           reusableResult? := some (old.value, old.state)
 
-    withRestoreOrSaveFull reusableResult? fun save => do
+    let (val, state) ← withRestoreOrSaveFull reusableResult? do
       withDeclName header.declName <| withLevelNames header.levelNames do
       let valStx ← liftMacroM <| declValToTerm header.value
       forallBoundedTelescope header.type header.numParams fun xs type => do
@@ -371,15 +358,15 @@ private def elabFunValues (headers : Array DefViewElabHeader) : TermElabM (Array
         -- NOTE: without this `instantiatedMVars`, `mkLambdaFVars` may leave around a redex that
         -- leads to more section variables being included than necessary
         let val ← instantiateMVars val
-        let val ← mkLambdaFVars xs val
-        if let some snap := header.bodySnap? then
-          snap.new.resolve <| some {
-            diagnostics :=
-              (← Language.Snapshot.Diagnostics.ofMessageLog (← Core.getAndEmptyMessageLog))
-            state := (← save)
-            value := val
-          }
-        return val
+        mkLambdaFVars xs val
+    if let some snap := header.bodySnap? then
+      snap.new.resolve <| some {
+        diagnostics :=
+          (← Language.Snapshot.Diagnostics.ofMessageLog (← Core.getAndEmptyMessageLog))
+        state
+        value := val
+      }
+    return val
 
 private def collectUsed (headers : Array DefViewElabHeader) (values : Array Expr) (toLift : List LetRecToLift)
     : StateRefT CollectFVars.State MetaM Unit := do

src/Lean/Elab/Tactic/Basic.lean

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

src/Lean/Elab/Tactic/BuiltinTactic.lean

@@ -71,7 +71,7 @@ where
         -- `tac` must be unchanged given the narrow above; let's reuse `finished`'s state!
         let oldParsed := old.val.get
         if let some state := oldParsed.data.finished.get.state? then
-          reusableResult? := some (state, state)
+          reusableResult? := some ((), state)
           -- only allow `next` reuse in this case
           oldNext? := oldParsed.next.get? 1 |>.map (⟨old.stx, ·⟩)
 
@@ -90,12 +90,11 @@ where
               {
                 range? := stxs |>.getRange?
                 task := next.result }]
-            let state ← withRestoreOrSaveFull reusableResult? fun save => do
+            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
-              save
             finished.resolve { state? := state }
 
         withTheReader Term.Context ({ · with tacSnap? := some {

src/Lean/Elab/Term.lean

@@ -306,14 +306,14 @@ def SavedState.restore (s : SavedState) (restoreInfo : Bool := false) : TermElab
     setInfoState infoState
 
 @[specialize, inherit_doc Core.withRestoreOrSaveFull]
-def withRestoreOrSaveFull (reusableResult? : Option (α × SavedState))
-    (cont : TermElabM SavedState → TermElabM α) : TermElabM α := do
+def withRestoreOrSaveFull (reusableResult? : Option (α × SavedState)) (act : TermElabM α) :
+    TermElabM (α × SavedState) := do
   if let some (_, state) := reusableResult? then
     set state.elab
   let reusableResult? := reusableResult?.map (fun (val, state) => (val, state.meta))
-  controlAt MetaM fun runInBase =>
-    Meta.withRestoreOrSaveFull reusableResult? fun restore =>
-      runInBase <| cont (return { meta := (← restore), «elab» := (← get) })
+  let (a, meta) ← controlAt MetaM fun runInBase => do
+    Meta.withRestoreOrSaveFull reusableResult? <| runInBase act
+  return (a, { meta, «elab» := (← get) })
 
 instance : MonadBacktrack SavedState TermElabM where
   saveState      := Term.saveState

src/Lean/Meta/Basic.lean

@@ -418,14 +418,14 @@ def SavedState.restore (b : SavedState) : MetaM Unit := do
   modify fun s => { s with mctx := b.meta.mctx, zetaDeltaFVarIds := b.meta.zetaDeltaFVarIds, postponed := b.meta.postponed }
 
 @[specialize, inherit_doc Core.withRestoreOrSaveFull]
-def withRestoreOrSaveFull (reusableResult? : Option (α × SavedState))
-    (cont : MetaM SavedState → MetaM α) : MetaM α := do
+def withRestoreOrSaveFull (reusableResult? : Option (α × SavedState)) (act : MetaM α) :
+    MetaM (α × SavedState) := do
   if let some (_, state) := reusableResult? then
     set state.meta
   let reusableResult? := reusableResult?.map (fun (val, state) => (val, state.core))
-  controlAt CoreM fun runInCoreM =>
-    Core.withRestoreOrSaveFull reusableResult? fun restore =>
-      runInCoreM <| cont (return { core := (← restore), meta := (← get) })
+  let (a, core) ← controlAt CoreM fun runInBase => do
+    Core.withRestoreOrSaveFull reusableResult? <| runInBase act
+  return (a, { core, meta := (← get) })
 
 instance : MonadBacktrack SavedState MetaM where
   saveState      := Meta.saveState

tests/lean/interactive/incrementalMutual.lean

@@ -29,3 +29,12 @@ def h1 : (by dbg_trace "h 1 0"; exact Nat) := (by dbg_trace "h 1 1"; exact 0)
                            --^ sync
                            --^ insert: ".5"
 end
+
+/-! #4328 incorrect info tree restore led to linter false-positives. -/
+
+-- RESET
+def map' {α β} (f : α → β) : List α → List β :=
+  List.map f
+          --^ collectDiagnostics
+          --^ insert: "\n"
+          --^ collectDiagnostics

tests/lean/interactive/incrementalMutual.lean.expected.out

@@ -17,3 +17,5 @@ h 1 1
 h 1 0.5
 h 0 1
 h 1 1
+{"version": 1, "uri": "file:///incrementalMutual.lean", "diagnostics": []}
+{"version": 2, "uri": "file:///incrementalMutual.lean", "diagnostics": []}