feat: snapshot trees and language processors (#3014)

This is the foundation for work on making processing in the language
server both more fine-grained (incremental tactics) as well as parallel.

CODEOWNERS

@@ -13,6 +13,7 @@
 /src/Lean/Data/Lsp/ @mhuisi
 /src/Lean/Elab/Deriving/ @semorrison
 /src/Lean/Elab/Tactic/ @semorrison
+/src/Lean/Language/ @Kha
 /src/Lean/Meta/Tactic/ @leodemoura
 /src/Lean/Parser/ @Kha
 /src/Lean/PrettyPrinter/ @Kha

src/Init/Data/Channel.lean

@@ -41,7 +41,7 @@ Sends a message on an `Channel`.
 
 This function does not block.
 -/
-def Channel.send (v : α) (ch : Channel α) : BaseIO Unit :=
+def Channel.send (ch : Channel α) (v : α) : BaseIO Unit :=
   ch.atomically do
     let st ← get
     if st.closed then return

src/Init/Data/Option/Basic.lean

@@ -37,6 +37,13 @@ def toMonad [Monad m] [Alternative m] : Option α → m α
   | none,   _ => none
   | some a, b => b a
 
+/-- Runs `f` on `o`'s value, if any, and returns its result, or else returns `none`.  -/
+@[inline] protected def bindM [Monad m] (f : α → m (Option β)) (o : Option α) : m (Option β) := do
+  if let some a := o then
+    return (← f a)
+  else
+    return none
+
 @[inline] protected def mapM [Monad m] (f : α → m β) (o : Option α) : m (Option β) := do
   if let some a := o then
     return some (← f a)

src/Lean/Data/JsonRpc.lean

@@ -183,6 +183,9 @@ structure ResponseError (α : Type u) where
 instance [ToJson α] : CoeOut (ResponseError α) Message :=
   ⟨fun r => Message.responseError r.id r.code r.message (r.data?.map toJson)⟩
 
+instance : CoeOut (ResponseError Unit) Message :=
+  ⟨fun r => Message.responseError r.id r.code r.message none⟩
+
 instance : Coe String RequestID := ⟨RequestID.str⟩
 instance : Coe JsonNumber RequestID := ⟨RequestID.num⟩
 

src/Lean/Data/Lsp/Internal.lean

@@ -84,6 +84,7 @@ structure RefInfo.Location where
   range       : Lsp.Range
   /-- Parent declaration of the reference. `none` if the reference is itself a declaration. -/
   parentDecl? : Option RefInfo.ParentDecl
+deriving Inhabited
 
 /-- Definition site and usage sites of a reference. Obtained from `Lean.Server.RefInfo`. -/
 structure RefInfo where

src/Lean/Data/Lsp/Ipc.lean

@@ -84,15 +84,22 @@ partial def readResponseAs (expectedID : RequestID) (α) [FromJson α] :
 def waitForExit : IpcM UInt32 := do
   (←read).wait
 
-/-- Waits for the worker to emit all diagnostics for the current document version
-and returns them as a list. -/
+/--
+Waits for the worker to emit all diagnostic notifications for the current document version and
+returns the last notification, if any.
+
+We used to return all notifications but with debouncing in the server, this would not be
+deterministic anymore as what messages are dropped depends on wall-clock timing.
+ -/
 partial def collectDiagnostics (waitForDiagnosticsId : RequestID := 0) (target : DocumentUri) (version : Nat)
-: IpcM (List (Notification PublishDiagnosticsParams)) := do
+: IpcM (Option (Notification PublishDiagnosticsParams)) := do
   writeRequest ⟨waitForDiagnosticsId, "textDocument/waitForDiagnostics", WaitForDiagnosticsParams.mk target version⟩
-  let rec loop : IpcM (List (Notification PublishDiagnosticsParams)) := do
+  loop
+where
+  loop := do
     match (←readMessage) with
     | Message.response id _ =>
-      if id == waitForDiagnosticsId then return []
+      if id == waitForDiagnosticsId then return none
       else loop
     | Message.responseError id _    msg _ =>
       if id == waitForDiagnosticsId then
@@ -100,10 +107,9 @@ partial def collectDiagnostics (waitForDiagnosticsId : RequestID := 0) (target :
       else loop
     | Message.notification "textDocument/publishDiagnostics" (some param) =>
       match fromJson? (toJson param) with
-      | Except.ok diagnosticParam => return ⟨"textDocument/publishDiagnostics", diagnosticParam⟩ :: (←loop)
+      | Except.ok diagnosticParam => return (← loop).getD ⟨"textDocument/publishDiagnostics", diagnosticParam⟩
       | Except.error inner => throw $ userError s!"Cannot decode publishDiagnostics parameters\n{inner}"
     | _ => loop
-  loop
 
 def runWith (lean : System.FilePath) (args : Array String := #[]) (test : IpcM α) : IO α := do
   let proc ← Process.spawn {

src/Lean/DeclarationRange.lean

@@ -48,14 +48,14 @@ def addDeclarationRanges [MonadEnv m] (declName : Name) (declRanges : Declaratio
 def findDeclarationRangesCore? [Monad m] [MonadEnv m] (declName : Name) : m (Option DeclarationRanges) :=
   return declRangeExt.find? (← getEnv) declName
 
-def findDeclarationRanges? [Monad m] [MonadEnv m] [MonadLiftT IO m] (declName : Name) : m (Option DeclarationRanges) := do
+def findDeclarationRanges? [Monad m] [MonadEnv m] [MonadLiftT BaseIO m] (declName : Name) : m (Option DeclarationRanges) := do
   let env ← getEnv
   let ranges ← if isAuxRecursor env declName || isNoConfusion env declName || (← isRec declName)  then
     findDeclarationRangesCore? declName.getPrefix
   else
     findDeclarationRangesCore? declName
   match ranges with
-  | none => return (← builtinDeclRanges.get (m := IO)).find? declName
+  | none => return (← builtinDeclRanges.get (m := BaseIO)).find? declName
   | some _ => return ranges
 
 end Lean

src/Lean/Elab/Command.lean

@@ -141,7 +141,8 @@ private def addTraceAsMessagesCore (ctx : Context) (log : MessageLog) (traceStat
   let mut log := log
   let traces' := traces.toArray.qsort fun ((a, _), _) ((b, _), _) => a < b
   for ((pos, endPos), traceMsg) in traces' do
-    log := log.add <| mkMessageCore ctx.fileName ctx.fileMap (.joinSep traceMsg.toList "\n") .information pos endPos
+    let data := .tagged `_traceMsg <| .joinSep traceMsg.toList "\n"
+    log := log.add <| mkMessageCore ctx.fileName ctx.fileMap data .information pos endPos
   return log
 
 private def addTraceAsMessages : CommandElabM Unit := do
@@ -268,11 +269,6 @@ instance : MonadRecDepth CommandElabM where
   getRecDepth      := return (← read).currRecDepth
   getMaxRecDepth   := return (← get).maxRecDepth
 
-register_builtin_option showPartialSyntaxErrors : Bool := {
-  defValue := false
-  descr    := "show elaboration errors from partial syntax trees (i.e. after parser recovery)"
-}
-
 builtin_initialize registerTraceClass `Elab.command
 
 partial def elabCommand (stx : Syntax) : CommandElabM Unit := do
@@ -321,11 +317,6 @@ def elabCommandTopLevel (stx : Syntax) : CommandElabM Unit := withRef stx do pro
     -- note the order: first process current messages & info trees, then add back old messages & trees,
     -- then convert new traces to messages
     let mut msgs := (← get).messages
-    -- `stx.hasMissing` should imply `initMsgs.hasErrors`, but the latter should be cheaper to check in general
-    if !showPartialSyntaxErrors.get (← getOptions) && initMsgs.hasErrors && stx.hasMissing then
-      -- discard elaboration errors, except for a few important and unlikely misleading ones, on parse error
-      msgs := ⟨msgs.msgs.filter fun msg =>
-        msg.data.hasTag (fun tag => tag == `Elab.synthPlaceholder || tag == `Tactic.unsolvedGoals || (`_traceMsg).isSuffixOf tag)⟩
     for tree in (← getInfoTrees) do
       trace[Elab.info] (← tree.format)
     modify fun st => { st with

src/Lean/Elab/Frontend.lean

@@ -4,8 +4,7 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura, Sebastian Ullrich
 -/
 prelude
-import Lean.Elab.Import
-import Lean.Elab.Command
+import Lean.Language.Lean
 import Lean.Util.Profile
 import Lean.Server.References
 
@@ -40,7 +39,19 @@ def setCommandState (commandState : Command.State) : FrontendM Unit :=
 
 def elabCommandAtFrontend (stx : Syntax) : FrontendM Unit := do
   runCommandElabM do
+    let initMsgs ← modifyGet fun st => (st.messages, { st with messages := {} })
     Command.elabCommandTopLevel stx
+    let mut msgs := (← get).messages
+    -- `stx.hasMissing` should imply `initMsgs.hasErrors`, but the latter should be cheaper to check
+    -- in general
+    if !Language.Lean.showPartialSyntaxErrors.get (← getOptions) && initMsgs.hasErrors &&
+        stx.hasMissing then
+      -- discard elaboration errors, except for a few important and unlikely misleading ones, on
+      -- parse error
+      msgs := ⟨msgs.msgs.filter fun msg =>
+        msg.data.hasTag (fun tag => tag == `Elab.synthPlaceholder ||
+          tag == `Tactic.unsolvedGoals || (`_traceMsg).isSuffixOf tag)⟩
+    modify ({ · with messages := initMsgs ++ msgs })
 
 def updateCmdPos : FrontendM Unit := do
   modify fun s => { s with cmdPos := s.parserState.pos }
@@ -86,12 +97,8 @@ def process (input : String) (env : Environment) (opts : Options) (fileName : Op
   pure (s.commandState.env, s.commandState.messages)
 
 builtin_initialize
-  registerOption `printMessageEndPos { defValue := false, descr := "print end position of each message in addition to start position" }
   registerTraceClass `Elab.info
 
-def getPrintMessageEndPos (opts : Options) : Bool :=
-  opts.getBool `printMessageEndPos false
-
 @[export lean_run_frontend]
 def runFrontend
     (input : String)
@@ -102,26 +109,50 @@ def runFrontend
     (ileanFileName? : Option String := none)
     : IO (Environment × Bool) := do
   let inputCtx := Parser.mkInputContext input fileName
-  let (header, parserState, messages) ← Parser.parseHeader inputCtx
-  -- allow `env` to be leaked, which would live until the end of the process anyway
-  let (env, messages) ← processHeader (leakEnv := true) header opts messages inputCtx trustLevel
-  let env := env.setMainModule mainModuleName
-  let mut commandState := Command.mkState env messages opts
+  -- TODO: replace with `#lang` processing
+  if /- Lean #lang? -/ true then
+    -- Temporarily keep alive old cmdline driver for the Lean language so that we don't pay the
+    -- overhead of passing the environment between snapshots until we actually make good use of it
+    -- outside the server
+    let (header, parserState, messages) ← Parser.parseHeader inputCtx
+    -- allow `env` to be leaked, which would live until the end of the process anyway
+    let (env, messages) ← processHeader (leakEnv := true) header opts messages inputCtx trustLevel
+    let env := env.setMainModule mainModuleName
+    let mut commandState := Command.mkState env messages opts
 
-  if ileanFileName?.isSome then
-    -- Collect InfoTrees so we can later extract and export their info to the ilean file
-    commandState := { commandState with infoState.enabled := true }
+    if ileanFileName?.isSome then
+      -- Collect InfoTrees so we can later extract and export their info to the ilean file
+      commandState := { commandState with infoState.enabled := true }
 
-  let s ← IO.processCommands inputCtx parserState commandState
-  for msg in s.commandState.messages.toList do
-    IO.print (← msg.toString (includeEndPos := getPrintMessageEndPos opts))
+    let s ← IO.processCommands inputCtx parserState commandState
+    for msg in s.commandState.messages.toList do
+      IO.print (← msg.toString (includeEndPos := Language.printMessageEndPos.get opts))
 
+    if let some ileanFileName := ileanFileName? then
+      let trees := s.commandState.infoState.trees.toArray
+      let references ←
+        Lean.Server.findModuleRefs inputCtx.fileMap trees (localVars := false) |>.toLspModuleRefs
+      let ilean := { module := mainModuleName, references : Lean.Server.Ilean }
+      IO.FS.writeFile ileanFileName $ Json.compress $ toJson ilean
+
+    return (s.commandState.env, !s.commandState.messages.hasErrors)
+
+  let ctx := { inputCtx with mainModuleName, opts, trustLevel }
+  let processor := Language.Lean.process
+  let snap ← processor none ctx
+  let snaps := Language.toSnapshotTree snap
+  snaps.runAndReport opts
   if let some ileanFileName := ileanFileName? then
-    let trees := s.commandState.infoState.trees.toArray
+    let trees := snaps.getAll.concatMap (match ·.infoTree? with | some t => #[t] | _ => #[])
     let references := Lean.Server.findModuleRefs inputCtx.fileMap trees (localVars := false)
     let ilean := { module := mainModuleName, references := ← references.toLspModuleRefs : Lean.Server.Ilean }
     IO.FS.writeFile ileanFileName $ Json.compress $ toJson ilean
 
-  pure (s.commandState.env, !s.commandState.messages.hasErrors)
+  let hasErrors := snaps.getAll.any (·.diagnostics.msgLog.hasErrors)
+  -- TODO: remove default when reworking cmdline interface in Lean; currently the only case
+  -- where we use the environment despite errors in the file is `--stats`
+  let env := Language.Lean.waitForFinalEnv? snap |>.getD (← mkEmptyEnvironment)
+  pure (env, !hasErrors)
+
 
 end Lean.Elab

src/Lean/Elab/InfoTree/Main.lean

@@ -94,16 +94,18 @@ partial def InfoTree.substitute (tree : InfoTree) (assignment : PersistentHashMa
     | none      => hole id
     | some tree => substitute tree assignment
 
-def ContextInfo.runMetaM (info : ContextInfo) (lctx : LocalContext) (x : MetaM α) : IO α := do
-  let x := x.run { lctx := lctx } { mctx := info.mctx }
+/-- Embeds a `CoreM` action in `IO` by supplying the information stored in `info`. -/
+def ContextInfo.runCoreM (info : ContextInfo) (x : CoreM α) : IO α := do
   /-
     We must execute `x` using the `ngen` stored in `info`. Otherwise, we may create `MVarId`s and `FVarId`s that
     have been used in `lctx` and `info.mctx`.
   -/
-  let ((a, _), _) ←
+  (·.1) <$>
     x.toIO { options := info.options, currNamespace := info.currNamespace, openDecls := info.openDecls, fileName := "<InfoTree>", fileMap := default }
            { env := info.env, ngen := info.ngen }
-  return a
+
+def ContextInfo.runMetaM (info : ContextInfo) (lctx : LocalContext) (x : MetaM α) : IO α := do
+  (·.1) <$> info.runCoreM (x.run { lctx := lctx } { mctx := info.mctx })
 
 def ContextInfo.toPPContext (info : ContextInfo) (lctx : LocalContext) : PPContext :=
   { env  := info.env, mctx := info.mctx, lctx := lctx,

src/Lean/Language/Basic.lean

@@ -0,0 +1,222 @@
+/-
+Copyright (c) 2023 Lean FRO. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+The generic interface of a `#lang` language processor used by the language server and cmdline
+driver. See the [server readme](../Server/README.md#worker-architecture) for an overview.
+
+Authors: Sebastian Ullrich
+-/
+
+prelude
+import Lean.Message
+import Lean.Parser.Types
+
+set_option linter.missingDocs true
+
+namespace Lean.Language
+
+/-- `MessageLog` with interactive diagnostics. -/
+structure Snapshot.Diagnostics where
+  /-- Non-interactive message log. -/
+  msgLog : MessageLog
+  /--
+  Dynamic mutable slot usable by the language server for memorizing interactive diagnostics. If
+  `none`, interactive diagnostics are not remembered, which should only be used for messages not
+  containing any interactive elements as client-side state will be lost on recreating a diagnostic.
+
+  See also section "Communication" in Lean/Server/README.md.
+  -/
+  interactiveDiagsRef? : Option (IO.Ref (Option Dynamic))
+deriving Inhabited
+
+/-- The empty set of diagnostics. -/
+def Snapshot.Diagnostics.empty : Snapshot.Diagnostics where
+  msgLog := .empty
+  -- nothing to memorize
+  interactiveDiagsRef? := none
+
+/--
+  The base class of all snapshots: all the generic information the language server needs about a
+  snapshot. -/
+structure Snapshot where
+  /--
+    The messages produced by this step. The union of message logs of all finished snapshots is
+    reported to the user. -/
+  diagnostics : Snapshot.Diagnostics
+  /-- General elaboration metadata produced by this step. -/
+  infoTree? : Option Elab.InfoTree := none
+  /--
+  Whether it should be indicated to the user that a fatal error (which should be part of
+  `diagnostics`) occurred that prevents processing of the remainder of the file.
+  -/
+  isFatal := false
+deriving Inhabited
+
+/-- A task producing some snapshot type (usually a subclass of `Snapshot`). -/
+-- Longer-term TODO: Give the server more control over the priority of tasks, depending on e.g. the
+-- cursor position. This may require starting the tasks suspended (e.g. in `Thunk`). The server may
+-- also need more dependency information for this in order to avoid priority inversion.
+structure SnapshotTask (α : Type) where
+  /-- Range that is marked as being processed by the server while the task is running. -/
+  range : String.Range
+  /-- Underlying task producing the snapshot. -/
+  task : Task α
+deriving Nonempty
+
+/-- Creates a snapshot task from a reporting range and a `BaseIO` action. -/
+def SnapshotTask.ofIO (range : String.Range) (act : BaseIO α) : BaseIO (SnapshotTask α) := do
+  return {
+    range
+    task := (← BaseIO.asTask act)
+  }
+
+/-- Creates a finished snapshot task. -/
+def SnapshotTask.pure (a : α) : SnapshotTask α where
+  -- irrelevant when already finished
+  range := default
+  task := .pure a
+
+/--
+  Explicitly cancels a tasks. Like with basic `Tasks`s, cancellation happens implicitly when the
+  last reference to the task is dropped *if* it is not an I/O task. -/
+def SnapshotTask.cancel (t : SnapshotTask α) : BaseIO Unit :=
+  IO.cancel t.task
+
+/-- Transforms a task's output without changing the reporting range. -/
+def SnapshotTask.map (t : SnapshotTask α) (f : α → β) (range : String.Range := t.range)
+    (sync := false) : SnapshotTask β :=
+  { range, task := t.task.map (sync := sync) f }
+
+/--
+  Chains two snapshot tasks. The range is taken from the first task if not specified; the range of
+  the second task is discarded. -/
+def SnapshotTask.bind (t : SnapshotTask α) (act : α → SnapshotTask β)
+    (range : String.Range := t.range) (sync := false) : SnapshotTask β :=
+  { range, task := t.task.bind (sync := sync) (act · |>.task) }
+
+/--
+  Chains two snapshot tasks. The range is taken from the first task if not specified; the range of
+  the second task is discarded. -/
+def SnapshotTask.bindIO (t : SnapshotTask α) (act : α → BaseIO (SnapshotTask β))
+    (range : String.Range := t.range) (sync := false) : BaseIO (SnapshotTask β) :=
+  return { range, task := (← BaseIO.bindTask (sync := sync) t.task fun a => (·.task) <$> (act a)) }
+
+/-- Synchronously waits on the result of the task. -/
+def SnapshotTask.get (t : SnapshotTask α) : α :=
+  t.task.get
+
+/-- Returns task result if already finished or else `none`. -/
+def SnapshotTask.get? (t : SnapshotTask α) : BaseIO (Option α) :=
+  return if (← IO.hasFinished t.task) then some t.task.get else none
+
+/--
+  Tree of snapshots where each snapshot comes with an array of asynchronous further subtrees. Used
+  for asynchronously collecting information about the entirety of snapshots in the language server.
+  The involved tasks may form a DAG on the `Task` dependency level but this is not captured by this
+  data structure. -/
+inductive SnapshotTree where
+  /-- Creates a snapshot tree node. -/
+  | mk (element : Snapshot) (children : Array (SnapshotTask SnapshotTree))
+deriving Nonempty
+
+/-- The immediately available element of the snapshot tree node. -/
+abbrev SnapshotTree.element : SnapshotTree → Snapshot
+  | mk s _ => s
+/-- The asynchronously available children of the snapshot tree node. -/
+abbrev SnapshotTree.children : SnapshotTree → Array (SnapshotTask SnapshotTree)
+  | mk _ children => children
+
+/--
+  Helper class for projecting a heterogeneous hierarchy of snapshot classes to a homogeneous
+  representation. -/
+class ToSnapshotTree (α : Type) where
+  /-- Transforms a language-specific snapshot to a homogeneous snapshot tree. -/
+  toSnapshotTree : α → SnapshotTree
+export ToSnapshotTree (toSnapshotTree)
+
+/--
+  Option for printing end position of each message in addition to start position. Used for testing
+  message ranges in the test suite. -/
+register_builtin_option printMessageEndPos : Bool := {
+  defValue := false, descr := "print end position of each message in addition to start position"
+}
+
+/--
+  Runs a tree of snapshots to conclusion and incrementally report messages on stdout. Messages are
+  reported in tree preorder.
+  This function is used by the cmdline driver; see `Lean.Server.FileWorker.reportSnapshots` for how
+  the language server reports snapshots asynchronously.  -/
+partial def SnapshotTree.runAndReport (s : SnapshotTree) (opts : Options) : IO Unit := do
+  s.element.diagnostics.msgLog.forM
+    (·.toString (includeEndPos := printMessageEndPos.get opts) >>= IO.print)
+  for t in s.children do
+    t.get.runAndReport opts
+
+/-- Waits on and returns all snapshots in the tree. -/
+partial def SnapshotTree.getAll (s : SnapshotTree) : Array Snapshot :=
+  go s |>.run #[] |>.2
+where
+  go s : StateM (Array Snapshot) Unit := do
+    modify (·.push s.element)
+    for t in s.children do
+      go t.get
+
+/-- Metadata that does not change during the lifetime of the language processing process. -/
+structure ModuleProcessingContext where
+  /-- Module name of the file being processed. -/
+  mainModuleName : Name
+  /-- Options provided outside of the file content, e.g. on the cmdline or in the lakefile. -/
+  opts : Options
+  /-- Kernel trust level. -/
+  trustLevel : UInt32 := 0
+
+/-- Context of an input processing invocation. -/
+structure ProcessingContext extends ModuleProcessingContext, Parser.InputContext
+
+/-- Monad transformer holding all relevant data for processing. -/
+abbrev ProcessingT m := ReaderT ProcessingContext m
+/-- Monad holding all relevant data for processing. -/
+abbrev ProcessingM := ProcessingT BaseIO
+
+instance : MonadLift ProcessingM (ProcessingT IO) where
+  monadLift := fun act ctx => act ctx
+
+/--
+Creates snapshot message log from non-interactive message log, also allocating a mutable cell
+that can be used by the server to memorize interactive diagnostics derived from the log.
+-/
+def Snapshot.Diagnostics.ofMessageLog (msgLog : Lean.MessageLog) :
+    ProcessingM Snapshot.Diagnostics := do
+  return { msgLog, interactiveDiagsRef? := some (← IO.mkRef none) }
+
+/-- Creates diagnostics from a single error message that should span the whole file. -/
+def diagnosticsOfHeaderError (msg : String) : ProcessingM Snapshot.Diagnostics := do
+  let msgLog := MessageLog.empty.add {
+    fileName := "<input>"
+    pos := ⟨0, 0⟩
+    endPos := (← read).fileMap.toPosition (← read).fileMap.source.endPos
+    data := msg
+  }
+  Snapshot.Diagnostics.ofMessageLog msgLog
+
+/--
+  Adds unexpected exceptions from header processing to the message log as a last resort; standard
+  errors should already have been caught earlier. -/
+def withHeaderExceptions (ex : Snapshot → α) (act : ProcessingT IO α) : ProcessingM α := do
+  match (← (act (← read)).toBaseIO) with
+  | .error e => return ex { diagnostics := (← diagnosticsOfHeaderError e.toString) }
+  | .ok a => return a
+
+end Language
+
+/--
+  Builds a function for processing a language using incremental snapshots by passing the previous
+  snapshot to `Language.process` on subsequent invocations. -/
+partial def Language.mkIncrementalProcessor (process : Option InitSnap → ProcessingM InitSnap)
+    (ctx : ModuleProcessingContext) : BaseIO (Parser.InputContext → BaseIO InitSnap) := do
+  let oldRef ← IO.mkRef none
+  return fun ictx => do
+    let snap ← process (← oldRef.get) { ctx, ictx with }
+    oldRef.set (some snap)
+    return snap

src/Lean/Language/Lean.lean

@@ -0,0 +1,473 @@
+/-
+Copyright (c) 2023 Lean FRO. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Implementation of the Lean language: parsing and processing of header and commands, incremental
+recompilation
+
+Authors: Sebastian Ullrich
+-/
+
+prelude
+import Lean.Language.Basic
+import Lean.Parser.Module
+import Lean.Elab.Command
+import Lean.Elab.Import
+
+/-!
+# Note [Incremental Parsing]
+
+In the language server, we want to minimize the work we do after each edit by reusing previous state
+where possible. This is true for both parsing, i.e. reusing syntax trees without running the parser,
+and elaboration. For both, we currently assume that we have to reprocess at least everything from
+the point of change downwards. This note is about how to find the correct starting point for
+incremental parsing; for elaboration, we then start with the first changed syntax tree.
+
+One initial thought about incremental parsing could be that it's not necessary as parsing is very
+fast compared to elaboration; on mathlib we average 41ms parsing per 1000 LoC. But there are quite a
+few files >= 1kloc (up to 4.5kloc) in there, so near the end of such files lag from always reparsing
+from the beginning may very well be noticeable.
+
+So if we do want incremental parsing, another thought might be that a user edit can only invalidate
+commands at or after the location of the change. Unfortunately, that's not true; take the (partial)
+input `def a := b private def c`. If we remove the space after `private`, the two commands
+syntactically become one with an application of `privatedef` to `b` even though the edit was
+strictly after the end of the first command.
+
+So obviously we must include at least the extent of the token that made the parser stop parsing a
+command as well such that invalidating the private token invalidates the preceding command.
+Unfortunately this is not sufficient either, given the following input:
+```
+structure a where /-- b -/ @[c] private axiom d : Nat
+```
+This is a syntactically valid sequence of a field-less structure and a declaration. If we again
+delete the space after private, it becomes a syntactically correct structure with a single field
+privateaxiom! So clearly, because of uses of atomic in the grammar, an edit can affect a command
+syntax tree even across multiple tokens.
+
+Now, what we do today, and have done since Lean 3, is to always reparse the last command completely
+preceding the edit location. If its syntax tree is unchanged, we preserve its data and reprocess all
+following commands only, otherwise we reprocess it fully as well. This seems to have worked well so
+far but it does seem a bit arbitrary given that even if it works for our current grammar, it can
+certainly be extended in ways that break the assumption.
+
+Finally, a more actually principled and generic solution would be to invalidate a syntax tree when
+the parser has reached the edit location during parsing. If it did not, surely the edit cannot have
+an effect on the syntax tree in question. Sadly such a "high-water mark" parser position does not
+exist currently and likely it could at best be approximated by e.g. "furthest `tokenFn` parse". Thus
+we remain at "go two commands up" at this point.
+-/
+
+set_option linter.missingDocs true
+
+namespace Lean.Language.Lean
+open Lean.Elab
+open Lean.Parser
+
+private def pushOpt (a? : Option α) (as : Array α) : Array α :=
+  match a? with
+  | some a => as.push a
+  | none   => as
+
+/-- Option for capturing output to stderr during elaboration. -/
+register_builtin_option stderrAsMessages : Bool := {
+  defValue := false
+  group    := "server"
+  descr    := "(server) capture output to the Lean stderr channel (such as from `dbg_trace`) during elaboration of a command as a diagnostic message"
+}
+
+/-- Option for showing elaboration errors from partial syntax errors. -/
+register_builtin_option showPartialSyntaxErrors : Bool := {
+  defValue := false
+  descr    := "show elaboration errors from partial syntax trees (i.e. after parser recovery)"
+}
+
+/-! The hierarchy of Lean snapshot types -/
+
+/-- Final state of processing of a command. -/
+structure CommandFinishedSnapshot extends Snapshot where
+  /-- Resulting elaboration state. -/
+  cmdState : Command.State
+deriving Nonempty
+instance : ToSnapshotTree CommandFinishedSnapshot where
+  toSnapshotTree s := ⟨s.toSnapshot, #[]⟩
+
+/--
+  State after processing a command's signature and before executing its tactic body, if any. Other
+  commands should immediately proceed to `finished`. -/
+-- TODO: tactics
+structure CommandSignatureProcessedSnapshot extends Snapshot where
+  /-- State after processing is finished. -/
+  finishedSnap : SnapshotTask CommandFinishedSnapshot
+deriving Nonempty
+instance : ToSnapshotTree CommandSignatureProcessedSnapshot where
+  toSnapshotTree s := ⟨s.toSnapshot, #[s.finishedSnap.map (sync := true) toSnapshotTree]⟩
+
+/-- State after a command has been parsed. -/
+structure CommandParsedSnapshotData extends Snapshot where
+  /-- Syntax tree of the command. -/
+  stx : Syntax
+  /-- Resulting parser state. -/
+  parserState : Parser.ModuleParserState
+  /-- Signature processing task. -/
+  sigSnap : SnapshotTask CommandSignatureProcessedSnapshot
+deriving Nonempty
+/-- State after a command has been parsed. -/
+-- workaround for lack of recursive structures
+inductive CommandParsedSnapshot where
+  /-- Creates a command parsed snapshot. -/
+  | mk (data : CommandParsedSnapshotData)
+    (nextCmdSnap? : Option (SnapshotTask CommandParsedSnapshot))
+deriving Nonempty
+/-- The snapshot data. -/
+abbrev CommandParsedSnapshot.data : CommandParsedSnapshot → CommandParsedSnapshotData
+  | mk data _ => data
+/-- Next command, unless this is a terminal command. -/
+-- It would be really nice to not make this depend on `sig.finished` where possible
+abbrev CommandParsedSnapshot.next? : CommandParsedSnapshot →
+    Option (SnapshotTask CommandParsedSnapshot)
+  | mk _ next? => next?
+partial instance : ToSnapshotTree CommandParsedSnapshot where
+  toSnapshotTree := go where
+    go s := ⟨s.data.toSnapshot,
+      #[s.data.sigSnap.map (sync := true) toSnapshotTree] |>
+        pushOpt (s.next?.map (·.map (sync := true) go))⟩
+
+/-- Cancels all significant computations from this snapshot onwards. -/
+partial def CommandParsedSnapshot.cancel (snap : CommandParsedSnapshot) : BaseIO Unit := do
+  -- This is the only relevant computation right now
+  -- TODO: cancel additional elaboration tasks if we add them without switching to implicit
+  -- cancellation
+  snap.data.sigSnap.cancel
+  if let some next := snap.next? then
+    -- recurse on next command (which may have been spawned just before we cancelled above)
+    let _ ← IO.mapTask (sync := true) (·.cancel) next.task
+
+/-- State after successful importing. -/
+structure HeaderProcessedState where
+  /-- The resulting initial elaboration state. -/
+  cmdState : Command.State
+  /-- First command task (there is always at least a terminal command). -/
+  firstCmdSnap : SnapshotTask CommandParsedSnapshot
+
+/-- State after the module header has been processed including imports. -/
+structure HeaderProcessedSnapshot extends Snapshot where
+  /-- State after successful importing. -/
+  result? : Option HeaderProcessedState
+  isFatal := result?.isNone
+instance : ToSnapshotTree HeaderProcessedSnapshot where
+  toSnapshotTree s := ⟨s.toSnapshot, #[] |>
+    pushOpt (s.result?.map (·.firstCmdSnap.map (sync := true) toSnapshotTree))⟩
+
+/-- State after successfully parsing the module header. -/
+structure HeaderParsedState where
+  /-- Resulting parser state. -/
+  parserState : Parser.ModuleParserState
+  /-- Header processing task. -/
+  processedSnap : SnapshotTask HeaderProcessedSnapshot
+
+/-- State after the module header has been parsed. -/
+structure HeaderParsedSnapshot extends Snapshot where
+  /-- Parser input context supplied by the driver, stored here for incremental parsing. -/
+  ictx : Parser.InputContext
+  /-- Resulting syntax tree. -/
+  stx : Syntax
+  /-- State after successful parsing. -/
+  result? : Option HeaderParsedState
+  isFatal := result?.isNone
+
+instance : ToSnapshotTree HeaderParsedSnapshot where
+  toSnapshotTree s := ⟨s.toSnapshot,
+    #[] |> pushOpt (s.result?.map (·.processedSnap.map (sync := true) toSnapshotTree))⟩
+
+/-- Shortcut accessor to the final header state, if successful. -/
+def HeaderParsedSnapshot.processedResult (snap : HeaderParsedSnapshot) :
+    SnapshotTask (Option HeaderProcessedState) :=
+  snap.result?.bind (·.processedSnap.map (sync := true) (·.result?)) |>.getD (.pure none)
+
+/-- Initial snapshot of the Lean language processor: a "header parsed" snapshot. -/
+abbrev InitialSnapshot := HeaderParsedSnapshot
+
+/-- Lean-specific processing context. -/
+structure LeanProcessingContext extends ProcessingContext where
+  /-- Position of the first file difference if there was a previous invocation. -/
+  firstDiffPos? : Option String.Pos
+
+/-- Monad transformer holding all relevant data for Lean processing. -/
+abbrev LeanProcessingT m := ReaderT LeanProcessingContext m
+/-- Monad holding all relevant data for Lean processing. -/
+abbrev LeanProcessingM := LeanProcessingT BaseIO
+
+instance : MonadLift LeanProcessingM (LeanProcessingT IO) where
+  monadLift := fun act ctx => act ctx
+
+instance : MonadLift (ProcessingT m) (LeanProcessingT m) where
+  monadLift := fun act ctx => act ctx.toProcessingContext
+
+/--
+Returns true if there was a previous run and the given position is before any textual change
+compared to it.
+-/
+def isBeforeEditPos (pos : String.Pos) : LeanProcessingM Bool := do
+  return (← read).firstDiffPos?.any (pos < ·)
+
+/--
+  Adds unexpected exceptions from header processing to the message log as a last resort; standard
+  errors should already have been caught earlier. -/
+private def withHeaderExceptions (ex : Snapshot → α) (act : LeanProcessingT IO α) :
+    LeanProcessingM α := do
+  match (← (act (← read)).toBaseIO) with
+  | .error e => return ex { diagnostics := (← diagnosticsOfHeaderError e.toString) }
+  | .ok a => return a
+
+/-- Entry point of the Lean language processor. -/
+/-
+General notes:
+* For each processing function we pass in the previous state, if any, in order to reuse still-valid
+  state. As there is no cheap way to check whether the `Environment` is unchanged, i.e. *semantic*
+  change detection is currently not possible, we must make sure to pass `none` as all follow-up
+  "previous states" from the first *syntactic* change onwards.
+* We must make sure to use `CommandParsedSnapshot.cancel` on such tasks when discarding them, i.e.
+  when not passing them along in `old?`.
+* Control flow up to finding the last still-valid snapshot (which should be quick) is synchronous so
+  as not to report this "fast forwarding" to the user as well as to make sure the next run sees all
+  fast-forwarded snapshots without having to wait on tasks.
+-/
+partial def process
+    (setupImports : Syntax → ProcessingT IO (Except HeaderProcessedSnapshot Options) :=
+      fun _ => pure <| .ok {})
+    (old? : Option InitialSnapshot) : ProcessingM InitialSnapshot := do
+  -- compute position of syntactic change once
+  let firstDiffPos? := old?.map (·.ictx.input.firstDiffPos (← read).input)
+  ReaderT.adapt ({ · with firstDiffPos? }) do
+    parseHeader old?
+where
+  parseHeader (old? : Option HeaderParsedSnapshot) : LeanProcessingM HeaderParsedSnapshot := do
+    let ctx ← read
+    let ictx := ctx.toInputContext
+    let unchanged old :=
+      -- when header syntax is unchanged, reuse import processing task as is and continue with
+      -- parsing the first command, synchronously if possible
+      if let some oldSuccess := old.result? then
+        return { old with ictx, result? := some { oldSuccess with
+          processedSnap := (← oldSuccess.processedSnap.bindIO (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) fun oldCmd =>
+                return .pure { oldProcessed with result? := some { oldProcSuccess with
+                  firstCmdSnap := (← parseCmd oldCmd oldSuccess.parserState oldProcSuccess.cmdState ctx) } }
+            else
+              return .pure oldProcessed) } }
+      else return old
+
+    -- fast path: if we have parsed the header successfully...
+    if let some old := old? then
+      if let some (some processed) ← old.processedResult.get? then
+        -- ...and the edit location is after the next command (see note [Incremental Parsing])...
+        if let some nextCom ← processed.firstCmdSnap.get? then
+          if (← isBeforeEditPos nextCom.data.parserState.pos) then
+            -- ...go immediately to next snapshot
+            return (← unchanged old)
+
+    withHeaderExceptions ({ · with ictx, stx := .missing, result? := none }) 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)
+          result? := none
+        }
+
+      -- semi-fast path: go to next snapshot if syntax tree is unchanged AND we're still in front
+      -- of the edit location
+      -- TODO: dropping the second condition would require adjusting positions in the state
+      -- NOTE: as `parserState.pos` includes trailing whitespace, this forces reprocessing even if
+      -- only that whitespace changes, which is wasteful but still necessary because it may
+      -- influence the range of error messages such as from a trailing `exact`
+      if let some old := old? then
+        if (← isBeforeEditPos parserState.pos) && old.stx == stx then
+          return (← unchanged old)
+        -- on first change, make sure to cancel all further old tasks
+        if let some oldSuccess := old.result? then
+          oldSuccess.processedSnap.cancel
+          let _ ← BaseIO.mapTask (t := oldSuccess.processedSnap.task) fun processed => do
+            if let some oldProcSuccess := processed.result? then
+              let _ ← BaseIO.mapTask (·.cancel) oldProcSuccess.firstCmdSnap.task
+
+      return {
+        ictx, stx
+        diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
+        result? := some {
+          parserState
+          processedSnap := (← processHeader stx parserState)
+        }
+      }
+
+  processHeader (stx : Syntax) (parserState : Parser.ModuleParserState) :
+      LeanProcessingM (SnapshotTask HeaderProcessedSnapshot) := do
+    let ctx ← read
+    SnapshotTask.ofIO ⟨0, ctx.input.endPos⟩ <|
+    ReaderT.run (r := ctx) <|  -- re-enter reader in new task
+    withHeaderExceptions (α := HeaderProcessedSnapshot) ({ · with result? := none }) do
+      let opts ← match (← setupImports stx) with
+        | .ok opts => pure opts
+        | .error snap => return snap
+      -- override context options with file options
+      let opts := ctx.opts.mergeBy (fun _ _ fileOpt => fileOpt) opts
+      -- allows `headerEnv` to be leaked, which would live until the end of the process anyway
+      let (headerEnv, msgLog) ← Elab.processHeader (leakEnv := true) stx opts .empty
+        ctx.toInputContext ctx.trustLevel
+      let diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
+      if msgLog.hasErrors then
+        return { diagnostics, result? := none }
+
+      let headerEnv := headerEnv.setMainModule ctx.mainModuleName
+      let cmdState := Elab.Command.mkState headerEnv msgLog opts
+      let cmdState := { cmdState with infoState := {
+        enabled := true
+        trees := #[Elab.InfoTree.context (.commandCtx {
+          env     := headerEnv
+          fileMap := ctx.fileMap
+          ngen    := { namePrefix := `_import }
+        }) (Elab.InfoTree.node
+            (Elab.Info.ofCommandInfo { elaborator := `header, stx })
+            (stx[1].getArgs.toList.map (fun importStx =>
+              Elab.InfoTree.node (Elab.Info.ofCommandInfo {
+                elaborator := `import
+                stx := importStx
+              }) #[].toPArray'
+            )).toPArray'
+        )].toPArray'
+      }}
+      return {
+        diagnostics
+        infoTree? := cmdState.infoState.trees[0]!
+        result? := some {
+          cmdState
+          firstCmdSnap := (← parseCmd none parserState cmdState)
+        }
+      }
+
+  parseCmd (old? : Option CommandParsedSnapshot) (parserState : Parser.ModuleParserState)
+      (cmdState : Command.State) : LeanProcessingM (SnapshotTask CommandParsedSnapshot) := do
+    let ctx ← read
+
+    -- check for cancellation, most likely during elaboration of previous command, before starting
+    -- processing of next command
+    if (← IO.checkCanceled) then
+      -- this is a bit ugly as we don't want to adjust our API with `Option`s just for cancellation
+      -- (as no-one should look at this result in that case) but anything containing `Environment`
+      -- is not `Inhabited`
+      return .pure <| .mk (nextCmdSnap? := none) {
+        diagnostics := .empty, stx := .missing, parserState
+        sigSnap := .pure {
+          diagnostics := .empty
+          finishedSnap := .pure { diagnostics := .empty, cmdState } } }
+
+    let unchanged old : BaseIO CommandParsedSnapshot :=
+      -- when syntax is unchanged, reuse command processing task as is
+      if let some oldNext := old.next? then
+        return .mk (data := old.data)
+          (nextCmdSnap? := (← old.data.sigSnap.bindIO (sync := true) fun oldSig =>
+            oldSig.finishedSnap.bindIO (sync := true) fun oldFinished =>
+              -- also wait on old command parse snapshot as parsing is cheap and may allow for
+              -- elaboration reuse
+              oldNext.bindIO (sync := true) fun oldNext => do
+                parseCmd oldNext old.data.parserState oldFinished.cmdState ctx))
+      else return old  -- terminal command, we're done!
+
+    -- fast path, do not even start new task for this snapshot
+    if let some old := old? then
+      if let some nextCom ← old.next?.bindM (·.get?) then
+        if (← isBeforeEditPos nextCom.data.parserState.pos) then
+          return .pure (← unchanged old)
+
+    SnapshotTask.ofIO ⟨parserState.pos, ctx.input.endPos⟩ do
+      let beginPos := parserState.pos
+      let scope := cmdState.scopes.head!
+      let pmctx := {
+        env := cmdState.env, options := scope.opts, currNamespace := scope.currNamespace
+        openDecls := scope.openDecls
+      }
+      let (stx, parserState, msgLog) := Parser.parseCommand ctx.toInputContext pmctx parserState
+        .empty
+
+      -- semi-fast path
+      if let some old := old? then
+        if (← isBeforeEditPos parserState.pos ctx) && old.data.stx == stx then
+          return (← unchanged old)
+        -- on first change, make sure to cancel all further old tasks
+        old.cancel
+
+      let sigSnap ← processCmdSignature stx cmdState msgLog.hasErrors beginPos ctx
+      let next? ← if Parser.isTerminalCommand stx then pure none
+        -- for now, wait on "command finished" snapshot before parsing next command
+        else some <$> (sigSnap.bind (·.finishedSnap)).bindIO fun finished =>
+          parseCmd none parserState finished.cmdState ctx
+      return .mk (nextCmdSnap? := next?) {
+        diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog ctx.toProcessingContext)
+        stx
+        parserState
+        sigSnap
+      }
+
+  processCmdSignature (stx : Syntax) (cmdState : Command.State) (hasParseError : Bool)
+      (beginPos : String.Pos) :
+      LeanProcessingM (SnapshotTask CommandSignatureProcessedSnapshot) := do
+    let ctx ← read
+
+    -- signature elaboration task; for now, does full elaboration
+    -- TODO: do tactic snapshots, reuse old state for them
+    SnapshotTask.ofIO (stx.getRange?.getD ⟨beginPos, beginPos⟩) do
+      let scope := cmdState.scopes.head!
+      let cmdStateRef ← IO.mkRef { cmdState with messages := .empty }
+      let cmdCtx : Elab.Command.Context := { ctx with cmdPos := beginPos, tacticCache? := none }
+      let (output, _) ←
+        IO.FS.withIsolatedStreams (isolateStderr := stderrAsMessages.get scope.opts) do
+          liftM (m := BaseIO) do
+            Elab.Command.catchExceptions
+              (getResetInfoTrees *> Elab.Command.elabCommandTopLevel stx)
+              cmdCtx cmdStateRef
+      let cmdState ← cmdStateRef.get
+      let mut messages := cmdState.messages
+      -- `stx.hasMissing` should imply `hasParseError`, but the latter should be cheaper to check in
+      -- general
+      if !showPartialSyntaxErrors.get cmdState.scopes[0]!.opts && hasParseError &&
+          stx.hasMissing then
+        -- discard elaboration errors, except for a few important and unlikely misleading ones, on
+        -- parse error
+        messages := ⟨messages.msgs.filter fun msg =>
+          msg.data.hasTag (fun tag => tag == `Elab.synthPlaceholder ||
+            tag == `Tactic.unsolvedGoals || (`_traceMsg).isSuffixOf tag)⟩
+      if !output.isEmpty then
+        messages := messages.add {
+          fileName := ctx.fileName
+          severity := MessageSeverity.information
+          pos      := ctx.fileMap.toPosition beginPos
+          data     := output
+        }
+      let cmdState := { cmdState with messages }
+      return {
+        diagnostics := .empty
+        finishedSnap := .pure {
+          diagnostics :=
+            (← Snapshot.Diagnostics.ofMessageLog cmdState.messages ctx.toProcessingContext)
+          infoTree? := some cmdState.infoState.trees[0]!
+          cmdState
+        }
+      }
+
+/-- Waits for and returns final environment, if importing was successful. -/
+partial def waitForFinalEnv? (snap : InitialSnapshot) : Option Environment := do
+  let snap ← snap.result?
+  let snap ← snap.processedSnap.get.result?
+  goCmd snap.firstCmdSnap.get
+where goCmd snap :=
+  if let some next := snap.next? then
+    goCmd next.get
+  else
+    snap.data.sigSnap.get.finishedSnap.get.cmdState.env
+
+end Lean

src/Lean/Message.lean

@@ -270,8 +270,13 @@ def getInfoMessages (log : MessageLog) : MessageLog :=
 def forM {m : Type → Type} [Monad m] (log : MessageLog) (f : Message → m Unit) : m Unit :=
   log.msgs.forM f
 
+/-- Converts the log to a list, oldest message first. -/
 def toList (log : MessageLog) : List Message :=
-  (log.msgs.foldl (fun acc msg => msg :: acc) []).reverse
+  log.msgs.toList
+
+/-- Converts the log to an array, oldest message first. -/
+def toArray (log : MessageLog) : Array Message :=
+  log.msgs.toArray
 
 end MessageLog
 

src/Lean/Server/Completion.lean

@@ -81,39 +81,46 @@ open Elab
 open Meta
 open FuzzyMatching
 
-/-- Cached header declarations for which `allowCompletion headerEnv decl` is true. -/
-builtin_initialize eligibleHeaderDeclsRef : IO.Ref (HashMap Name ConstantInfo) ← IO.mkRef {}
+abbrev EligibleHeaderDecls := HashMap Name ConstantInfo
 
-/-- Caches the declarations in the header for which `allowCompletion headerEnv decl` is true. -/
-def fillEligibleHeaderDecls (headerEnv : Environment) : IO Unit := do
-  eligibleHeaderDeclsRef.modify fun startEligibleHeaderDecls =>
-    (·.2) <| StateT.run (m := Id) (s := startEligibleHeaderDecls) do
-      headerEnv.constants.forM fun declName c => do
-        modify fun eligibleHeaderDecls =>
-          if allowCompletion headerEnv declName then
-            eligibleHeaderDecls.insert declName c
-          else
-            eligibleHeaderDecls
+/-- Cached header declarations for which `allowCompletion headerEnv decl` is true. -/
+builtin_initialize eligibleHeaderDeclsRef : IO.Ref (Option EligibleHeaderDecls) ←
+  IO.mkRef none
+
+/--
+Returns the declarations in the header for which `allowCompletion env decl` is true, caching them
+if not already cached.
+-/
+def getEligibleHeaderDecls (env : Environment) : IO EligibleHeaderDecls := do
+  eligibleHeaderDeclsRef.modifyGet fun
+    | some eligibleHeaderDecls => (eligibleHeaderDecls, some eligibleHeaderDecls)
+    | none =>
+      let (_, eligibleHeaderDecls) :=
+        StateT.run (m := Id) (s := {}) do
+          -- `map₁` are the header decls
+          env.constants.map₁.forM fun declName c => do
+            modify fun eligibleHeaderDecls =>
+              if allowCompletion env declName then
+                eligibleHeaderDecls.insert declName c
+              else
+                eligibleHeaderDecls
+      (eligibleHeaderDecls, some eligibleHeaderDecls)
 
 /-- Iterate over all declarations that are allowed in completion results. -/
 private def forEligibleDeclsM [Monad m] [MonadEnv m] [MonadLiftT (ST IO.RealWorld) m]
-    (f : Name → ConstantInfo → m PUnit) : m PUnit := do
+    [MonadLiftT IO m] (f : Name → ConstantInfo → m PUnit) : m PUnit := do
   let env ← getEnv
-  (← eligibleHeaderDeclsRef.get).forM f
+  (← getEligibleHeaderDecls env).forM f
   -- map₂ are exactly the local decls
   env.constants.map₂.forM fun name c => do
     if allowCompletion env name then
       f name c
 
 /-- Checks whether this declaration can appear in completion results. -/
-private def allowCompletion [Monad m] [MonadEnv m] [MonadLiftT (ST IO.RealWorld) m]
-    (declName : Name) : m Bool := do
-  let env ← getEnv
-  if (← eligibleHeaderDeclsRef.get).contains declName then
-    return true
-  if env.constants.map₂.contains declName && Lean.Meta.allowCompletion env declName then
-    return true
-  return false
+private def allowCompletion (eligibleHeaderDecls : EligibleHeaderDecls) (env : Environment)
+    (declName : Name) : Bool :=
+  eligibleHeaderDecls.contains declName ||
+    env.constants.map₂.contains declName && Lean.Meta.allowCompletion env declName
 
 /--
 Sorts `items` descendingly according to their score and ascendingly according to their label
@@ -371,16 +378,17 @@ private def idCompletionCore
       matchAtomic id (alias.replacePrefix ns Name.anonymous)
     else
       none
+  let eligibleHeaderDecls ← getEligibleHeaderDecls env
   -- Auxiliary function for `alias`
   let addAlias (alias : Name) (declNames : List Name) (score : Float) : M Unit := do
     declNames.forM fun declName => do
-      if ← allowCompletion declName then
+      if allowCompletion eligibleHeaderDecls env declName then
         addUnresolvedCompletionItemForDecl alias.getString! declName score
   -- search explicitly open `ids`
   for openDecl in ctx.openDecls do
     match openDecl with
     | OpenDecl.explicit openedId resolvedId =>
-      if ← allowCompletion resolvedId then
+      if allowCompletion eligibleHeaderDecls env resolvedId then
         if let some score := matchAtomic id openedId then
           addUnresolvedCompletionItemForDecl openedId.getString! resolvedId score
     | OpenDecl.simple ns _      =>

src/Lean/Server/FileWorker.lean

@@ -6,6 +6,7 @@ Authors: Marc Huisinga, Wojciech Nawrocki
 -/
 prelude
 import Init.System.IO
+import Init.Data.Channel
 
 import Lean.Data.RBMap
 import Lean.Environment
@@ -15,9 +16,9 @@ import Lean.Data.Json.FromToJson
 
 import Lean.Util.FileSetupInfo
 import Lean.LoadDynlib
+import Lean.Language.Basic
 
 import Lean.Server.Utils
-import Lean.Server.Snapshots
 import Lean.Server.AsyncList
 import Lean.Server.References
 
@@ -59,90 +60,152 @@ open Snapshots
 open JsonRpc
 
 structure WorkerContext where
-  hIn              : FS.Stream
-  hOut             : FS.Stream
+  /-- Synchronized output channel for LSP messages. Notifications for outdated versions are
+    discarded on read. -/
+  chanOut          : IO.Channel JsonRpc.Message
+  /--
+  Latest document version received by the client, used for filtering out notifications from
+  previous versions.
+  -/
+  maxDocVersionRef : IO.Ref Int
   hLog             : FS.Stream
-  headerTask       : Task (Except Error (Snapshot × SearchPath))
   initParams       : InitializeParams
+  processor        : Parser.InputContext → BaseIO Lean.Language.Lean.InitialSnapshot
   clientHasWidgets : Bool
+  /--
+  Options defined on the worker cmdline (i.e. not including options from `setup-file`), used for
+  context-free tasks such as editing delay.
+  -/
+  cmdlineOpts      : Options
 
 /-! # Asynchronous snapshot elaboration -/
 
 section Elab
-  structure AsyncElabState where
-    snaps : Array Snapshot
-
-  abbrev AsyncElabM := StateT AsyncElabState <| EIO ElabTaskError
-
   -- Placed here instead of Lean.Server.Utils because of an import loop
-  private def publishIleanInfo (method : String) (m : DocumentMeta) (hOut : FS.Stream)
-      (snaps : Array Snapshot) : IO Unit := do
-    let trees := snaps.map fun snap => snap.infoTree
+  private def mkIleanInfoNotification (method : String) (m : DocumentMeta)
+      (trees : Array Elab.InfoTree) : BaseIO (JsonRpc.Notification Lsp.LeanIleanInfoParams) := do
     let references ← findModuleRefs m.text trees (localVars := true) |>.toLspModuleRefs
-    let param := { version := m.version, references : LeanIleanInfoParams }
-    hOut.writeLspNotification { method, param }
+    let param := { version := m.version, references }
+    return { method, param }
 
-  private def publishIleanInfoUpdate : DocumentMeta → FS.Stream → Array Snapshot → IO Unit :=
-    publishIleanInfo "$/lean/ileanInfoUpdate"
+  private def mkIleanInfoUpdateNotification : DocumentMeta → Array Elab.InfoTree →
+      BaseIO (JsonRpc.Notification Lsp.LeanIleanInfoParams) :=
+    mkIleanInfoNotification "$/lean/ileanInfoUpdate"
 
-  private def publishIleanInfoFinal : DocumentMeta → FS.Stream → Array Snapshot → IO Unit :=
-    publishIleanInfo "$/lean/ileanInfoFinal"
+  private def mkIleanInfoFinalNotification : DocumentMeta → Array Elab.InfoTree →
+      BaseIO (JsonRpc.Notification Lsp.LeanIleanInfoParams) :=
+    mkIleanInfoNotification "$/lean/ileanInfoFinal"
 
-  /-- Elaborates the next command after `parentSnap` and emits diagnostics into `hOut`. -/
-  private def nextCmdSnap (ctx : WorkerContext) (m : DocumentMeta) (cancelTk : CancelToken)
-      : AsyncElabM (Option Snapshot) := do
-    cancelTk.check
-    let s ← get
-    let .some lastSnap := s.snaps.back? | panic! "empty snapshots"
-    if lastSnap.isAtEnd then
-      publishDiagnostics m lastSnap.diagnostics.toArray ctx.hOut
-      publishProgressDone m ctx.hOut
-      -- This will overwrite existing ilean info for the file, in case something
-      -- went wrong during the incremental updates.
-      publishIleanInfoFinal m ctx.hOut s.snaps
-      return none
-    publishProgressAtPos m lastSnap.endPos ctx.hOut
-    let snap ← compileNextCmd m.mkInputContext lastSnap ctx.clientHasWidgets
-    set { s with snaps := s.snaps.push snap }
-    -- TODO(MH): check for interrupt with increased precision
-    cancelTk.check
-    /- NOTE(MH): This relies on the client discarding old diagnostics upon receiving new ones
-      while preferring newer versions over old ones. The former is necessary because we do
-      not explicitly clear older diagnostics, while the latter is necessary because we do
-      not guarantee that diagnostics are emitted in order. Specifically, it may happen that
-      we interrupted this elaboration task right at this point and a newer elaboration task
-      emits diagnostics, after which we emit old diagnostics because we did not yet detect
-      the interrupt. Explicitly clearing diagnostics is difficult for a similar reason,
-      because we cannot guarantee that no further diagnostics are emitted after clearing
-      them. -/
-    -- NOTE(WN): this is *not* redundant even if there are no new diagnostics in this snapshot
-    -- because empty diagnostics clear existing error/information squiggles. Therefore we always
-    -- want to publish in case there was previously a message at this position.
-    publishDiagnostics m snap.diagnostics.toArray ctx.hOut
-    publishIleanInfoUpdate m ctx.hOut #[snap]
-    return some snap
+  /-- State of `reportSnapshots`. -/
+  private structure ReportSnapshotsState where
+    /-- Whether we have waited for a snapshot to finish at least once (see debouncing below). -/
+    hasBlocked := false
+    /-- All info trees encountered so far. -/
+    allInfoTrees : Array Elab.InfoTree := #[]
+    /-- New info trees encountered since we last sent a .ilean update notification. -/
+    newInfoTrees : Array Elab.InfoTree := #[]
+    /-- Whether we encountered any snapshot with `Snapshot.isFatal`. -/
+    hasFatal := false
+  deriving Inhabited
 
-  /-- Elaborates all commands after the last snap (at least the header snap is assumed to exist), emitting the diagnostics into `hOut`. -/
-  def unfoldCmdSnaps (m : DocumentMeta) (snaps : Array Snapshot) (cancelTk : CancelToken) (startAfterMs : UInt32)
-      : ReaderT WorkerContext IO (AsyncList ElabTaskError Snapshot) := do
-    let ctx ← read
-    let some headerSnap := snaps[0]? | panic! "empty snapshots"
-    if headerSnap.msgLog.hasErrors then
-      -- Treats header processing errors as fatal so users aren't swamped with
-      -- followup errors
-      publishProgressAtPos m headerSnap.beginPos ctx.hOut (kind := LeanFileProgressKind.fatalError)
-      publishIleanInfoFinal m ctx.hOut #[headerSnap]
-      return AsyncList.ofList [headerSnap]
-    else
-      -- This will overwrite existing ilean info for the file since this has a
-      -- higher version number.
-      publishIleanInfoUpdate m ctx.hOut snaps
-      return AsyncList.ofList snaps.toList ++ AsyncList.delayed (← EIO.asTask (ε := ElabTaskError) (prio := .dedicated) do
-        IO.sleep startAfterMs
-        AsyncList.unfoldAsync (nextCmdSnap ctx m cancelTk) { snaps })
+  register_builtin_option server.reportDelayMs : Nat := {
+    defValue := 200
+    group := "server"
+    descr := "(server) time in milliseconds to wait before reporting progress and diagnostics on \
+      document edit in order to reduce flickering
+
+This option can only be set on the command line, not in the lakefile or via `set_option`."
+  }
+
+  /--
+  Type of state stored in `Snapshot.Diagnostics.cacheRef?`.
+
+  See also section "Communication" in Lean/Server/README.md.
+  -/
+  structure MemorizedInteractiveDiagnostics where
+    diags : Array Widget.InteractiveDiagnostic
+  deriving TypeName
+
+  open Language in
+  /--
+    Reports status of a snapshot tree incrementally to the user: progress,
+    diagnostics, .ilean reference information.
+
+    See also section "Communication" in Lean/Server/README.md.
+
+    Debouncing: we only report information
+    * after first waiting for `reportDelayMs`, to give trivial tasks a chance to finish
+    * when first blocking, i.e. not before skipping over any unchanged snapshots and such trival
+      tasks
+    * afterwards, each time new information is found in a snapshot
+    * at the very end, if we never blocked (e.g. emptying a file should make
+      sure to empty diagnostics as well eventually) -/
+  private partial def reportSnapshots (ctx : WorkerContext) (doc : EditableDocumentCore)
+      (cancelTk : CancelToken) : BaseIO (Task Unit) := do
+    let t ← BaseIO.asTask do
+      IO.sleep (server.reportDelayMs.get ctx.cmdlineOpts).toUInt32
+    BaseIO.bindTask t fun _ => do
+      BaseIO.bindTask (← go (toSnapshotTree doc.initSnap) {}) fun st => do
+        if (← cancelTk.isSet) then
+          return .pure ()
+
+        -- callback at the end of reporting
+        if st.hasFatal then
+          ctx.chanOut.send <| mkFileProgressAtPosNotification doc.meta 0 .fatalError
+        else
+          ctx.chanOut.send <| mkFileProgressDoneNotification doc.meta
+        unless st.hasBlocked do
+          publishDiagnostics
+        -- This will overwrite existing ilean info for the file, in case something
+        -- went wrong during the incremental updates.
+        ctx.chanOut.send (← mkIleanInfoFinalNotification doc.meta st.allInfoTrees)
+        return .pure ()
+  where
+    publishDiagnostics := do
+      ctx.chanOut.send <| mkPublishDiagnosticsNotification doc.meta <|
+        (← doc.diagnosticsRef.get).map (·.toDiagnostic)
+    go (node : SnapshotTree) (st : ReportSnapshotsState) : BaseIO (Task ReportSnapshotsState) := do
+      if !node.element.diagnostics.msgLog.isEmpty then
+        let diags ←
+          if let some memorized ← node.element.diagnostics.interactiveDiagsRef?.bindM fun ref => do
+              return (← ref.get).bind (·.get? MemorizedInteractiveDiagnostics) then
+            pure memorized.diags
+          else
+            let diags ← node.element.diagnostics.msgLog.toArray.mapM
+              (Widget.msgToInteractiveDiagnostic doc.meta.text · ctx.clientHasWidgets)
+            if let some cacheRef := node.element.diagnostics.interactiveDiagsRef? then
+              cacheRef.set <| some <| .mk { diags : MemorizedInteractiveDiagnostics }
+            pure diags
+        doc.diagnosticsRef.modify (· ++ diags)
+        if st.hasBlocked then
+          publishDiagnostics
+
+      let mut st := { st with hasFatal := st.hasFatal || node.element.isFatal }
+
+      if let some itree := node.element.infoTree? then
+        let mut newInfoTrees := st.newInfoTrees.push itree
+        if st.hasBlocked then
+          ctx.chanOut.send (← mkIleanInfoUpdateNotification doc.meta newInfoTrees)
+          newInfoTrees := #[]
+        st := { st with newInfoTrees, allInfoTrees := st.allInfoTrees.push itree }
+
+      goSeq st node.children.toList
+
+    goSeq (st : ReportSnapshotsState) :
+        List (SnapshotTask SnapshotTree) → BaseIO (Task ReportSnapshotsState)
+      | [] => return .pure st
+      | t::ts => do
+        let mut st := st
+        unless (← IO.hasFinished t.task) do
+          ctx.chanOut.send <| mkFileProgressAtPosNotification doc.meta t.range.start
+          if !st.hasBlocked then
+            publishDiagnostics
+            st := { st with hasBlocked := true }
+        BaseIO.bindTask t.task fun t => do
+          BaseIO.bindTask (← go t st) (goSeq · ts)
 end Elab
 
--- Pending requests are tracked so they can be cancelled
+-- Pending requests are tracked so they can be canceled
 abbrev PendingRequestMap := RBMap RequestID (Task (Except IO.Error Unit)) compare
 
 structure AvailableImportsCache where
@@ -151,11 +214,9 @@ structure AvailableImportsCache where
 
 structure WorkerState where
   doc                : EditableDocument
-  -- The initial header syntax tree that the file worker was started with.
-  initHeaderStx      : Syntax
-  -- The current header syntax tree. Changing the header from `initHeaderStx` initiates a restart
-  -- that only completes after a while, so `currHeaderStx` tracks the modified syntax until then.
-  currHeaderStx      : Syntax
+  /-- Token flagged for aborting `doc.reporter` when a new document version comes in. -/
+  reporterCancelTk   : CancelToken
+  srcSearchPathTask  : Task SearchPath
   importCachingTask? : Option (Task (Except Error AvailableImportsCache))
   pendingRequests    : PendingRequestMap
   /-- A map of RPC session IDs. We allow asynchronous elab tasks and request handlers
@@ -164,180 +225,142 @@ structure WorkerState where
 
 abbrev WorkerM := ReaderT WorkerContext <| StateRefT WorkerState IO
 
+/-- Makes sure we load imports at most once per process as they cannot be unloaded. -/
+private builtin_initialize importsLoadedRef : IO.Ref Bool ← IO.mkRef false
+
+open Language Lean in
+/--
+Callback from Lean language processor after parsing imports that requests necessary information from
+Lake for processing imports.
+-/
+def setupImports (meta : DocumentMeta) (chanOut : Channel JsonRpc.Message)
+    (srcSearchPathPromise : Promise SearchPath) (stx : Syntax) :
+    Language.ProcessingT IO (Except Language.Lean.HeaderProcessedSnapshot Options) := do
+  let importsAlreadyLoaded ← importsLoadedRef.modifyGet ((·, true))
+  if importsAlreadyLoaded then
+    -- As we never unload imports in the server, we should not run the code below twice in the
+    -- same process and instead ask the watchdog to restart the worker
+    IO.sleep 200  -- give user time to make further edits before restart
+    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 }
+
+  let imports := Elab.headerToImports stx
+  let fileSetupResult ← setupFile meta imports fun stderrLine => do
+    let progressDiagnostic := {
+      range      := ⟨⟨0, 0⟩, ⟨0, 0⟩⟩
+      -- make progress visible anywhere in the file
+      fullRange? := some ⟨⟨0, 0⟩, meta.text.utf8PosToLspPos meta.text.source.endPos⟩
+      severity?  := DiagnosticSeverity.information
+      message    := stderrLine
+    }
+    chanOut.send <| mkPublishDiagnosticsNotification meta #[progressDiagnostic]
+  -- clear progress notifications in the end
+  chanOut.send <| mkPublishDiagnosticsNotification meta #[]
+  match fileSetupResult.kind with
+  | .importsOutOfDate =>
+    return .error {
+      diagnostics := (← Language.diagnosticsOfHeaderError
+        "Imports are out of date and must be rebuilt; \
+          use the \"Restart File\" command in your editor.")
+      result? := none
+    }
+  | .error msg =>
+    return .error {
+      diagnostics := (← diagnosticsOfHeaderError msg)
+      result? := none
+    }
+  | _ => pure ()
+
+  srcSearchPathPromise.resolve fileSetupResult.srcSearchPath
+  return .ok fileSetupResult.fileOptions
+
 /- Worker initialization sequence. -/
 section Initialization
-  def buildHeaderEnv (m : DocumentMeta) (headerStx : Syntax) (fileSetupResult : FileSetupResult) : IO (Environment × MessageLog) := do
-    let (headerEnv, msgLog) ←
-      match fileSetupResult.kind with
-      | .success | .noLakefile =>
-        -- allows `headerEnv` to be leaked, which would live until the end of the process anyway
-        Elab.processHeader (leakEnv := true) headerStx fileSetupResult.fileOptions MessageLog.empty m.mkInputContext
-      | .importsOutOfDate =>
-        mkErrorEnvironment "Imports are out of date and must be rebuilt; use the \"Restart File\" command in your editor."
-      | .error msg =>
-        mkErrorEnvironment msg
-
-    let mut headerEnv := headerEnv
-    try
-      if let some path := System.Uri.fileUriToPath? m.uri then
-        headerEnv := headerEnv.setMainModule (← moduleNameOfFileName path none)
-    catch _ =>
-      pure ()
-
-    return (headerEnv, msgLog)
-
-  where
-    mkErrorEnvironment (errorMsg : String) : IO (Environment × MessageLog) := do
-      let msgs := MessageLog.empty.add { fileName := "<ignored>", pos := ⟨0, 0⟩, data := errorMsg }
-      return (← mkEmptyEnvironment, msgs)
-
-  def buildCommandState
-      (m : DocumentMeta)
-      (headerStx : Syntax)
-      (headerEnv : Environment)
-      (headerMsgLog : MessageLog)
-      (opts : Options)
-      : Elab.Command.State :=
-    let headerContextInfo : Elab.CommandContextInfo := {
-      env     := headerEnv
-      fileMap := m.text
-      ngen    := { namePrefix := `_worker }
-    }
-    let headerInfo := Elab.Info.ofCommandInfo { elaborator := `header, stx := headerStx }
-    let headerInfoNodes := headerStx[1].getArgs.toList.map fun importStx =>
-      Elab.InfoTree.node (Elab.Info.ofCommandInfo {
-        elaborator := `import
-        stx := importStx
-      }) #[].toPArray'
-    let headerInfoTree := Elab.InfoTree.node headerInfo headerInfoNodes.toPArray'
-    let headerInfoState := {
-      enabled := true
-      trees := #[Elab.InfoTree.context (.commandCtx headerContextInfo) headerInfoTree].toPArray'
-    }
-    { Elab.Command.mkState headerEnv headerMsgLog opts with infoState := headerInfoState }
-
-  def compileHeader (m : DocumentMeta) (hOut : FS.Stream) (globalOptions : Options) (hasWidgets : Bool)
-      : IO (Syntax × Task (Except Error (Snapshot × SearchPath))) := do
-    -- parsing should not take long, do synchronously
-    let (headerStx, headerParserState, parseMsgLog) ← Parser.parseHeader m.mkInputContext
-    (headerStx, ·) <$> EIO.asTask do
-      let imports := Lean.Elab.headerToImports headerStx
-      let fileSetupResult ← setupFile m imports fun stderrLine =>
-        let progressDiagnostic := {
-          range     := ⟨⟨0, 0⟩, ⟨0, 0⟩⟩
-          severity? := DiagnosticSeverity.information
-          message   := stderrLine
-        }
-        publishDiagnostics m #[progressDiagnostic] hOut
-      let fileSetupResult := fileSetupResult.addGlobalOptions globalOptions
-      let (headerEnv, envMsgLog) ← buildHeaderEnv m headerStx fileSetupResult
-      -- Prepare header-based caches that requests may use
-      runHeaderCachingHandlers headerEnv
-      let headerMsgLog := parseMsgLog.append envMsgLog
-      let cmdState := buildCommandState m headerStx headerEnv headerMsgLog fileSetupResult.fileOptions
-      let headerSnap := {
-        beginPos         := 0
-        stx              := headerStx
-        mpState          := headerParserState
-        cmdState         := cmdState
-        interactiveDiags := ← cmdState.messages.msgs.mapM (Widget.msgToInteractiveDiagnostic m.text · hasWidgets)
-        tacticCache      := (← IO.mkRef {})
-      }
-      publishDiagnostics m headerSnap.diagnostics.toArray hOut
-      return (headerSnap, fileSetupResult.srcSearchPath)
-
-  def initializeWorker (meta : DocumentMeta) (i o e : FS.Stream) (initParams : InitializeParams) (opts : Options)
+  def initializeWorker (meta : DocumentMeta) (o e : FS.Stream) (initParams : InitializeParams) (opts : Options)
       : IO (WorkerContext × WorkerState) := do
     let clientHasWidgets := initParams.initializationOptions?.bind (·.hasWidgets?) |>.getD false
-    let (headerStx, headerTask) ← compileHeader meta o opts (hasWidgets := clientHasWidgets)
-    let cancelTk ← CancelToken.new
+    let mut mainModuleName := Name.anonymous
+    try
+      if let some path := System.Uri.fileUriToPath? meta.uri then
+        mainModuleName ← moduleNameOfFileName path none
+    catch _ => pure ()
+    let maxDocVersionRef ← IO.mkRef 0
+    let chanOut ← mkLspOutputChannel maxDocVersionRef
+    let srcSearchPathPromise ← IO.Promise.new
+
+    let processor := Language.Lean.process (setupImports meta chanOut srcSearchPathPromise)
+    let processor ← Language.mkIncrementalProcessor processor { opts, mainModuleName }
+    let initSnap ← processor meta.mkInputContext
     let ctx := {
-      hIn  := i
-      hOut := o
+      chanOut
       hLog := e
-      headerTask
       initParams
+      processor
       clientHasWidgets
+      maxDocVersionRef
+      cmdlineOpts := opts
     }
-    let cmdSnaps ← EIO.mapTask (t := headerTask) (match · with
-      | Except.ok (s, _) => unfoldCmdSnaps meta #[s] cancelTk ctx (startAfterMs := 0)
-      | Except.error e   => throw (e : ElabTaskError))
-    let doc : EditableDocument := { meta, cmdSnaps := AsyncList.delayed cmdSnaps, cancelTk }
+    let doc : EditableDocumentCore := {
+      meta, initSnap
+      diagnosticsRef := (← IO.mkRef ∅)
+    }
+    let reporterCancelTk ← CancelToken.new
+    let reporter ← reportSnapshots ctx doc reporterCancelTk
     return (ctx, {
-      doc                := doc
-      initHeaderStx      := headerStx
-      currHeaderStx      := headerStx
-      importCachingTask? := none
+      doc := { doc with reporter }
+      reporterCancelTk
+      srcSearchPathTask  := srcSearchPathPromise.result
       pendingRequests    := RBMap.empty
       rpcSessions        := RBMap.empty
+      importCachingTask? := none
     })
+  where
+    /-- Creates an LSP message output channel along with a reader that sends out read messages on
+        the output FS stream after discarding outdated notifications. This is the only component of
+        the worker with access to the output stream, so we can synchronize messages from parallel
+        elaboration tasks here. -/
+    mkLspOutputChannel maxDocVersion : IO (IO.Channel JsonRpc.Message) := do
+      let chanOut ← IO.Channel.new
+      let _ ← chanOut.forAsync (prio := .dedicated) fun msg => do
+        -- discard outdated notifications; note that in contrast to responses, notifications can
+        -- always be silently discarded
+        let version? : Option Int := do match msg with
+          | .notification "textDocument/publishDiagnostics" (some params) =>
+            let params : PublishDiagnosticsParams ← fromJson? (toJson params) |>.toOption
+            params.version?
+          | .notification "$/lean/fileProgress" (some params) =>
+            let params : LeanFileProgressParams ← fromJson? (toJson params) |>.toOption
+            params.textDocument.version?
+          | _ => none
+        if let some version := version? then
+          if version < (← maxDocVersion.get) then
+            return
+          -- note that because of `server.reportDelayMs`, we cannot simply set `maxDocVersion` here
+          -- as that would allow outdated messages to be reported until the delay is over
+        o.writeLspMessage msg |>.catchExceptions (fun _ => pure ())
+      return chanOut
 end Initialization
 
 section Updates
   def updatePendingRequests (map : PendingRequestMap → PendingRequestMap) : WorkerM Unit := do
     modify fun st => { st with pendingRequests := map st.pendingRequests }
 
-  def determineValidSnapshots (oldDoc : EditableDocument) (newMeta : DocumentMeta) (newHeaderSnap : Snapshot) : IO (List Snapshot) := do
-    let changePos := oldDoc.meta.text.source.firstDiffPos newMeta.text.source
-    -- Ignores exceptions, we are only interested in the successful snapshots
-    let (cmdSnaps, _) ← oldDoc.cmdSnaps.getFinishedPrefix
-    oldDoc.cmdSnaps.cancel
-    -- NOTE(WN): we invalidate eagerly as `endPos` consumes input greedily. To re-elaborate only
-    -- when really necessary, we could do a whitespace-aware `Syntax` comparison instead.
-    let mut validSnaps ← pure (cmdSnaps.takeWhile (fun s => s.endPos < changePos))
-    if h : validSnaps.length ≤ 1 then
-      validSnaps := [newHeaderSnap]
-    else
-      /- When at least one valid non-header snap exists, it may happen that a change does not fall
-          within the syntactic range of that last snap but still modifies it by appending tokens.
-          We check for this here. We do not currently handle crazy grammars in which an appended
-          token can merge two or more previous commands into one. To do so would require reparsing
-          the entire file. -/
-      have : validSnaps.length ≥ 2 := Nat.gt_of_not_le h
-      let mut lastSnap := validSnaps.getLast (by subst ·; simp at h)
-      let preLastSnap :=
-        have : 0 < validSnaps.length := Nat.lt_of_lt_of_le (by decide) this
-        have : validSnaps.length - 2 < validSnaps.length := Nat.sub_lt this (by decide)
-        validSnaps[validSnaps.length - 2]
-      let newLastStx ← parseNextCmd newMeta.mkInputContext preLastSnap
-      if newLastStx != lastSnap.stx then
-        validSnaps := validSnaps.dropLast
-    return validSnaps
-
-  def startNewSnapshotTasks (newMeta : DocumentMeta) : WorkerM (AsyncList ElabTaskError Snapshot × CancelToken) := do
-    let ctx ← read
-    let oldDoc := (← get).doc
-    oldDoc.cancelTk.set
-    let initHeaderStx := (← get).initHeaderStx
-    let (newHeaderStx, newMpState, _) ← Parser.parseHeader newMeta.mkInputContext
-    let cancelTk ← CancelToken.new
-
-    if initHeaderStx != newHeaderStx then
-      set { ← get with currHeaderStx := newHeaderStx }
-      let terminationTask ← EIO.asTask (ε := ElabTaskError) (prio := .dedicated) do
-        IO.sleep ctx.initParams.editDelay.toUInt32
-        cancelTk.check
-        IO.Process.exit 2
-      return (AsyncList.delayed terminationTask, cancelTk)
-
-    let headSnapTask := oldDoc.cmdSnaps.waitHead?
-    let newSnapTasks ← EIO.mapTask (ε := ElabTaskError) (t := headSnapTask) (prio := .dedicated) fun headSnap?? => do
-      -- There is always at least one snapshot absent exceptions
-      let some headSnap ← MonadExcept.ofExcept headSnap?? | panic! "empty snapshots"
-      let newHeaderSnap := { headSnap with stx := newHeaderStx, mpState := newMpState }
-      let validSnaps ← determineValidSnapshots oldDoc newMeta newHeaderSnap
-      -- wait for a bit, giving the initial `cancelTk.check` in `nextCmdSnap` time to trigger
-      -- before kicking off any expensive elaboration (TODO: make expensive elaboration cancelable)
-      unfoldCmdSnaps newMeta validSnaps.toArray cancelTk ctx
-        (startAfterMs := ctx.initParams.editDelay.toUInt32)
-
-    return (AsyncList.delayed newSnapTasks, cancelTk)
-
   /-- Given the new document, updates editable doc state. -/
-  def updateDocument (newMeta : DocumentMeta) : WorkerM Unit := do
-    let (newSnaps, cancelTk) ← startNewSnapshotTasks newMeta
-    modify fun st => { st with doc := { meta := newMeta, cmdSnaps := newSnaps, cancelTk } }
-
+  def updateDocument (meta : DocumentMeta) : WorkerM Unit := do
+    (← get).reporterCancelTk.set
+    let ctx ← read
+    let initSnap ← ctx.processor meta.mkInputContext
+    let doc : EditableDocumentCore := {
+      meta, initSnap
+      diagnosticsRef := (← IO.mkRef ∅)
+    }
+    let reporterCancelTk ← CancelToken.new
+    let reporter ← reportSnapshots ctx doc reporterCancelTk
+    modify fun st => { st with doc := { doc with reporter }, reporterCancelTk }
+    -- we assume version updates are monotonous and that we are on the main thread
+    ctx.maxDocVersionRef.set meta.version
 end Updates
 
 /- Notifications are handled in the main thread. They may change global worker state
@@ -407,6 +430,30 @@ section MessageHandling
       : WorkerM Unit := do
     updatePendingRequests (fun pendingRequests => pendingRequests.insert id requestTask)
 
+  open Widget RequestM Language in
+  def handleGetInteractiveDiagnosticsRequest (params : GetInteractiveDiagnosticsParams) :
+      WorkerM (Array InteractiveDiagnostic) := do
+    let st ← get
+    -- NOTE: always uses latest document (which is the only one we can retrieve diagnostics for);
+    -- any race should be temporary as the client should re-request interactive diagnostics when
+    -- they receive the non-interactive diagnostics for the new document
+    let diags ← st.doc.diagnosticsRef.get
+    -- NOTE: does not wait for `lineRange?` to be fully elaborated, which would be problematic with
+    -- fine-grained incremental reporting anyway; instead, the client is obligated to resend the
+    -- request when the non-interactive diagnostics of this range have changed
+    return diags.filter fun diag =>
+      let r := diag.fullRange
+      let diagStartLine := r.start.line
+      let diagEndLine   :=
+        if r.end.character == 0 then
+          r.end.line
+        else
+          r.end.line + 1
+      params.lineRange?.all fun ⟨s, e⟩ =>
+        -- does [s,e) intersect [diagStartLine,diagEndLine)?
+        s ≤ diagStartLine ∧ diagStartLine < e ∨
+        diagStartLine ≤ s ∧ s < diagEndLine
+
   def handleImportCompletionRequest (id : RequestID) (params : CompletionParams)
       : WorkerM (Task (Except Error AvailableImportsCache)) := do
     let ctx ← read
@@ -417,8 +464,8 @@ section MessageHandling
     | none => IO.asTask do
       let availableImports ← ImportCompletion.collectAvailableImports
       let lastRequestTimestampMs ← IO.monoMsNow
-      let completions := ImportCompletion.find text st.currHeaderStx params availableImports
-      ctx.hOut.writeLspResponse ⟨id, completions⟩
+      let completions := ImportCompletion.find text st.doc.initSnap.stx params availableImports
+      ctx.chanOut.send <| .response id (toJson completions)
       pure { availableImports, lastRequestTimestampMs : AvailableImportsCache }
 
     | some task => IO.mapTask (t := task) fun result => do
@@ -427,8 +474,8 @@ section MessageHandling
       if timestampNowMs - lastRequestTimestampMs >= 10000 then
         availableImports ← ImportCompletion.collectAvailableImports
       lastRequestTimestampMs := timestampNowMs
-      let completions := ImportCompletion.find text st.currHeaderStx params availableImports
-      ctx.hOut.writeLspResponse ⟨id, completions⟩
+      let completions := ImportCompletion.find text st.doc.initSnap.stx params availableImports
+      ctx.chanOut.send <| .response id (toJson completions)
       pure { availableImports, lastRequestTimestampMs : AvailableImportsCache }
 
   def handleRequest (id : RequestID) (method : String) (params : Json)
@@ -436,53 +483,68 @@ section MessageHandling
     let ctx ← read
     let st ← get
 
-    if method == "$/lean/rpc/connect" then
-      try
+    -- special cases
+    try
+      match method with
+      -- needs access to `WorkerState.rpcSessions`
+      | "$/lean/rpc/connect" =>
         let ps ← parseParams RpcConnectParams params
         let resp ← handleRpcConnect ps
-        ctx.hOut.writeLspResponse ⟨id, resp⟩
-      catch e =>
-        ctx.hOut.writeLspResponseError
-          { id
-            code := ErrorCode.internalError
-            message := toString e }
+        ctx.chanOut.send <| .response id (toJson resp)
+        return
+      | "$/lean/rpc/call" =>
+        let params ← parseParams Lsp.RpcCallParams params
+        -- needs access to `EditableDocumentCore.diagnosticsRef`
+        if params.method == `Lean.Widget.getInteractiveDiagnostics then
+          let some seshRef := st.rpcSessions.find? params.sessionId
+            | ctx.chanOut.send <| .responseError id .rpcNeedsReconnect "Outdated RPC session" none
+          let params ← IO.ofExcept (fromJson? params.params)
+          let resp ← handleGetInteractiveDiagnosticsRequest params
+
+          let resp ← seshRef.modifyGet fun st =>
+            rpcEncode resp st.objects |>.map (·) ({st with objects := ·})
+          ctx.chanOut.send <| .response id resp
+          return
+      | "textDocument/completion" =>
+        let params ← parseParams CompletionParams params
+        -- must not wait on import processing snapshot
+        if ImportCompletion.isImportCompletionRequest st.doc.meta.text st.doc.initSnap.stx params
+        then
+          let importCachingTask ← handleImportCompletionRequest id params
+          set { st with importCachingTask? := some importCachingTask }
+          return
+      | _ => pure ()
+    catch e =>
+      ctx.chanOut.send <| .responseError id .internalError (toString e) none
       return
 
-    if method == "textDocument/completion" then
-      let params ← parseParams CompletionParams params
-      if ImportCompletion.isImportCompletionRequest st.doc.meta.text st.currHeaderStx params then
-        let importCachingTask ← handleImportCompletionRequest id params
-        set <| { st with importCachingTask? := some importCachingTask }
-        return
-
-    -- we assume that every request requires at least the header snapshot or the search path
-    let t ← IO.bindTask ctx.headerTask fun x => do
-      let (_, srcSearchPath) ← IO.ofExcept x
+    -- we assume that any other request requires at least the the search path
+    -- TODO: move into language-specific request handling
+    let t ← IO.bindTask st.srcSearchPathTask fun srcSearchPath => do
       let rc : RequestContext :=
         { rpcSessions := st.rpcSessions
           srcSearchPath
           doc := st.doc
           hLog := ctx.hLog
-          hOut := ctx.hOut
           initParams := ctx.initParams }
       let t? ← EIO.toIO' <| handleLspRequest method params rc
       let t₁ ← match t? with
         | Except.error e =>
           IO.asTask do
-            ctx.hOut.writeLspResponseError <| e.toLspResponseError id
+            ctx.chanOut.send <| e.toLspResponseError id
         | Except.ok t => (IO.mapTask · t) fun
           | Except.ok resp =>
-            ctx.hOut.writeLspResponse ⟨id, resp⟩
+            ctx.chanOut.send <| .response id (toJson resp)
           | Except.error e =>
-            ctx.hOut.writeLspResponseError <| e.toLspResponseError id
+            ctx.chanOut.send <| e.toLspResponseError id
     queueRequest id t
 end MessageHandling
 
 section MainLoop
+  variable (hIn : FS.Stream) in
   partial def mainLoop : WorkerM Unit := do
-    let ctx ← read
     let mut st ← get
-    let msg ← ctx.hIn.readLspMessage
+    let msg ← hIn.readLspMessage
     let filterFinishedTasks (acc : PendingRequestMap) (id : RequestID) (task : Task (Except IO.Error Unit))
         : IO PendingRequestMap := do
       if (← hasFinished task) then
@@ -507,9 +569,6 @@ section MainLoop
       handleRequest id method (toJson params)
       mainLoop
     | Message.notification "exit" none =>
-      let doc := st.doc
-      doc.cancelTk.set
-      doc.cmdSnaps.cancel
       return ()
     | Message.notification method (some params) =>
       handleNotification method (toJson params)
@@ -532,12 +591,15 @@ def initAndRunWorker (i o e : FS.Stream) (opts : Options) : IO UInt32 := do
   let e := e.withPrefix s!"[{param.textDocument.uri}] "
   let _ ← IO.setStderr e
   try
-    let (ctx, st) ← initializeWorker meta i o e initParams.param opts
-    let _ ← StateRefT'.run (s := st) <| ReaderT.run (r := ctx) mainLoop
+    let (ctx, st) ← initializeWorker meta o e initParams.param opts
+    let _ ← StateRefT'.run (s := st) <| ReaderT.run (r := ctx) (mainLoop i)
     return (0 : UInt32)
-  catch e =>
-    IO.eprintln e
-    publishDiagnostics meta #[{ range := ⟨⟨0, 0⟩, ⟨0, 0⟩⟩, severity? := DiagnosticSeverity.error, message := e.toString }] o
+  catch err =>
+    IO.eprintln err
+    e.writeLspMessage <| mkPublishDiagnosticsNotification meta #[{
+      range := ⟨⟨0, 0⟩, ⟨0, 0⟩⟩
+      severity? := DiagnosticSeverity.error
+      message := err.toString }]
     return (1 : UInt32)
 
 @[export lean_server_worker_main]

src/Lean/Server/FileWorker/RequestHandling.lean

@@ -35,9 +35,9 @@ def handleCompletion (p : CompletionParams)
   -- such as a trailing dot after an option name. This shouldn't be a problem since any subsequent
   -- command starts with a keyword that (currently?) does not participate in completion.
   withWaitFindSnap doc (·.endPos + ' ' >= pos)
-    (notFoundX := pure { items := #[], isIncomplete := true })
-    (abortedX :=
+    (notFoundX :=
       -- work around https://github.com/microsoft/vscode/issues/155738
+      -- this is important when a snapshot cannot be found because it was aborted
       pure { items := #[{label := "-"}], isIncomplete := true })
     (x := fun snap => do
       if let some r ← Completion.find? p doc.meta.text pos snap.infoTree caps then
@@ -62,7 +62,6 @@ def handleCompletionItemResolve (item : CompletionItem)
   let pos := text.lspPosToUtf8Pos data.params.position
   withWaitFindSnap doc (·.endPos + ' ' >= pos)
     (notFoundX := pure item)
-    (abortedX := pure item)
     (x := fun snap => Completion.resolveCompletionItem? text pos snap.infoTree item id)
 
 open Elab in
@@ -627,8 +626,7 @@ partial def handleWaitForDiagnostics (p : WaitForDiagnosticsParams)
   let t ← RequestM.asTask waitLoop
   RequestM.bindTask t fun doc? => do
     let doc ← liftExcept doc?
-    let t₁ := doc.cmdSnaps.waitAll
-    return t₁.map fun _ => pure WaitForDiagnostics.mk
+    return doc.reporter.map fun _ => pure WaitForDiagnostics.mk
 
 builtin_initialize
   registerLspRequestHandler
@@ -641,7 +639,6 @@ builtin_initialize
     CompletionParams
     CompletionList
     handleCompletion
-    Completion.fillEligibleHeaderDecls
   registerLspRequestHandler
     "completionItem/resolve"
     CompletionItem

src/Lean/Server/FileWorker/SetupFile.lean

@@ -52,15 +52,24 @@ partial def runLakeSetupFile
   let exitCode ← lakeProc.wait
   return ⟨spawnArgs, exitCode, stdout, stderr⟩
 
+/-- Categorizes possible outcomes of running `lake setup-file`. -/
 inductive FileSetupResultKind where
+  /-- File configuration loaded and dependencies updated successfully. -/
   | success
+  /-- No Lake project found, no setup was done. -/
   | noLakefile
+  /-- Imports must be rebuilt but `--no-build` was specified. -/
   | importsOutOfDate
+  /-- Other error during Lake invocation. -/
   | error (msg : String)
 
+/-- Result of running `lake setup-file`. -/
 structure FileSetupResult where
+  /-- Kind of outcome. -/
   kind          : FileSetupResultKind
+  /-- Search path from successful setup, or else empty. -/
   srcSearchPath : SearchPath
+  /-- Additional options from successful setup, or else empty. -/
   fileOptions   : Options
 
 def FileSetupResult.ofSuccess (pkgSearchPath : SearchPath) (fileOptions : Options)
@@ -88,11 +97,6 @@ def FileSetupResult.ofError (msg : String) : IO FileSetupResult := do return {
   fileOptions   := Options.empty
 }
 
-def FileSetupResult.addGlobalOptions (result : FileSetupResult) (globalOptions : Options)
-    : FileSetupResult :=
-  let fileOptions := globalOptions.mergeBy (fun _ _ fileOpt => fileOpt) result.fileOptions
-  { result with fileOptions := fileOptions }
-
 /-- Uses `lake setup-file` to compile dependencies on the fly and add them to `LEAN_PATH`.
 Compilation progress is reported to `handleStderr`. Returns the search path for
 source files and the options for the file. -/

src/Lean/Server/FileWorker/Utils.lean

@@ -5,29 +5,16 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Wojciech Nawrocki, Marc Huisinga
 -/
 prelude
+import Lean.Language.Lean
 import Lean.Server.Utils
 import Lean.Server.Snapshots
 import Lean.Server.AsyncList
-
 import Lean.Server.Rpc.Basic
 
 namespace Lean.Server.FileWorker
 open Snapshots
 open IO
 
-def logSnapContent (s : Snapshot) (text : FileMap) : IO Unit :=
-  IO.eprintln s!"[{s.beginPos}, {s.endPos}]: ```\n{text.source.extract s.beginPos (s.endPos - ⟨1⟩)}\n```"
-
-inductive ElabTaskError where
-  | aborted
-  | ioError (e : IO.Error)
-
-instance : Coe IO.Error ElabTaskError :=
-  ⟨ElabTaskError.ioError⟩
-
-instance : MonadLift IO (EIO ElabTaskError) where
-  monadLift act := act.toEIO (↑ ·)
-
 structure CancelToken where
   ref : IO.Ref Bool
 
@@ -36,24 +23,59 @@ namespace CancelToken
 def new : IO CancelToken :=
   CancelToken.mk <$> IO.mkRef false
 
-def check [MonadExceptOf ElabTaskError m] [MonadLiftT (ST RealWorld) m] [Monad m] (tk : CancelToken) : m Unit := do
-  let c ← tk.ref.get
-  if c then
-    throw ElabTaskError.aborted
-
-def set (tk : CancelToken) : IO Unit :=
+def set (tk : CancelToken) : BaseIO Unit :=
   tk.ref.set true
 
+def isSet (tk : CancelToken) : BaseIO Bool :=
+  tk.ref.get
+
 end CancelToken
 
-/-- A document editable in the sense that we track the environment
-and parser state after each command so that edits can be applied
-without recompiling code appearing earlier in the file. -/
-structure EditableDocument where
-  meta       : DocumentMeta
-  /-- State snapshots after header and each command. -/
-  cmdSnaps   : AsyncList ElabTaskError Snapshot
-  cancelTk   : CancelToken
+-- TEMP: translate from new heterogeneous snapshot tree to old homogeneous async list
+private partial def mkCmdSnaps (initSnap : Language.Lean.InitialSnapshot) :
+    AsyncList IO.Error Snapshot := Id.run do
+  let some headerParsed := initSnap.result? | return .nil
+  .delayed <| headerParsed.processedSnap.task.bind fun headerProcessed => Id.run do
+    let some headerSuccess := headerProcessed.result? | return .pure <| .ok .nil
+    return .pure <| .ok <| .cons {
+      stx := initSnap.stx
+      mpState := headerParsed.parserState
+      cmdState := headerSuccess.cmdState
+    } <| .delayed <| headerSuccess.firstCmdSnap.task.bind go
+where go cmdParsed :=
+  cmdParsed.data.sigSnap.task.bind fun sig =>
+    sig.finishedSnap.task.map fun finished =>
+      .ok <| .cons {
+        stx := cmdParsed.data.stx
+        mpState := cmdParsed.data.parserState
+        cmdState := finished.cmdState
+      } (match cmdParsed.next? with
+        | some next => .delayed <| next.task.bind go
+        | none => .nil)
+
+/--
+A document bundled with processing information. Turned into `EditableDocument` as soon as the
+reporter task has been started.
+-/
+structure EditableDocumentCore where
+  /-- The document. -/
+  meta     : DocumentMeta
+  /-- Initial processing snapshot. -/
+  initSnap : Language.Lean.InitialSnapshot
+  /-- Old representation for backward compatibility. -/
+  cmdSnaps : AsyncList IO.Error Snapshot := mkCmdSnaps initSnap
+  /--
+  Interactive versions of diagnostics reported so far. Filled by `reportSnapshots` and read by
+  `handleGetInteractiveDiagnosticsRequest`.
+  -/
+  diagnosticsRef : IO.Ref (Array Widget.InteractiveDiagnostic)
+
+/-- `EditableDocumentCore` with reporter task. -/
+structure EditableDocument extends EditableDocumentCore where
+  /--
+    Task reporting processing status back to client. We store it here for implementing
+    `waitForDiagnostics`. -/
+  reporter : Task Unit
 
 namespace EditableDocument
 

src/Lean/Server/FileWorker/WidgetRequests.lean

@@ -101,35 +101,6 @@ structure GetInteractiveDiagnosticsParams where
   lineRange? : Option Lsp.LineRange
   deriving Inhabited, FromJson, ToJson
 
-open RequestM in
-def getInteractiveDiagnostics (params : GetInteractiveDiagnosticsParams) : RequestM (RequestTask (Array InteractiveDiagnostic)) := do
-  let doc ← readDoc
-  let rangeEnd := params.lineRange?.map fun range =>
-    doc.meta.text.lspPosToUtf8Pos ⟨range.«end», 0⟩
-  let t := doc.cmdSnaps.waitUntil fun snap => rangeEnd.any (snap.endPos >= ·)
-  pure <| t.map fun (snaps, _) =>
-    let diags? := snaps.getLast?.map fun snap =>
-      snap.interactiveDiags.toArray.filter fun diag =>
-        let r := diag.fullRange
-        let diagStartLine := r.start.line
-        let diagEndLine   :=
-          if r.end.character == 0 then
-            r.end.line
-          else
-            r.end.line + 1
-        params.lineRange?.all fun ⟨s, e⟩ =>
-          -- does [s,e) intersect [diagStartLine,diagEndLine)?
-          s ≤ diagStartLine ∧ diagStartLine < e ∨
-          diagStartLine ≤ s ∧ s < diagEndLine
-    pure <| diags?.getD #[]
-
-builtin_initialize
-  registerBuiltinRpcProcedure
-    `Lean.Widget.getInteractiveDiagnostics
-    GetInteractiveDiagnosticsParams
-    (Array InteractiveDiagnostic)
-    getInteractiveDiagnostics
-
 structure GetGoToLocationParams where
   kind : GoToKind
   info : WithRpcRef InfoWithCtx

src/Lean/Server/README.md

@@ -1,3 +1,7 @@
+# Language Server
+
+This directory contains the implementation of the Lean Language Server, an implementation as well as extension of the [Language Server Protocol](https://microsoft.github.io/language-server-protocol/) used for communicating with Lean extensions in editors.
+
 ## Building & Developing
 
 Both watchdog and worker (see below) are part of the main `lean` binary.
@@ -43,19 +47,37 @@ In Lean 4, the situation is different as `.olean` artifacts are required to exis
 
 ### Worker architecture
 
-A central concept in the worker is the *snapshot*.
-Lean files are processed (elaborated) strictly from top to bottom, with each command being fully processed before processing of subsequent commands is started.
-The worker implements this same processing order, but saves a `Snapshot` of the entire elaboration state after the imports and each command, which is cheap and easy to do because the state is all functional data structures.
-Thanks to these snapshots, we can restart processing the file at the point of a change by discarding and rebuilding all snapshots after it.
-Snapshots are computed asynchronously and stored in an `AsyncList`, which is a list whose tail is potentially still being processed.
-Request handlers usually locate and access a single snapshot in the list based on the cursor position using `withWaitFindSnap`, which will wait for elaboration if it is not sufficiently progressed yet.
+The actual processing of the opened file is left to the `Lean.Language.Lean` processor.
+The interface is shared with the cmdline driver that is used for building Lean files but the core concept of *snapshots* produced by a processor is mostly of interest to the worker: snapshots are how processing data is saved and reused between versions of a file such as to make processing on file edits incremental.
+How exactly incrementality is implemented is left completely to the processor; in a strictly ordered language like Lean, there may be a chain of snapshots for each top-level command, potentially with nested snapshots for increased granularity of incrementality.
+In languages with less strict ordering and less syntax extensibility, there may be a single snapshot for the full syntax tree of the file, and then nested snapshots for processing each declaration in it.
+In the simplest case, there may be a single snapshot after processing the full file, without any incrementality.
+All the worker needs to know is that `Lean.Language.Lean.process` returns a root snapshot of some type that can be transformed into an asynchronously constructed tree of the generic `Lean.Language.Snapshot` type via `Lean.Language.ToSnapshotTree`.
+We use a tree and not an asynchronous list (which would basically be a channel) in order to accommodate parallel processing where it's not clear a priori which of a number of child snapshots will be available first.
+After loading the file and after each edit, the server will obtain this tree from the processor given the new file content and asynchronously wait on all its nodes and report the processing status (diagnostics and progress information) stored therein to the client incrementally (`Lean.Server.FileWorker.reportSnapshots`).
+
+(TODO: the remainder is currently hard-coded for the `Lean` language)
+
+Request handlers usually locate and access a single snapshot in the tree based on the cursor position using `withWaitFindSnap`, which will wait for elaboration if it is not sufficiently progressed yet.
 After the snapshot is available, they can access its data, in particular the command's `Syntax` tree and elaboration `InfoTree`, in order to respond to the request.
 
-The `InfoTree` is the second central server data structure.
+The `InfoTree` is the second central server metadata structure.
 It is filled during elaboration with various metadata that cannot (easily) be recovered from the kernel declarations in the environment: goal & subterm infos including the precise local & metavariable contexts used during elaboration, macro expansion steps, ...
 Once a relevant `Snapshot` `snap` has been located, `snap.infoTree.smallestInfo?` and other functions from `Lean.Server.InfoUtils` can be used to further locate information about a document position.
 The command `set_option trace.Elab.info true` instructs Lean to display the `InfoTree` for the declarations occurring after it.
 
+### Communication
+
+The asynchronous nature of the file worker complicates communication.
+As mentioned above, a language processor does not directly talk to the worker but merely places diagnostics in the asynchronously constructed tree of snapshots where they will be picked up by the reporting task in the worker.
+From there the diagnostics are passed to the channel `WorkerContext.chanOut` read by a single dedicated thread and finally written to stdout, which ensures that reporting tasks from multiple document versions cannot race on the eventual output.
+Request responses are sent to this channel as well.
+
+A further complication in communication is interactive diagnostics, which unlike most Lean objects have relevant *identity* as the client can hold references to them, and replacing an interactive diagnostic with an equal but not identical diagnostic can lead to the client reloading the view and losing local state such as the unfolding of a trace tree.
+Thus we have to make sure that when we reuse snapshots, we reuse interactive diagnostic objects as is.
+On the other hand, we do not want language processors to think about interactive diagnostics for simplicity and module dependency reasons, so we transform diagnostics into interactive ones in the reporting task in the worker and have language processors merely store a dynamic `IO.Ref` in `Snapshot.Diagnostics` that the reporting task then uses to store and reuse interactive diagnostics.
+We initially stored unique IDs in `Snapshot.Diagnostics` for this that would be associated with the cached value in a map in the worker but there was no practical upside to this additional bookkeeping overhead.
+
 ## Code style
 
 Comments should exist to denote specifics of our implementation but, for

src/Lean/Server/References.lean

@@ -38,13 +38,15 @@ def addRef : RefInfo → Reference → RefInfo
     { i with usages := usages.push ref }
   | i, _ => i
 
-def toLspRefInfo (i : RefInfo) : IO Lsp.RefInfo := do
-  let refToRefInfoLocation (ref : Reference) : IO RefInfo.Location := do
+def toLspRefInfo (i : RefInfo) : BaseIO Lsp.RefInfo := do
+  let refToRefInfoLocation (ref : Reference) : BaseIO RefInfo.Location := do
     let parentDeclName? := ref.ci.parentDecl?
-    let parentDeclRanges? ← ref.ci.runMetaM ref.info.lctx do
-      let some parentDeclName := parentDeclName?
-        | return none
-      findDeclarationRanges? parentDeclName
+    let .ok parentDeclRanges? ← EIO.toBaseIO <| ref.ci.runCoreM do
+        let some parentDeclName := parentDeclName?
+          | return none
+        findDeclarationRanges? parentDeclName
+      -- we only use `CoreM` to get access to a `MonadEnv`, but these are currently all `IO`
+      | unreachable!
     return {
       range := ref.range
       parentDecl? := do
@@ -70,7 +72,7 @@ def addRef (self : ModuleRefs) (ref : Reference) : ModuleRefs :=
   let refInfo := self.findD ref.ident RefInfo.empty
   self.insert ref.ident (refInfo.addRef ref)
 
-def toLspModuleRefs (refs : ModuleRefs) : IO Lsp.ModuleRefs := do
+def toLspModuleRefs (refs : ModuleRefs) : BaseIO Lsp.ModuleRefs := do
   let refs ← refs.toList.mapM fun (k, v) => do
     return (k, ← v.toLspRefInfo)
   return HashMap.ofList refs

src/Lean/Server/Requests.lean

@@ -64,7 +64,6 @@ structure RequestContext where
   srcSearchPath : SearchPath
   doc           : FileWorker.EditableDocument
   hLog          : IO.FS.Stream
-  hOut          : IO.FS.Stream
   initParams    : Lsp.InitializeParams
 
 abbrev RequestTask α := Task (Except RequestError α)
@@ -106,14 +105,12 @@ def bindTask (t : Task α) (f : α → RequestM (RequestTask β)) : RequestM (Re
   let rc ← readThe RequestContext
   EIO.bindTask t (f · rc)
 
-def waitFindSnapAux (notFoundX abortedX : RequestM α) (x : Snapshot → RequestM α)
-    : Except ElabTaskError (Option Snapshot) → RequestM α
+def waitFindSnapAux (notFoundX : RequestM α) (x : Snapshot → RequestM α)
+    : Except IO.Error (Option Snapshot) → RequestM α
   /- The elaboration task that we're waiting for may be aborted if the file contents change.
   In that case, we reply with the `fileChanged` error by default. Thanks to this, the server doesn't
   get bogged down in requests for an old state of the document. -/
-  | Except.error FileWorker.ElabTaskError.aborted => abortedX
-  | Except.error (FileWorker.ElabTaskError.ioError e) =>
-    throw (RequestError.ofIoError e)
+  | Except.error e => throw (RequestError.ofIoError e)
   | Except.ok none => notFoundX
   | Except.ok (some snap) => x snap
 
@@ -123,19 +120,17 @@ and if a matching snapshot was found executes `x` with it. If not found, the tas
 def withWaitFindSnap (doc : EditableDocument) (p : Snapshot → Bool)
     (notFoundX : RequestM β)
     (x : Snapshot → RequestM β)
-    (abortedX : RequestM β := throwThe RequestError .fileChanged)
     : RequestM (RequestTask β) := do
   let findTask := doc.cmdSnaps.waitFind? p
-  mapTask findTask <| waitFindSnapAux notFoundX abortedX x
+  mapTask findTask <| waitFindSnapAux notFoundX x
 
 /-- See `withWaitFindSnap`. -/
 def bindWaitFindSnap (doc : EditableDocument) (p : Snapshot → Bool)
     (notFoundX : RequestM (RequestTask β))
     (x : Snapshot → RequestM (RequestTask β))
-    (abortedX : RequestM (RequestTask β) := throwThe RequestError .fileChanged)
     : RequestM (RequestTask β) := do
   let findTask := doc.cmdSnaps.waitFind? p
-  bindTask findTask <| waitFindSnapAux notFoundX abortedX x
+  bindTask findTask <| waitFindSnapAux notFoundX x
 
 /-- Create a task which waits for the snapshot containing `lspPos` and executes `f` with it.
 If no such snapshot exists, the request fails with an error. -/
@@ -182,13 +177,8 @@ section HandlerTable
 open Lsp
 
 structure RequestHandler where
-  fileSource          : Json → Except RequestError Lsp.DocumentUri
-  handle              : Json → RequestM (RequestTask Json)
-  /--
-  Handler that is called by the file worker after processing the header with the header environment.
-  Enables request handlers to cache data related to imports.
-  -/
-  handleHeaderCaching : Environment → IO Unit
+  fileSource : Json → Except RequestError Lsp.DocumentUri
+  handle : Json → RequestM (RequestTask Json)
 
 builtin_initialize requestHandlers : IO.Ref (PersistentHashMap String RequestHandler) ←
   IO.mkRef {}
@@ -208,9 +198,7 @@ as LSP error responses. -/
 def registerLspRequestHandler (method : String)
     paramType [FromJson paramType] [FileSource paramType]
     respType [ToJson respType]
-    (handler : paramType → RequestM (RequestTask respType))
-    (headerCachingHandler : Environment → IO Unit := fun _ => pure ())
-    : IO Unit := do
+    (handler : paramType → RequestM (RequestTask respType)) : IO Unit := do
   if !(← Lean.initializing) then
     throw <| IO.userError s!"Failed to register LSP request handler for '{method}': only possible during initialization"
   if (← requestHandlers.get).contains method then
@@ -221,8 +209,8 @@ def registerLspRequestHandler (method : String)
     let params ← liftExcept <| parseRequestParams paramType j
     let t ← handler params
     pure <| t.map <| Except.map ToJson.toJson
-  let handleHeaderCaching := headerCachingHandler
-  requestHandlers.modify fun rhs => rhs.insert method { fileSource, handle, handleHeaderCaching }
+
+  requestHandlers.modify fun rhs => rhs.insert method { fileSource, handle }
 
 def lookupLspRequestHandler (method : String) : IO (Option RequestHandler) :=
   return (← requestHandlers.get).find? method
@@ -253,14 +241,6 @@ def chainLspRequestHandler (method : String)
   else
     throw <| IO.userError s!"Failed to chain LSP request handler for '{method}': no initial handler registered"
 
-/--
-Runs the header caching handler for every single registered request handler using the given
-`headerEnv`.
--/
-def runHeaderCachingHandlers (headerEnv : Environment) : IO Unit := do
-  (← requestHandlers.get).forM fun _ handler =>
-    handler.handleHeaderCaching headerEnv
-
 def routeLspRequest (method : String) (params : Json) : IO (Except RequestError DocumentUri) := do
   match (← lookupLspRequestHandler method) with
   | none => return Except.error <| RequestError.methodNotFound method

src/Lean/Server/Snapshots.lean

@@ -9,6 +9,7 @@ import Init.System.IO
 
 import Lean.Elab.Import
 import Lean.Elab.Command
+import Lean.Language.Lean
 
 import Lean.Widget.InteractiveDiagnostic
 
@@ -26,19 +27,9 @@ builtin_initialize dummyTacticCache : IO.Ref Tactic.Cache ← IO.mkRef {}
 
 /-- What Lean knows about the world after the header and each command. -/
 structure Snapshot where
-  /-- Where the command which produced this snapshot begins. Note that
-  neighbouring snapshots are *not* necessarily attached beginning-to-end,
-  since inputs outside the grammar advance the parser but do not produce
-  snapshots. -/
-  beginPos : String.Pos
   stx : Syntax
   mpState : Parser.ModuleParserState
   cmdState : Command.State
-  /-- We cache interactive diagnostics in order not to invoke the pretty-printer again on messages
-  from previous snapshots when publishing diagnostics for every new snapshot (this is quadratic),
-  as well as not to invoke it once again when handling `$/lean/interactiveDiagnostics`. -/
-  interactiveDiags : PersistentArray Widget.InteractiveDiagnostic
-  tacticCache : IO.Ref Tactic.Cache
 
 namespace Snapshot
 
@@ -51,9 +42,6 @@ def env (s : Snapshot) : Environment :=
 def msgLog (s : Snapshot) : MessageLog :=
   s.cmdState.messages
 
-def diagnostics (s : Snapshot) : PersistentArray Lsp.Diagnostic :=
-  s.interactiveDiags.map fun d => d.toDiagnostic
-
 def infoTree (s : Snapshot) : InfoTree :=
   -- the parser returns exactly one command per snapshot, and the elaborator creates exactly one node per command
   assert! s.cmdState.infoState.trees.size == 1
@@ -66,10 +54,10 @@ open Command in
 /-- Use the command state in the given snapshot to run a `CommandElabM`.-/
 def runCommandElabM (snap : Snapshot) (meta : DocumentMeta) (c : CommandElabM α) : EIO Exception α := do
   let ctx : Command.Context := {
-    cmdPos := snap.beginPos,
+    cmdPos := snap.stx.getPos? |>.getD 0,
     fileName := meta.uri,
     fileMap := meta.text,
-    tacticCache? := snap.tacticCache,
+    tacticCache? := none
   }
   c.run ctx |>.run' snap.cmdState
 
@@ -83,82 +71,4 @@ def runTermElabM (snap : Snapshot) (meta : DocumentMeta) (c : TermElabM α) : EI
 
 end Snapshot
 
-/-- Parses the next command occurring after the given snapshot
-without elaborating it. -/
-def parseNextCmd (inputCtx : Parser.InputContext) (snap : Snapshot) : IO Syntax := do
-  let cmdState := snap.cmdState
-  let scope := cmdState.scopes.head!
-  let pmctx := { env := cmdState.env, options := scope.opts, currNamespace := scope.currNamespace, openDecls := scope.openDecls }
-  let (cmdStx, _, _) :=
-    Parser.parseCommand inputCtx pmctx snap.mpState snap.msgLog
-  return cmdStx
-
-register_builtin_option server.stderrAsMessages : Bool := {
-  defValue := true
-  group    := "server"
-  descr    := "(server) capture output to the Lean stderr channel (such as from `dbg_trace`) during elaboration of a command as a diagnostic message"
-}
-
-/-- Compiles the next command occurring after the given snapshot. If there is no next command
-(file ended), `Snapshot.isAtEnd` will hold of the return value. -/
--- NOTE: This code is really very similar to Elab.Frontend. But generalizing it
--- over "store snapshots"/"don't store snapshots" would likely result in confusing
--- isServer? conditionals and not be worth it due to how short it is.
-def compileNextCmd (inputCtx : Parser.InputContext) (snap : Snapshot) (hasWidgets : Bool) : IO Snapshot := do
-  let cmdState := snap.cmdState
-  let scope := cmdState.scopes.head!
-  let pmctx := { env := cmdState.env, options := scope.opts, currNamespace := scope.currNamespace, openDecls := scope.openDecls }
-  let (cmdStx, cmdParserState, msgLog) :=
-    Parser.parseCommand inputCtx pmctx snap.mpState snap.msgLog
-  let cmdPos := cmdStx.getPos?.get!
-  let cmdStateRef ← IO.mkRef { snap.cmdState with messages := msgLog }
-  /- The same snapshot may be executed by different tasks. So, to make sure `elabCommandTopLevel` has exclusive
-      access to the cache, we create a fresh reference here. Before this change, the
-      following `snap.tacticCache.modify` would reset the tactic post cache while another snapshot was still using it. -/
-  let tacticCacheNew ← IO.mkRef (← snap.tacticCache.get)
-  let cmdCtx : Elab.Command.Context := {
-    cmdPos       := snap.endPos
-    fileName     := inputCtx.fileName
-    fileMap      := inputCtx.fileMap
-    tacticCache? := some tacticCacheNew
-  }
-  let (output, _) ← IO.FS.withIsolatedStreams (isolateStderr := server.stderrAsMessages.get scope.opts) <| liftM (m := BaseIO) do
-    Elab.Command.catchExceptions
-      (getResetInfoTrees *> Elab.Command.elabCommandTopLevel cmdStx)
-      cmdCtx cmdStateRef
-  let postNew := (← tacticCacheNew.get).post
-  snap.tacticCache.modify fun _ => { pre := postNew, post := {} }
-  let mut postCmdState ← cmdStateRef.get
-  if !output.isEmpty then
-    postCmdState := {
-      postCmdState with
-      messages := postCmdState.messages.add {
-        fileName := inputCtx.fileName
-        severity := MessageSeverity.information
-        pos      := inputCtx.fileMap.toPosition snap.endPos
-        data     := output
-      }
-    }
-  let postCmdSnap : Snapshot := {
-    beginPos := cmdPos
-    stx := cmdStx
-    mpState := cmdParserState
-    cmdState := postCmdState
-    interactiveDiags := ← withNewInteractiveDiags postCmdState.messages
-    tacticCache := (← IO.mkRef {})
-  }
-  return postCmdSnap
-
-where
-  /-- Compute the current interactive diagnostics log by finding a "diff" relative to the parent
-  snapshot. We need to do this because unlike the `MessageLog` itself, interactive diags are not
-  part of the command state. -/
-  withNewInteractiveDiags (msgLog : MessageLog) : IO (PersistentArray Widget.InteractiveDiagnostic) := do
-    let newMsgCount := msgLog.msgs.size - snap.msgLog.msgs.size
-    let mut ret := snap.interactiveDiags
-    for i in List.iota newMsgCount do
-      let newMsg := msgLog.msgs.get! (msgLog.msgs.size - i)
-      ret := ret.push (← Widget.msgToInteractiveDiagnostic inputCtx.fileMap newMsg hasWidgets)
-    return ret
-
 end Lean.Server.Snapshots

src/Lean/Server/Utils.lean

@@ -114,36 +114,39 @@ def applyDocumentChange (oldText : FileMap) : (change : Lsp.TextDocumentContentC
 def foldDocumentChanges (changes : Array Lsp.TextDocumentContentChangeEvent) (oldText : FileMap) : FileMap :=
   changes.foldl applyDocumentChange oldText
 
-def publishDiagnostics (m : DocumentMeta) (diagnostics : Array Lsp.Diagnostic) (hOut : FS.Stream) : IO Unit :=
-  hOut.writeLspNotification {
-    method := "textDocument/publishDiagnostics"
-    param  := {
-      uri         := m.uri
-      version?    := some m.version
-      diagnostics := diagnostics
-      : PublishDiagnosticsParams
-    }
+/-- Constructs a `textDocument/publishDiagnostics` notification. -/
+def mkPublishDiagnosticsNotification (m : DocumentMeta) (diagnostics : Array Lsp.Diagnostic) :
+    JsonRpc.Notification Lsp.PublishDiagnosticsParams where
+  method := "textDocument/publishDiagnostics"
+  param  := {
+    uri         := m.uri
+    version?    := some m.version
+    diagnostics := diagnostics
   }
 
-def publishProgress (m : DocumentMeta) (processing : Array LeanFileProgressProcessingInfo) (hOut : FS.Stream) : IO Unit :=
-  hOut.writeLspNotification {
-    method := "$/lean/fileProgress"
-    param := {
-      textDocument := { uri := m.uri, version? := m.version }
-      processing
-      : LeanFileProgressParams
-    }
+/-- Constructs a `$/lean/fileProgress` notification. -/
+def mkFileProgressNotification (m : DocumentMeta) (processing : Array LeanFileProgressProcessingInfo) :
+    JsonRpc.Notification Lsp.LeanFileProgressParams where
+  method := "$/lean/fileProgress"
+  param := {
+    textDocument := { uri := m.uri, version? := m.version }
+    processing
   }
 
-def publishProgressAtPos (m : DocumentMeta) (pos : String.Pos) (hOut : FS.Stream) (kind : LeanFileProgressKind := LeanFileProgressKind.processing) : IO Unit :=
-  publishProgress m #[{ range := ⟨m.text.utf8PosToLspPos pos, m.text.utf8PosToLspPos m.text.source.endPos⟩, kind := kind }] hOut
+/-- Constructs a `$/lean/fileProgress` notification from the given position onwards. -/
+def mkFileProgressAtPosNotification (m : DocumentMeta) (pos : String.Pos)
+  (kind : LeanFileProgressKind := LeanFileProgressKind.processing) :
+    JsonRpc.Notification Lsp.LeanFileProgressParams :=
+  mkFileProgressNotification m #[{ range := ⟨m.text.utf8PosToLspPos pos, m.text.utf8PosToLspPos m.text.source.endPos⟩, kind := kind }]
 
-def publishProgressDone (m : DocumentMeta) (hOut : FS.Stream) : IO Unit :=
-  publishProgress m #[] hOut
+/-- Constructs a `$/lean/fileProgress` notification marking processing as done. -/
+def mkFileProgressDoneNotification (m : DocumentMeta) : JsonRpc.Notification Lsp.LeanFileProgressParams :=
+  mkFileProgressNotification m #[]
 
 -- TODO: should return a request ID (or task?) when we add response handling
-def applyWorkspaceEdit (params : ApplyWorkspaceEditParams) (hOut : FS.Stream) : IO Unit :=
-  hOut.writeLspRequest ⟨"workspace/applyEdit", "workspace/applyEdit", params⟩
+def mkApplyWorkspaceEditRequest (params : ApplyWorkspaceEditParams) :
+    JsonRpc.Request ApplyWorkspaceEditParams :=
+  ⟨"workspace/applyEdit", "workspace/applyEdit", params⟩
 
 end Lean.Server
 

src/Lean/Server/Watchdog.lean

@@ -235,7 +235,7 @@ section ServerM
             ++ " or the server is shutting down.")
           -- one last message to clear the diagnostics for this file so that stale errors
           -- do not remain in the editor forever.
-          publishDiagnostics fw.doc #[] o
+          o.writeLspMessage <| mkPublishDiagnosticsNotification fw.doc #[]
           return WorkerEvent.terminated
         | 2 =>
           return .importsChanged
@@ -248,7 +248,7 @@ section ServerM
               (ErrorCode.workerCrashed, "likely due to a stack overflow or a bug")
           fw.errorPendingRequests o errorCode
             s!"Server process for {fw.doc.uri} crashed, {errorCausePointer}."
-          publishProgressAtPos fw.doc 0 o (kind := LeanFileProgressKind.fatalError)
+          o.writeLspMessage <| mkFileProgressAtPosNotification fw.doc 0 (kind := LeanFileProgressKind.fatalError)
           return WorkerEvent.crashed err
       loop
     let task ← IO.asTask (loop $ ←read) Task.Priority.dedicated
@@ -257,7 +257,7 @@ section ServerM
       | Except.error e => WorkerEvent.ioError e
 
   def startFileWorker (m : DocumentMeta) : ServerM Unit := do
-    publishProgressAtPos m 0 (← read).hOut
+    (← read).hOut.writeLspMessage <| mkFileProgressAtPosNotification m 0
     let st ← read
     let workerProc ← Process.spawn {
       toStdioConfig := workerCfg

src/Lean/Widget/InteractiveDiagnostic.lean

@@ -190,7 +190,8 @@ partial def msgToInteractive (msgData : MessageData) (hasWidgets : Bool) (indent
   fmtToTT fmt indent
 
 /-- Transform a Lean Message concerning the given text into an LSP Diagnostic. -/
-def msgToInteractiveDiagnostic (text : FileMap) (m : Message) (hasWidgets : Bool) : IO InteractiveDiagnostic := do
+def msgToInteractiveDiagnostic (text : FileMap) (m : Message) (hasWidgets : Bool) :
+    BaseIO InteractiveDiagnostic := do
   let low : Lsp.Position := text.leanPosToLspPos m.pos
   let fullHigh := text.leanPosToLspPos <| m.endPos.getD m.pos
   let high : Lsp.Position := match m.endPos with
@@ -212,10 +213,9 @@ def msgToInteractiveDiagnostic (text : FileMap) (m : Message) (hasWidgets : Bool
     if m.data.isDeprecationWarning then some #[.deprecated]
     else if m.data.isUnusedVariableWarning then some #[.unnecessary]
     else none
-  let message ← try
-      msgToInteractive m.data hasWidgets
-    catch ex =>
-      pure <| TaggedText.text s!"[error when printing message: {ex.toString}]"
+  let message := match (← msgToInteractive m.data hasWidgets |>.toBaseIO) with
+    | .ok msg => msg
+    | .error ex => TaggedText.text s!"[error when printing message: {ex.toString}]"
   pure { range, fullRange? := some fullRange, severity?, source?, message, tags? }
 
 end Lean.Widget

tests/lean/interactive/editAfterError.lean.expected.out

@@ -22,39 +22,3 @@
    "fullRange":
    {"start": {"line": 1, "character": 7},
     "end": {"line": 1, "character": 11}}}]}
-{"version": 2,
- "uri": "file:///editAfterError.lean",
- "diagnostics":
- [{"source": "Lean 4",
-   "severity": 1,
-   "range":
-   {"start": {"line": 0, "character": 7}, "end": {"line": 0, "character": 10}},
-   "message": "unknown identifier 'tru'",
-   "fullRange":
-   {"start": {"line": 0, "character": 7}, "end": {"line": 0, "character": 10}}},
-  {"source": "Lean 4",
-   "severity": 1,
-   "range":
-   {"start": {"line": 1, "character": 7}, "end": {"line": 1, "character": 11}},
-   "message": "unknown identifier 'fals'",
-   "fullRange":
-   {"start": {"line": 1, "character": 7},
-    "end": {"line": 1, "character": 11}}}]}
-{"version": 2,
- "uri": "file:///editAfterError.lean",
- "diagnostics":
- [{"source": "Lean 4",
-   "severity": 1,
-   "range":
-   {"start": {"line": 0, "character": 7}, "end": {"line": 0, "character": 10}},
-   "message": "unknown identifier 'tru'",
-   "fullRange":
-   {"start": {"line": 0, "character": 7}, "end": {"line": 0, "character": 10}}},
-  {"source": "Lean 4",
-   "severity": 1,
-   "range":
-   {"start": {"line": 1, "character": 7}, "end": {"line": 1, "character": 11}},
-   "message": "unknown identifier 'fals'",
-   "fullRange":
-   {"start": {"line": 1, "character": 7},
-    "end": {"line": 1, "character": 11}}}]}

tests/lean/interactive/run.lean

@@ -127,9 +127,8 @@ partial def main (args : List String) : IO Unit := do
           requestNo := requestNo + 1
           versionNo := versionNo + 1
         | "collectDiagnostics" =>
-          let diags ← Ipc.collectDiagnostics requestNo uri (versionNo - 1)
-          for diag in diags do
-            IO.eprintln (toJson diag.param)
+          if let some diags ← Ipc.collectDiagnostics requestNo uri (versionNo - 1) then
+            IO.eprintln (toJson diags.param)
           requestNo := requestNo + 1
         | "codeAction" =>
           let params : CodeActionParams := {

tests/lean/interactive/unterminatedDocComment.lean.expected.out

@@ -4,5 +4,3 @@
    "range":
    {"start": {"line": 5, "character": 2}, "end": {"line": 5, "character": 3}}}]}
 {"version": 2, "uri": "file:///unterminatedDocComment.lean", "diagnostics": []}
-{"version": 2, "uri": "file:///unterminatedDocComment.lean", "diagnostics": []}
-{"version": 2, "uri": "file:///unterminatedDocComment.lean", "diagnostics": []}

tests/lean/partialSyntaxTraces.lean

@@ -0,0 +1,5 @@
+/-! Partial syntax should not suppress trace output. -/
+
+set_option pp.rawOnError true
+set_option trace.Elab.command true
+def f

tests/lean/partialSyntaxTraces.lean.expected.out

@@ -0,0 +1,6 @@
+partialSyntaxTraces.lean:6:0: error: unexpected end of input; expected ':=', 'where' or '|'
+[Elab.command] [Error pretty printing syntax: parenthesize: uncaught backtrack exception. Falling back to raw printer.]
+    (Command.declaration
+     (Command.declModifiers [] [] [] [] [] [])
+     (Command.definition "def" (Command.declId `f []) (Command.optDeclSig [] []) (Command.whereStructInst <missing>)))
+[Elab.command] 

tests/lean/server/diags.lean

@@ -11,19 +11,16 @@ def main : IO Unit := do
 
     hIn.write (←FS.readBinFile "open_content.log")
     hIn.flush
-    let diags ← Ipc.collectDiagnostics 1 "file:///test.lean" 1
-    if diags.isEmpty then
-      throw $ userError "Test failed, no diagnostics received."
+    let some diag ← Ipc.collectDiagnostics 1 "file:///test.lean" 1
+      | throw $ userError "Test failed, no diagnostics received."
+    FS.writeFile "content_diag.json.produced" (toString <| toJson (diag : JsonRpc.Message))
+
+    if let Except.ok (refDiag : JsonRpc.Notification PublishDiagnosticsParams) :=
+      (Json.parse $ ←FS.readFile "content_diag.json") >>= fromJson?
+    then
+      assert! (diag == refDiag)
     else
-      let diag := diags.getLast!
-      FS.writeFile "content_diag.json.produced" (toString <| toJson (diag : JsonRpc.Message))
+      throw $ userError "Failed parsing test file."
 
-      if let Except.ok (refDiag : JsonRpc.Notification PublishDiagnosticsParams) :=
-        (Json.parse $ ←FS.readFile "content_diag.json") >>= fromJson?
-      then
-        assert! (diag == refDiag)
-      else
-        throw $ userError "Failed parsing test file."
-
-      Ipc.shutdown 2
-      discard $ Ipc.waitForExit
+    Ipc.shutdown 2
+    discard $ Ipc.waitForExit

tests/lean/server/edits.lean

@@ -16,19 +16,16 @@ def main : IO Unit := do
     hIn.write (←FS.readBinFile "content_changes.log")
     hIn.flush
 
-    let diags ← Ipc.collectDiagnostics 2 "file:///test.lean" 7
-    if diags.isEmpty then
-      throw $ userError "Test failed, no diagnostics received."
+    let some diag ← Ipc.collectDiagnostics 2 "file:///test.lean" 7
+      | throw $ userError "Test failed, no diagnostics received."
+    FS.writeFile "edits_diag.json.produced" (toString <| toJson (diag : JsonRpc.Message))
+
+    if let Except.ok (refDiag : JsonRpc.Notification PublishDiagnosticsParams) :=
+      (Json.parse $ ←FS.readFile "edits_diag.json") >>= fromJson?
+    then
+      assert! (diag == refDiag)
     else
-      let diag := diags.getLast!
-      FS.writeFile "edits_diag.json.produced" (toString <| toJson (diag : JsonRpc.Message))
+      throw $ userError "Failed parsing test file."
 
-      if let Except.ok (refDiag : JsonRpc.Notification PublishDiagnosticsParams) :=
-        (Json.parse $ ←FS.readFile "edits_diag.json") >>= fromJson?
-      then
-        assert! (diag == refDiag)
-      else
-        throw $ userError "Failed parsing test file."
-
-      Ipc.shutdown 3
-      discard $ Ipc.waitForExit
+    Ipc.shutdown 3
+    discard $ Ipc.waitForExit

tests/lean/server/init_exit_worker.lean

@@ -9,8 +9,8 @@ def main : IO Unit := do
 
     hIn.write (←FS.readBinFile "open_empty.log")
     hIn.flush
-    let diags ← Ipc.collectDiagnostics 1 "file:///empty" 1
-    assert! diags.length >= 1
+    let diags? ← Ipc.collectDiagnostics 1 "file:///empty" 1
+    assert! diags?.isSome
 
     Ipc.writeNotification ⟨"exit", Json.null⟩
     discard Ipc.waitForExit