fix: cancellation of synchronous part of previous elaboration (#7882)

This PR fixes a regression where elaboration of a previous document
version is not cancelled on changes to the document.

Done by removing the default from `SnapshotTask.cancelTk?` and
consistently passing the current thread's token for synchronous
elaboration steps.

(cherry picked from commit 1421b61)

Author: Kha

src/Lean/Elab/Command.lean

@@ -510,14 +510,15 @@ where go := do
                 let oldCmds? := oldSnap?.map fun old =>
                   if old.newStx.isOfKind nullKind then old.newStx.getArgs else #[old.newStx]
                 let cmdPromises ← cmds.mapM fun _ => IO.Promise.new
+                let cancelTk? := (← read).cancelTk?
                 snap.new.resolve <| .ofTyped {
                   diagnostics := .empty
                   macroDecl := decl
                   newStx := stxNew
                   newNextMacroScope := nextMacroScope
                   hasTraces
                   next := Array.zipWith (fun cmdPromise cmd =>
-                    { stx? := some cmd, task := cmdPromise.resultD default }) cmdPromises cmds
+                    { stx? := some cmd, task := cmdPromise.resultD default, cancelTk? }) cmdPromises cmds
                   : MacroExpandedSnapshot
                 }
                 -- After the first command whose syntax tree changed, we must disable

src/Lean/Elab/MutualDef.lean

