refactor: make syntax covering snapshot tasks more precise on the top level (leanprover#8744)

Author: Kha

src/Lean/Elab/Frontend.lean

@@ -112,9 +112,9 @@ where
         |>.foldl (· ++ ·) {}
       -- In contrast to messages, we should collect info trees only from the top-level command
       -- snapshots as they subsume any info trees reported incrementally by their children.
-      let trees := commands.map (·.infoTreeSnap.get.infoTree?) |>.filterMap id |>.toPArray'
+      let trees := commands.map (·.elabSnap.infoTreeSnap.get.infoTree?) |>.filterMap id |>.toPArray'
       return {
-        commandState := { snap.resultSnap.get.cmdState with messages, infoState.trees := trees }
+        commandState := { snap.elabSnap.resultSnap.get.cmdState with messages, infoState.trees := trees }
         parserState := snap.parserState
         cmdPos := snap.parserState.pos
         commands := commands.map (·.stx)

src/Lean/Language/Basic.lean

@@ -218,7 +218,7 @@ partial def SnapshotTask.cancelRec [ToSnapshotTree α] (t : SnapshotTask α) : B
 
 /-- Snapshot type without child nodes. -/
 structure SnapshotLeaf extends Snapshot
-deriving Nonempty, TypeName
+deriving Inhabited, TypeName
 
 instance : ToSnapshotTree SnapshotLeaf where
   toSnapshotTree s := SnapshotTree.mk s.toSnapshot #[]

src/Lean/Language/Lean.lean

@@ -392,21 +392,27 @@ where
         return {
           ictx
           stx := newStx
-          diagnostics := old.diagnostics
+          diagnostics := .empty
+          metaSnap := .finished newStx {
+            diagnostics := old.diagnostics
+          }
           result? := some {
             parserState := newParserState
-            processedSnap := (← oldSuccess.processedSnap.bindIO (stx? := newStx)
+            processedSnap := (← oldSuccess.processedSnap.bindIO
                 (cancelTk? := none) (reportingRange? := progressRange?) (sync := true) fun oldProcessed => do
               if let some oldProcSuccess := oldProcessed.result? then
                 -- also wait on old command parse snapshot as parsing is cheap and may allow for
                 -- elaboration reuse
-                oldProcSuccess.firstCmdSnap.bindIO (sync := true) (stx? := newStx)
+                oldProcSuccess.firstCmdSnap.bindIO (sync := true)
                     (cancelTk? := none) (reportingRange? := progressRange?) fun oldCmd => do
                   let prom ← IO.Promise.new
                   let cancelTk ← IO.CancelToken.new
                   parseCmd oldCmd newParserState oldProcSuccess.cmdState prom (sync := true) cancelTk ctx
-                  return .finished newStx {
-                    diagnostics := oldProcessed.diagnostics
+                  return .finished none {
+                    diagnostics := .empty
+                    metaSnap := .finished newStx {
+                      diagnostics := oldProcessed.diagnostics
+                    }
                     result? := some {
                       cmdState := oldProcSuccess.cmdState
                       firstCmdSnap := { stx? := none, task := prom.result!, cancelTk? := cancelTk } } }
@@ -425,13 +431,16 @@ where
                 -- ...go immediately to next snapshot
                 return (← unchanged old old.stx oldSuccess.parserState)
 
-    withHeaderExceptions ({ · with ictx, stx := .missing, result? := none }) do
+    withHeaderExceptions ({ · with ictx, stx := .missing, result? := none, metaSnap := default }) do
       -- parsing the header should be cheap enough to do synchronously
       let (stx, parserState, msgLog) ← Parser.parseHeader ictx
       if msgLog.hasErrors then
         return {
           ictx, stx
-          diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
+          diagnostics := .empty
+          metaSnap := .finished stx {
+            diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
+          }
           result? := none
         }
 
@@ -452,7 +461,10 @@ where
         old.result?.forM (·.processedSnap.cancelRec)
       return {
         ictx, stx
-        diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
+        diagnostics := .empty
+        metaSnap := .finished stx {
+          diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
+        }
         result? := some {
           parserState
           processedSnap := (← processHeader ⟨trimmedStx⟩ parserState)
@@ -462,9 +474,9 @@ where
   processHeader (stx : HeaderSyntax) (parserState : Parser.ModuleParserState) :
       LeanProcessingM (SnapshotTask HeaderProcessedSnapshot) := do
     let ctx ← read
-    SnapshotTask.ofIO stx none (some ⟨0, ctx.input.endPos⟩) <|
+    SnapshotTask.ofIO none none (some ⟨0, ctx.input.endPos⟩) <|
     ReaderT.run (r := ctx) <|  -- re-enter reader in new task
-    withHeaderExceptions (α := HeaderProcessedSnapshot) ({ · with result? := none }) do
+    withHeaderExceptions (α := HeaderProcessedSnapshot) ({ · with result? := none, metaSnap := default }) do
       let setup ← match (← setupImports stx) with
         | .ok setup => pure setup
         | .error snap => return snap
@@ -483,7 +495,7 @@ where
       let stopTime := (← IO.monoNanosNow).toFloat / 1000000000
       let diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
       if msgLog.hasErrors then
-        return { diagnostics, result? := none }
+        return { diagnostics, result? := none, metaSnap := default }
 
       let mut traceState := default
       if trace.profiler.output.get? setup.opts |>.isSome then
@@ -521,8 +533,11 @@ where
       let cancelTk ← IO.CancelToken.new
       parseCmd none parserState cmdState prom (sync := true) cancelTk ctx
       return {
-        diagnostics
-        infoTree? := cmdState.infoState.trees[0]!
+        diagnostics := .empty
+        metaSnap := .finished stx {
+          diagnostics
+          infoTree? := cmdState.infoState.trees[0]!
+        }
         result? := some {
           cmdState
           firstCmdSnap := { stx? := none, task := prom.result!, cancelTk? := cancelTk }
@@ -543,7 +558,7 @@ where
         let newProm ← IO.Promise.new
         let cancelTk ← IO.CancelToken.new
         -- can reuse range, syntax unchanged
-        BaseIO.chainTask (sync := true) old.resultSnap.task fun oldResult =>
+        BaseIO.chainTask (sync := true) old.elabSnap.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
@@ -601,10 +616,13 @@ where
       -- is not `Inhabited`
       prom.resolve <| {
         diagnostics := .empty, stx := .missing, parserState
-        elabSnap := default
-        resultSnap := .finished none { diagnostics := .empty, cmdState }
-        infoTreeSnap := .finished none { diagnostics := .empty }
-        reportSnap := default
+        elabSnap := {
+          diagnostics := .empty
+          elabSnap := default
+          resultSnap := .finished none { diagnostics := .empty, cmdState }
+          infoTreeSnap := .finished none { diagnostics := .empty }
+          reportSnap := default
+        }
         nextCmdSnap? := none
       }
       return
@@ -641,13 +659,16 @@ where
       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!, cancelTk? := none }
-        infoTreeSnap := { stx? := stx', reportingRange? := initRange?, task := finishedPromise.result!, cancelTk? := none }
-        reportSnap := { stx? := none, reportingRange? := initRange?, task := reportPromise.result!, cancelTk? := none }
+        elabSnap := {
+          diagnostics := .empty
+          elabSnap := { stx? := stx', task := elabPromise.result!, cancelTk? := some elabCmdCancelTk }
+          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 }
+        { old? := old?.map fun old => ⟨old.stx, old.elabSnap.elabSnap⟩, new := elabPromise }
         elabCmdCancelTk ctx
 
       let mut reportedCmdState := cmdState
@@ -777,6 +798,6 @@ where goCmd snap :=
   if let some next := snap.nextCmdSnap? then
     goCmd next.get
   else
-    snap.resultSnap.get.cmdState
+    snap.elabSnap.resultSnap.get.cmdState
 
 end Lean

src/Lean/Language/Lean/Types.lean

@@ -36,12 +36,11 @@ deriving Nonempty
 instance : ToSnapshotTree CommandResultSnapshot where
   toSnapshotTree s := ⟨s.toSnapshot, #[]⟩
 
-/-- State after a command has been parsed. -/
-structure CommandParsedSnapshot extends Snapshot where
-  /-- Syntax tree of the command. -/
-  stx : Syntax
-  /-- Resulting parser state. -/
-  parserState : Parser.ModuleParserState
+/--
+State before a command is elaborated. This is separate from `CommandParsedSnapshot` so that all
+snapshots belonging to a command are grouped below a task with the command's syntax tree.
+-/
+structure CommandElaboratingSnapshot extends Snapshot where
   /--
   Snapshot for incremental reporting and reuse during elaboration, type dependent on specific
   elaborator.
@@ -55,16 +54,30 @@ structure CommandParsedSnapshot extends Snapshot where
   infoTreeSnap : SnapshotTask SnapshotLeaf
   /-- Additional, untyped snapshots used for reporting, not reuse. -/
   reportSnap : SnapshotTask SnapshotTree
-  /-- Next command, unless this is a terminal command. -/
-  nextCmdSnap? : Option (SnapshotTask CommandParsedSnapshot)
 deriving Nonempty
-partial instance : ToSnapshotTree CommandParsedSnapshot where
+instance : ToSnapshotTree CommandElaboratingSnapshot where
   toSnapshotTree := go where
     go s := ⟨s.toSnapshot,
       #[s.elabSnap.map (sync := true) toSnapshotTree,
         s.resultSnap.map (sync := true) toSnapshotTree,
         s.infoTreeSnap.map (sync := true) toSnapshotTree,
-        s.reportSnap] |>
+        s.reportSnap]⟩
+
+/-- State after a command has been parsed. -/
+structure CommandParsedSnapshot extends Snapshot where
+  /-- Syntax tree of the command. -/
+  stx : Syntax
+  /-- Resulting parser state. -/
+  parserState : Parser.ModuleParserState
+  /-- State before the command is elaborated. This snapshot is always fulfilled immediately. -/
+  elabSnap : CommandElaboratingSnapshot
+  /-- Next command, unless this is a terminal command. -/
+  nextCmdSnap? : Option (SnapshotTask CommandParsedSnapshot)
+deriving Nonempty
+partial instance : ToSnapshotTree CommandParsedSnapshot where
+  toSnapshotTree := go where
+    go s := ⟨s.toSnapshot,
+      #[.finished s.stx (toSnapshotTree s.elabSnap)] |>
         pushOpt (s.nextCmdSnap?.map (·.map (sync := true) go))⟩
 
 /-- State after successful importing. -/
@@ -76,11 +89,16 @@ structure HeaderProcessedState where
 
 /-- State after the module header has been processed including imports. -/
 structure HeaderProcessedSnapshot extends Snapshot where
+  /--
+  Holds produced diagnostics and info tree. Separate snapshot so that it can be tagged with the
+  header syntax, which should not be done for this snapshot containing `firstCmdSnap`.
+  -/
+  metaSnap : SnapshotTask SnapshotLeaf
   /-- State after successful importing. -/
   result? : Option HeaderProcessedState
   isFatal := result?.isNone
 instance : ToSnapshotTree HeaderProcessedSnapshot where
-  toSnapshotTree s := ⟨s.toSnapshot, #[] |>
+  toSnapshotTree s := ⟨s.toSnapshot, #[s.metaSnap.map (sync := true) toSnapshotTree] |>
     pushOpt (s.result?.map (·.firstCmdSnap.map (sync := true) toSnapshotTree))⟩
 
 /-- State after successfully parsing the module header. -/
@@ -92,6 +110,11 @@ structure HeaderParsedState where
 
 /-- State after the module header has been parsed. -/
 structure HeaderParsedSnapshot extends Snapshot where
+  /--
+  Holds produced diagnostics. Separate snapshot so that it can be tagged with the header syntax,
+  which should not be done for this snapshot containing `firstCmdSnap`.
+  -/
+  metaSnap : SnapshotTask SnapshotLeaf
   /-- Parser input context supplied by the driver, stored here for incremental parsing. -/
   ictx : Parser.InputContext
   /-- Resulting syntax tree. -/
@@ -101,8 +124,8 @@ structure HeaderParsedSnapshot extends Snapshot where
   isFatal := result?.isNone
 
 instance : ToSnapshotTree HeaderParsedSnapshot where
-  toSnapshotTree s := ⟨s.toSnapshot,
-    #[] |> pushOpt (s.result?.map (·.processedSnap.map (sync := true) toSnapshotTree))⟩
+  toSnapshotTree s := ⟨s.toSnapshot, #[s.metaSnap.map (sync := true) toSnapshotTree] |>
+    pushOpt (s.result?.map (·.processedSnap.map (sync := true) toSnapshotTree))⟩
 
 /-- Shortcut accessor to the final header state, if successful. -/
 def HeaderParsedSnapshot.processedResult (snap : HeaderParsedSnapshot) :

src/Lean/Server/FileWorker.lean

@@ -382,7 +382,7 @@ def setupImports
     unless (← IO.checkCanceled) do
       IO.Process.exit 2  -- signal restart request to watchdog
     -- should not be visible to user as task is already canceled
-    return .error { diagnostics := .empty, result? := none }
+    return .error { diagnostics := .empty, result? := none, metaSnap := default }
 
   chanOut.sync.send <| mkInitialIleanInfoUpdateNotification doc <| collectImports stx
   let imports := Elab.headerToImports stx
@@ -404,11 +404,13 @@ def setupImports
         "Imports are out of date and must be rebuilt; \
           use the \"Restart File\" command in your editor.")
       result? := none
+      metaSnap := default
     }
   | .error msg =>
     return .error {
       diagnostics := (← diagnosticsOfHeaderError msg)
       result? := none
+      metaSnap := default
     }
   | _ => pure ()
 

src/Lean/Server/FileWorker/Utils.lean

@@ -28,7 +28,7 @@ private partial def mkCmdSnaps (initSnap : Language.Lean.InitialSnapshot) :
     } <| .delayed <| headerSuccess.firstCmdSnap.task.asServerTask.bindCheap go
 where
   go cmdParsed :=
-    cmdParsed.resultSnap.task.asServerTask.mapCheap fun result =>
+    cmdParsed.elabSnap.resultSnap.task.asServerTask.mapCheap fun result =>
       .ok <| .cons {
         stx := cmdParsed.stx
         mpState := cmdParsed.parserState

src/Lean/Server/Requests.lean

@@ -396,7 +396,7 @@ def findCmdDataAtPos
   findCmdParsedSnap doc hoverPos |>.bindCheap fun
     | some cmdParsed => toSnapshotTree cmdParsed |>.findInfoTreeAtPos doc.meta.text hoverPos includeStop |>.bindCheap fun
       | some infoTree => .pure <| some (cmdParsed.stx, infoTree)
-      | none          => cmdParsed.infoTreeSnap.task.asServerTask.mapCheap fun s =>
+      | none          => cmdParsed.elabSnap.infoTreeSnap.task.asServerTask.mapCheap fun s =>
         assert! s.infoTree?.isSome
         some (cmdParsed.stx, s.infoTree?.get!)
     | none => .pure none