refactor: move module header-specific data into syntax-tagged snapshots

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 }

src/Lean/Language/Lean/Types.lean

@@ -89,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. -/
@@ -105,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. -/
@@ -114,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