@@ -165,6 +165,8 @@ private def elabHeaders (views : Array DefView) (expandedDeclIds : Array ExpandD
           -- no syntax guard to store, we already did the necessary checks
           oldBodySnap? := guard reuseBody *> pure ⟨.missing, old.bodySnap⟩
           if oldBodySnap?.isNone then
+            -- NOTE: this will eagerly cancel async tasks not associated with an inner snapshot, most
+            -- importantly kernel checking and compilation of the top-level declaration
             old.bodySnap.cancelRec
           oldTacSnap? := do
               guard reuseTac
@@ -217,6 +219,7 @@ private def elabHeaders (views : Array DefView) (expandedDeclIds : Array ExpandD
             return newHeader
       if let some snap := view.headerSnap? then
         let (tacStx?, newTacTask?) ← mkTacTask view.value tacPromise
+        let cancelTk? := (← readThe Core.Context).cancelTk?
         let bodySnap := {
           stx? := view.value
           reportingRange? :=
@@ -227,6 +230,8 @@ private def elabHeaders (views : Array DefView) (expandedDeclIds : Array ExpandD
             else
               getBodyTerm? view.value |>.getD view.value |>.getRange?
           task := bodyPromise.resultD default
+          -- We should not cancel the entire body early if we have tactics
+          cancelTk? := guard newTacTask?.isNone *> cancelTk?
         }
         snap.new.resolve <| some {
           diagnostics :=
@@ -269,7 +274,8 @@ where
    := do
     if let some e := getBodyTerm? body then
       if let `(by $tacs*) := e then
-        return (e, some { stx? := mkNullNode tacs, task := tacPromise.resultD default })
+        let cancelTk? := (← readThe Core.Context).cancelTk?
+        return (e, some { stx? := mkNullNode tacs, task := tacPromise.resultD default, cancelTk? })
     tacPromise.resolve default
     return (none, none)
 
@@ -432,8 +438,7 @@ private def elabFunValues (headers : Array DefViewElabHeader) (vars : Array Expr
           snap.new.resolve <| some old
           reusableResult? := some (old.value, old.state)
         else
-          -- NOTE: this will eagerly cancel async tasks not associated with an inner snapshot, most
-          -- importantly kernel checking and compilation of the top-level declaration
+          -- make sure to cancel any async tasks that may still be running (e.g. kernel and codegen)
           old.val.cancelRec
 
     let (val, state) ← withRestoreOrSaveFull reusableResult? header.tacSnap? do
@@ -1197,6 +1202,7 @@ private def logGoalsAccomplishedSnapshotTask (views : Array DefView)
     -- Use first line of the mutual block to avoid covering the progress of the whole mutual block
     reportingRange? := (← getRef).getPos?.map fun pos => ⟨pos, pos⟩
     task := logGoalsAccomplishedTask
+    cancelTk? := none
   }
 
 end Term
@@ -1235,9 +1241,10 @@ def elabMutualDef (ds : Array Syntax) : CommandElabM Unit := do
       } }
       if snap.old?.isSome && (view.headerSnap?.bind (·.old?)).isNone then
         snap.old?.forM (·.val.cancelRec)
+      let cancelTk? := (← read).cancelTk?
       defs := defs.push {
         fullHeaderRef
-        headerProcessedSnap := { stx? := d, task := headerPromise.resultD default }
+        headerProcessedSnap := { stx? := d, task := headerPromise.resultD default, cancelTk? }
       }
       reusedAllHeaders := reusedAllHeaders && view.headerSnap?.any (·.old?.isSome)
     views := views.push view

src/Lean/Elab/Tactic/Basic.lean

@@ -165,6 +165,7 @@ structure EvalTacticFailure where
   state : SavedState
 
 partial def evalTactic (stx : Syntax) : TacticM Unit := do
+  checkSystem "tactic execution"
   profileitM Exception "tactic execution" (decl := stx.getKind) (← getOptions) <|
   withRef stx <| withIncRecDepth <| withFreshMacroScope <| match stx with
     | .node _ k _    =>
@@ -240,6 +241,7 @@ where
                     snap.old?.forM (·.val.cancelRec)
                   let promise ← IO.Promise.new
                   -- Store new unfolding in the snapshot tree
+                  let cancelTk? := (← readThe Core.Context).cancelTk?
                   snap.new.resolve {
                     stx := stx'
                     diagnostics := .empty
@@ -249,7 +251,7 @@ where
                       state? := (← Tactic.saveState)
                       moreSnaps := #[]
                     }
-                    next := #[{ stx? := stx', task := promise.resultD default }]
+                    next := #[{ stx? := stx', task := promise.resultD default, cancelTk? }]
                   }
                   -- Update `tacSnap?` to old unfolding
                   withTheReader Term.Context ({ · with tacSnap? := some {

src/Lean/Elab/Tactic/BuiltinTactic.lean

@@ -78,13 +78,14 @@ where
       let next ← IO.Promise.new
       let finished ← IO.Promise.new
       let inner ← IO.Promise.new
+      let cancelTk? := (← readThe Core.Context).cancelTk?
       snap.new.resolve {
         desc := tac.getKind.toString
         diagnostics := .empty
         stx := tac
-        inner? := some { stx? := tac, task := inner.resultD default }
-        finished := { stx? := tac, task := finished.resultD default }
-        next := #[{ stx? := stxs, task := next.resultD default }]
+        inner? := some { stx? := tac, task := inner.resultD default, cancelTk? }
+        finished := { stx? := tac, task := finished.resultD default, cancelTk? }
+        next := #[{ stx? := stxs, task := next.resultD default, cancelTk? }]
       }
       -- Run `tac` in a fresh info tree state and store resulting state in snapshot for
       -- incremental reporting, then add back saved trees. Here we rely on `evalTactic`

src/Lean/Elab/Tactic/Induction.lean

@@ -286,14 +286,15 @@ where
         -- them, eventually put each of them back in `Context.tacSnap?` in `applyAltStx`
         let finished ← IO.Promise.new
         let altPromises ← altStxs.mapM fun _ => IO.Promise.new
+        let cancelTk? := (← readThe Core.Context).cancelTk?
         tacSnap.new.resolve {
           -- save all relevant syntax here for comparison with next document version
           stx := mkNullNode altStxs
           diagnostics := .empty
           inner? := none
-          finished := { stx? := mkNullNode altStxs, reportingRange? := none, task := finished.resultD default }
+          finished := { stx? := mkNullNode altStxs, reportingRange? := none, task := finished.resultD default, cancelTk? }
           next := Array.zipWith
-            (fun stx prom => { stx? := some stx, task := prom.resultD default })
+            (fun stx prom => { stx? := some stx, task := prom.resultD default, cancelTk? })
             altStxs altPromises
         }
         goWithIncremental <| altPromises.mapIdx fun i prom => {

src/Lean/Language/Basic.lean

@@ -93,18 +93,17 @@ structure SnapshotTask (α : Type) where
   Cancellation token that can be set by the server to cancel the task when it detects the results
   are not needed anymore.
   -/
-  cancelTk? : Option IO.CancelToken := none
+  cancelTk? : Option IO.CancelToken
   /-- Underlying task producing the snapshot. -/
   task : Task α
 deriving Nonempty, Inhabited
 
 /-- Creates a snapshot task from the syntax processed by the task and a `BaseIO` action. -/
-def SnapshotTask.ofIO (stx? : Option Syntax)
+def SnapshotTask.ofIO (stx? : Option Syntax) (cancelTk? : Option IO.CancelToken)
     (reportingRange? : Option String.Range := defaultReportingRange? stx?) (act : BaseIO α) :
     BaseIO (SnapshotTask α) := do
   return {
-    stx?
-    reportingRange?
+    stx?, reportingRange?, cancelTk?
     task := (← BaseIO.asTask act)
   }
 
@@ -114,6 +113,7 @@ def SnapshotTask.finished (stx? : Option Syntax) (a : α) : SnapshotTask α wher
   -- irrelevant when already finished
   reportingRange? := none
   task := .pure a
+  cancelTk? := none
 
 /-- Transforms a task's output without changing the processed syntax. -/
 def SnapshotTask.map (t : SnapshotTask α) (f : α → β) (stx? : Option Syntax := t.stx?)

src/Lean/Language/Lean.lean

@@ -397,7 +397,7 @@ where
                     diagnostics := oldProcessed.diagnostics
                     result? := some {
                       cmdState := oldProcSuccess.cmdState
-                      firstCmdSnap := { stx? := none, task := prom.result! } } }
+                      firstCmdSnap := { stx? := none, task := prom.result!, cancelTk? := cancelTk } } }
               else
                 return .finished newStx oldProcessed) } }
       else return old
@@ -450,7 +450,7 @@ where
   processHeader (stx : Syntax) (parserState : Parser.ModuleParserState) :
       LeanProcessingM (SnapshotTask HeaderProcessedSnapshot) := do
     let ctx ← read
-    SnapshotTask.ofIO stx (some ⟨0, ctx.input.endPos⟩) <|
+    SnapshotTask.ofIO stx none (some ⟨0, ctx.input.endPos⟩) <|
     ReaderT.run (r := ctx) <|  -- re-enter reader in new task
     withHeaderExceptions (α := HeaderProcessedSnapshot) ({ · with result? := none }) do
       let setup ← match (← setupImports stx) with
@@ -507,7 +507,7 @@ where
         infoTree? := cmdState.infoState.trees[0]!
         result? := some {
           cmdState
-          firstCmdSnap := { stx? := none, task := prom.result! }
+          firstCmdSnap := { stx? := none, task := prom.result!, cancelTk? := cancelTk }
         }
       }
 
@@ -523,17 +523,19 @@ where
       -- from `old`
       if let some oldNext := old.nextCmdSnap? then do
         let newProm ← IO.Promise.new
+        let cancelTk ← IO.CancelToken.new
         -- can reuse range, syntax unchanged
         BaseIO.chainTask (sync := true) old.resultSnap.task fun oldResult =>
           -- also wait on old command parse snapshot as parsing is cheap and may allow for
           -- elaboration reuse
           BaseIO.chainTask (sync := true) oldNext.task fun oldNext => do
-            let cancelTk ← IO.CancelToken.new
             parseCmd oldNext newParserState oldResult.cmdState newProm sync cancelTk ctx
         prom.resolve <| { old with nextCmdSnap? := some {
           stx? := none
           reportingRange? := some ⟨newParserState.pos, ctx.input.endPos⟩
-          task := newProm.result! } }
+          task := newProm.result!
+          cancelTk? := cancelTk
+        } }
       else prom.resolve old  -- terminal command, we're done!
 
     -- fast path, do not even start new task for this snapshot (see [Incremental Parsing])
@@ -615,15 +617,16 @@ where
       })
       let diagnostics ← Snapshot.Diagnostics.ofMessageLog msgLog
 
-      -- use per-command cancellation token for elaboration so that
+      -- use per-command cancellation token for elaboration so that cancellation of further commands
+      -- does not affect current command
       let elabCmdCancelTk ← IO.CancelToken.new
       prom.resolve {
         diagnostics, nextCmdSnap?
         stx := stx', parserState := parserState'
         elabSnap := { stx? := stx', task := elabPromise.result!, cancelTk? := some elabCmdCancelTk }
-        resultSnap := { stx? := stx', reportingRange? := initRange?, task := resultPromise.result! }
-        infoTreeSnap := { stx? := stx', reportingRange? := initRange?, task := finishedPromise.result! }
-        reportSnap := { stx? := none, reportingRange? := initRange?, task := reportPromise.result! }
+        resultSnap := { stx? := stx', reportingRange? := initRange?, task := resultPromise.result!, cancelTk? := none }
+        infoTreeSnap := { stx? := stx', reportingRange? := initRange?, task := finishedPromise.result!, cancelTk? := none }
+        reportSnap := { stx? := none, reportingRange? := initRange?, task := reportPromise.result!, cancelTk? := none }
       }
       let cmdState ← doElab stx cmdState beginPos
         { old? := old?.map fun old => ⟨old.stx, old.elabSnap⟩, new := elabPromise }
@@ -665,8 +668,8 @@ where
           -- We want to trace all of `CommandParsedSnapshot` but `traceTask` is part of it, so let's
           -- create a temporary snapshot tree containing all tasks but it
           let snaps := #[
-            { stx? := stx', task := elabPromise.result!.map (sync := true) toSnapshotTree },
-            { stx? := stx', task := resultPromise.result!.map (sync := true) toSnapshotTree }] ++
+            { stx? := stx', task := elabPromise.result!.map (sync := true) toSnapshotTree, cancelTk? := none },
+            { stx? := stx', task := resultPromise.result!.map (sync := true) toSnapshotTree, cancelTk? := none }] ++
             cmdState.snapshotTasks
           let tree := SnapshotTree.mk { diagnostics := .empty } snaps
           BaseIO.bindTask (← tree.waitAll) fun _ => do
@@ -690,6 +693,7 @@ where
             stx? := none
             reportingRange? := initRange?
             task := traceTask
+            cancelTk? := none
           }
       if let some next := next? then
         -- We're definitely off the fast-forwarding path now

src/Lean/Meta/Basic.lean

@@ -2279,13 +2279,15 @@ def realizeConst (forConst : Name) (constName : Name) (realize : MetaM Unit) :
       initHeartbeats := (← IO.getNumHeartbeats)
     }
     let (env, exTask, dyn) ← env.realizeConst forConst constName (realizeAndReport coreCtx)
+    -- Realizations cannot be cancelled as their result is shared across elaboration runs
     let exAct ← Core.wrapAsyncAsSnapshot (cancelTk? := none) fun
       | none => return
       | some ex => do
         logError <| ex.toMessageData (← getOptions)
     Core.logSnapshotTask {
       stx? := none
       task := (← BaseIO.mapTask (t := exTask) exAct)
+      cancelTk? := none
     }
     if let some res := dyn.get? RealizeConstantResult then
       let mut snap := res.snap

src/Lean/Server/Test/Cancel.lean

@@ -56,6 +56,36 @@ elab_rules : tactic
 -- can't use a naked promise in `initialize` as marking it persistent would block
 initialize unblockedCancelTk : IO.CancelToken ← IO.CancelToken.new
 
+/--
+Waits for `unblock` to be called, which is expected to happen in a subsequent document version that
+does not invalidate this tactic. Complains if cancellation token was set before unblocking, i.e. if
+the tactic was invalidated after all.
+-/
+scoped syntax "wait_for_unblock" : tactic
+@[incremental]
+elab_rules : tactic
+| `(tactic| wait_for_unblock) => do
+  let ctx ← readThe Core.Context
+  let some cancelTk := ctx.cancelTk? | unreachable!
+
+  dbg_trace "blocked!"
+  log "blocked"
+  let ctx ← readThe Elab.Term.Context
+  let some tacSnap := ctx.tacSnap? | unreachable!
+  tacSnap.new.resolve {
+    diagnostics := (← Language.Snapshot.Diagnostics.ofMessageLog (← Core.getMessageLog))
+    stx := default
+    finished := default
+  }
+
+  while true do
+    if (← unblockedCancelTk.isSet) then
+      break
+    IO.sleep 30
+  if (← cancelTk.isSet) then
+    IO.eprintln "cancelled!"
+    log "cancelled (should never be visible)"
+
 /--
 Spawns a `logSnapshotTask` that waits for `unblock` to be called, which is expected to happen in a
 subsequent document version that does not invalidate this tactic. Complains if cancellation token
@@ -83,6 +113,10 @@ scoped elab "unblock" : tactic => do
   dbg_trace "unblocking!"
   unblockedCancelTk.set
 
+/--
+Like `wait_for_cancel_once` but does the waiting in a separate task and waits for its
+cancellation.
+-/
 scoped syntax "wait_for_cancel_once_async" : tactic
 @[incremental]
 elab_rules : tactic
@@ -110,3 +144,35 @@ elab_rules : tactic
 
   dbg_trace "blocked!"
   log "blocked"
+
+/--
+Like `wait_for_cancel_once_async` but waits for the main thread's cancellation token. This is useful
+to test main thread cancellation in non-incremental contexts because we otherwise wouldn't be able
+to send out the "blocked" message from there.
+-/
+scoped syntax "wait_for_main_cancel_once_async" : tactic
+@[incremental]
+elab_rules : tactic
+| `(tactic| wait_for_main_cancel_once_async) => do
+  let prom ← IO.Promise.new
+  if let some t := (← onceRef.modifyGet (fun old => (old, old.getD prom.result!))) then
+    IO.wait t
+    return
+
+  let some cancelTk := (← readThe Core.Context).cancelTk? | unreachable!
+  let act ← Elab.Term.wrapAsyncAsSnapshot (cancelTk? := none) fun _ => do
+    let ctx ← readThe Core.Context
+    -- TODO: `CancelToken` should probably use `Promise`
+    while true do
+      if (← cancelTk.isSet) then
+        break
+      IO.sleep 30
+    IO.eprintln "cancelled!"
+    log "cancelled (should never be visible)"
+    prom.resolve ()
+    Core.checkInterrupted
+  let t ← BaseIO.asTask (act ())
+  Core.logSnapshotTask { stx? := none, task := t, cancelTk? := cancelTk }
+
+  dbg_trace "blocked!"
+  log "blocked"