From ceeaff75684ad5e55fb04c30d0b540a39e0a7ef5 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 12 Oct 2023 16:08:28 +0200
Subject: [PATCH 01/60] feat: start integrating new snapshot types into cmdline
 driver

---
 src/Lean/Elab/Frontend.lean    |  29 +---
 src/Lean/Language/Basic.lean   | 151 ++++++++++++++++++
 src/Lean/Language/Lean.lean    | 279 +++++++++++++++++++++++++++++++++
 src/Lean/Server/Snapshots.lean |   6 -
 4 files changed, 438 insertions(+), 27 deletions(-)
 create mode 100644 src/Lean/Language/Basic.lean
 create mode 100644 src/Lean/Language/Lean.lean

diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index 080687cae980..99d84d391b54 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -3,8 +3,7 @@ Copyright (c) 2019 Microsoft Corporation. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura, Sebastian Ullrich
 -/
-import Lean.Elab.Import
-import Lean.Elab.Command
+import Lean.Language.Lean
 import Lean.Util.Profile
 import Lean.Server.References
 
@@ -85,12 +84,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)
@@ -100,26 +95,18 @@ def runFrontend
     (trustLevel : UInt32 := 0)
     (ileanFileName? : Option String := none)
     : IO (Environment × Bool) := do
+  let process ← Lean.Language.Lean.mkProcessor { mainModuleName, opts }
   let inputCtx := Parser.mkInputContext input fileName
-  let (header, parserState, messages) ← Parser.parseHeader inputCtx
-  let (env, messages) ← processHeader 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 }
-
-  let s ← IO.processCommands inputCtx parserState commandState
-  for msg in s.commandState.messages.toList do
-    IO.print (← msg.toString (includeEndPos := getPrintMessageEndPos opts))
-
+  let snaps ← process inputCtx
+  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 : Lean.Server.Ilean }
     IO.FS.writeFile ileanFileName $ Json.compress $ toJson ilean
 
-  pure (s.commandState.env, !s.commandState.messages.hasErrors)
+  let hasErrors := snaps.getAll.any (·.msgLog.hasErrors)
+  let env := sorry  -- TODO
+  pure (env, !hasErrors)
 
 end Lean.Elab
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
new file mode 100644
index 000000000000..1925bb99079a
--- /dev/null
+++ b/src/Lean/Language/Basic.lean
@@ -0,0 +1,151 @@
+/-
+Copyright (c) 2023 Lean FRO. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Authors: Sebastian Ullrich
+-/
+
+import Lean.Message
+import Lean.Parser.Types
+
+namespace Lean.Language
+
+/--
+  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. -/
+  msgLog : MessageLog
+  /-- General elaboration metadata produced by this step. -/
+  infoTree? : Option Elab.InfoTree := none
+  -- (`InfoTree` is quite Lean-specific at this point, but we want to make it more generic)
+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
+  task : Task α
+deriving Nonempty
+
+def SnapshotTask.ofIO (range : String.Range) (act : BaseIO α) : BaseIO (SnapshotTask α) := do
+  return {
+    range
+    task := (← BaseIO.asTask act)
+  }
+
+def SnapshotTask.pure (a : α) : SnapshotTask α where
+  -- irrelevant when already finished
+  range := default
+  task := .pure a
+
+def SnapshotTask.cancel (t : SnapshotTask α) : BaseIO Unit :=
+  IO.cancel t.task
+
+def SnapshotTask.map (t : SnapshotTask α) (f : α → β) : SnapshotTask β :=
+  { t with task := t.task.map f }
+
+/--
+  Chain 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) :
+    BaseIO (SnapshotTask β) :=
+  return { range, task := (← BaseIO.bindTask t.task fun a => (·.task) <$> (act a)) }
+
+/--
+  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
+  | mk (element : Snapshot) (children : Array (SnapshotTask SnapshotTree))
+deriving Nonempty
+abbrev SnapshotTree.element : SnapshotTree → Snapshot
+  | mk s _ => s
+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
+  toSnapshotTree : α → SnapshotTree
+export ToSnapshotTree (toSnapshotTree)
+
+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. -/
+partial def SnapshotTree.runAndReport (s : SnapshotTree) (opts : Options) : IO Unit := do
+  s.element.msgLog.forM (·.toString (includeEndPos := printMessageEndPos.get opts) >>= IO.println)
+  for t in s.children do
+    t.task.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.task.get
+
+/-- Reports `IO` exceptions as single snapshot message. -/
+def withFatalExceptions (ex : Snapshot → α) (act : IO α) : BaseIO α := do
+  match (← act.toBaseIO) with
+  | .error e => return ex {
+    msgLog := MessageLog.empty.add { fileName := "TODO", pos := ⟨0, 0⟩, data := e.toString }
+  }
+  | .ok a => return a
+
+def get? (t? : Option (SnapshotTask α)) : BaseIO (Option α) := do
+  let some t := t? | return none
+  return if (← IO.hasFinished t.task) then t.task.get else none
+
+def getOrCancel (t? : Option (SnapshotTask α)) : BaseIO (Option α) := do
+  let some t := t? | return none
+  if (← IO.hasFinished t.task) then
+    return t.task.get
+  IO.cancel t.task
+  return none
+
+/--
+  Checks whether reparsing `stx` in `newSource` should result in `stx` again.
+  `stopPos` is the furthest position the parser has looked at for producing
+  `stx`, which usually is the extent of the subsequent token (if any) plus one. -/
+-- This is a fundamentally different approach from the "go up one snapshot from
+-- the point of change" in the old implementation which was failing in edge
+-- cases and is not generalizable to a DAG of snapshots.
+def unchangedParse (stx : Syntax) (stopPos : String.Pos) (newSource : String) : Bool :=
+  if let .original start .. := stx.getHeadInfo then
+    let oldSubstr := { start with stopPos }
+    oldSubstr == { oldSubstr with str := newSource }
+  else false
+
+/-- Metadata that does not change during the lifetime of the language processing process. -/
+structure ProcessingContext where
+  mainModuleName : Name
+  opts : Options
+
+structure Language where
+  mkProcessor : ProcessingContext → BaseIO (Parser.InputContext → BaseIO SnapshotTree)
+
+/-- Returns a function for processing a language using incremental snapshots. -/
+partial def mkIncrementalProcessor [ToSnapshotTree α]
+  (process : ProcessingContext → Option α → Parser.InputContext → BaseIO α) :
+    ProcessingContext → BaseIO (Parser.InputContext → BaseIO SnapshotTree) := fun ctx => do
+  let oldRef ← IO.mkRef none
+  return fun doc => do
+    let snap ← process ctx (← oldRef.get) doc
+    oldRef.set (some snap)
+    return toSnapshotTree snap
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
new file mode 100644
index 000000000000..5a26ce8e89ef
--- /dev/null
+++ b/src/Lean/Language/Lean.lean
@@ -0,0 +1,279 @@
+/-
+Copyright (c) 2023 Lean FRO. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+
+Authors: Sebastian Ullrich
+-/
+
+import Lean.Language.Basic
+import Lean.Parser.Module
+import Lean.Elab.Command
+import Lean.Elab.Import
+
+namespace Lean.Language.Lean
+open Lean.Elab
+open Lean.Parser
+
+def pushOpt (a? : Option α) (as : Array α) : Array α :=
+  match a? with
+  | some a => as.push a
+  | none   => as
+
+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"
+}
+
+/-! The hierarchy of Lean snapshot types -/
+
+/-- Final state of processing of a command. -/
+structure CommandFinishedSnapshot extends Snapshot where
+  cmdState : Command.State
+deriving Nonempty
+instance : ToSnapshotTree CommandFinishedSnapshot where
+  toSnapshotTree s := ⟨s.toSnapshot, #[]⟩
+
+/-- State after execution of a single synchronous tactic step. -/
+inductive TacticEvaluatedSnapshot where
+  | mk (toSnapshot : Snapshot) (next : Array (SnapshotTask TacticEvaluatedSnapshot))
+/-- Potential, potentially parallel, follow-up tactic executions. -/
+-- In the first, non-parallel version, each task will depend on its predecessor
+abbrev TacticEvaluatedSnapshot.next : TacticEvaluatedSnapshot → Array (SnapshotTask TacticEvaluatedSnapshot)
+  | mk _ next => next
+partial instance : ToSnapshotTree TacticEvaluatedSnapshot where
+  toSnapshotTree := go where
+    go := fun ⟨s, next⟩ => ⟨s, next.map (·.map go)⟩
+
+/--
+  State after processing a command's signature and before executing its tactic
+  body, if any. Other commands should immediately proceed to `finished`. -/
+structure CommandSignatureProcessedSnapshot extends Snapshot where
+  /-- Potential, potentially parallel, follow-up tactic executions. -/
+  tacs : Array (SnapshotTask TacticEvaluatedSnapshot)
+  /-- State after processing is finished. -/
+  -- If we make proofs completely opaque, this will not have to depend on `tacs`
+  finished : SnapshotTask CommandFinishedSnapshot
+deriving Nonempty
+instance : ToSnapshotTree CommandSignatureProcessedSnapshot where
+  toSnapshotTree s := ⟨s.toSnapshot, s.tacs.map (·.map toSnapshotTree) |>.push (s.finished.map toSnapshotTree)⟩
+
+/-- State after a command has been parsed. -/
+structure CommandParsedSnapshotData extends Snapshot where
+  stx : Syntax
+  parserState : Parser.ModuleParserState
+  /-- Furthest position the parser has looked at; used for incremental parsing. -/
+  stopPos : String.Pos
+  sig : SnapshotTask CommandSignatureProcessedSnapshot
+deriving Nonempty
+inductive CommandParsedSnapshot where
+  | mk (data : CommandParsedSnapshotData) (next? : Option (SnapshotTask CommandParsedSnapshot))
+deriving Nonempty
+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.sig.map toSnapshotTree] |> pushOpt (s.next?.map (·.map go))⟩
+
+structure HeaderProcessedSucessfully where
+  cmdState : Command.State
+  next : SnapshotTask CommandParsedSnapshot
+
+/-- State after the module header has been processed including imports. -/
+structure HeaderProcessedSnapshot extends Snapshot where
+  success? : Option HeaderProcessedSucessfully
+instance : ToSnapshotTree HeaderProcessedSnapshot where
+  toSnapshotTree s := ⟨s.toSnapshot, #[] |> pushOpt (s.success?.map (·.next.map toSnapshotTree))⟩
+
+structure HeaderParsedSucessfully where
+  parserState : Parser.ModuleParserState
+  /-- Furthest position the parser has looked at; used for incremental parsing. -/
+  stopPos : String.Pos
+  next : SnapshotTask HeaderProcessedSnapshot
+
+/-- State after the module header has been parsed. -/
+structure HeaderParsedSnapshot extends Snapshot where
+  stx : Syntax
+  success? : Option HeaderParsedSucessfully
+instance : ToSnapshotTree HeaderParsedSnapshot where
+  toSnapshotTree s := ⟨s.toSnapshot, #[] |> pushOpt (s.success?.map (·.next.map toSnapshotTree))⟩
+
+abbrev InitialSnapshot := HeaderParsedSnapshot
+
+/-- Entry point of the Lean language processor. -/
+-- As a general note, for each processing function we pass in the previous
+-- state, if any, in order to reuse still-valid state. Thus the logic of
+-- what snapshots to reuse is completely moved out of the server into the
+-- language-specific processor. As there is no cheap way to check whether
+-- the `Environment` is unchanged, we must make sure to pass `none` as all
+-- follow-up "previous states" when we do change it.
+partial def processLean
+  (ctx : ProcessingContext) (old? : Option InitialSnapshot) (ictx : Parser.InputContext) :
+    BaseIO InitialSnapshot :=
+  parseHeader old?
+where
+  parseHeader (old? : Option HeaderParsedSnapshot) := do
+    let unchanged old success :=
+      -- when header syntax is unchanged, reuse import processing task as is
+      return { old with success? := some { success with
+        next := (← success.next.bindIO (processHeader · old.stx success.parserState)) } }
+
+    -- fast path: if we have parsed the header sucessfully...
+    if let some old@{ success? := some success, .. } := old? then
+      -- ...and the header syntax appears unchanged...
+      if unchangedParse old.stx success.stopPos ictx.input then
+        -- ...go immediately to next snapshot
+        return (← unchanged old success)
+
+    withFatalExceptions ({ · with stx := .missing, success? := none }) do
+      let (stx, parserState, msgLog) ← Parser.parseHeader ictx
+      -- TODO: move into `parseHeader`; it should add the length of `peekToken`,
+      -- which is the full extent the parser considered before stopping
+      let stopPos := parserState.pos + (⟨1⟩ : String.Pos)
+
+      if msgLog.hasErrors then
+        return { stx, msgLog, success? := none }
+
+      -- semi-fast path: go to next snapshot if syntax tree is unchanged
+      if let some old@{ success? := some success, .. } := old? then
+        if old.stx == stx then
+          return (← unchanged old success)
+        success.next.cancel
+
+      return { stx, msgLog, success? := some {
+        parserState
+        stopPos
+        next := (← processHeader none stx parserState) } }
+
+  processHeader (old? : Option HeaderProcessedSnapshot) (stx : Syntax) (parserState : Parser.ModuleParserState) := do
+    -- fast path, do not even start new task for this snapshot
+    if let some old := old? then
+      if let some success := old.success? then
+        return .pure { old with success? := some { success with
+          next := (← success.next.bindIO (parseCmd · parserState success.cmdState)) } }
+      else
+        return .pure old
+
+    SnapshotTask.ofIO ⟨0, ictx.input.endPos⟩ do
+    withFatalExceptions ({ · with success? := none }) do
+    -- discard existing continuation if any, there is nothing to reuse
+    let _ ← old?.bind (·.success?) |>.map (·.next) |> getOrCancel
+    -- function copied from current implementation, see below
+    -- TODO: we should do this at most once per process
+    let cmdState ← doImport stx
+    return {
+      msgLog := cmdState.messages
+      infoTree? := cmdState.infoState.trees[0]!
+      success? := some { cmdState, next := (← parseCmd none parserState cmdState) } }
+
+  parseCmd (old? : Option CommandParsedSnapshot) (parserState : Parser.ModuleParserState) (cmdState : Command.State) :
+      BaseIO (SnapshotTask CommandParsedSnapshot) := do
+    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)
+          (next? := (← old.data.sig.bindIO fun sig =>
+            sig.finished.bindIO fun finished => do
+              parseCmd (← getOrCancel oldNext) old.data.parserState finished.cmdState))
+      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 unchangedParse old.data.stx old.data.stopPos ictx.input then
+        return .pure (← unchanged old)
+
+    SnapshotTask.ofIO ⟨parserState.pos, ictx.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 ictx pmctx parserState cmdState.messages
+
+      -- TODO: move into `parseCommand`; it should add the length of `peekToken`,
+      -- which is the full extent the parser considered before stopping
+      let stopPos := parserState.pos + (⟨1⟩ : String.Pos)
+      if let some old := old? then
+        if old.data.stx == stx then
+          return (← unchanged old)
+
+      -- signature elaboration task; for now, does full elaboration
+      -- TODO: do tactic snapshots, reuse old state for them
+      let _ ← old?.map (·.data.sig) |> getOrCancel
+      let sig ← SnapshotTask.ofIO (stx.getRange?.getD ⟨parserState.pos, parserState.pos⟩) do
+        let cmdState ← doElab stx cmdState beginPos scope
+        return {
+          msgLog := cmdState.messages
+          infoTree? := some cmdState.infoState.trees[0]!
+          -- TODO
+          tacs := #[]
+          finished := .pure {
+            msgLog := .empty
+            cmdState
+          }
+        }
+
+      return .mk (data := {
+        msgLog
+        stx
+        parserState
+        stopPos
+        sig
+      }) (next? :=
+          if Parser.isTerminalCommand stx then none
+          -- for now, wait on "command finished" snapshot before parsing next command
+          else some (← sig.bindIO fun sig => do
+            sig.finished.bindIO fun finished =>
+              parseCmd none parserState finished.cmdState))
+
+  doImport stx := do
+    let (headerEnv, msgLog) ← Elab.processHeader stx ctx.opts .empty ictx
+    let mut headerEnv := headerEnv
+    headerEnv := headerEnv.setMainModule ctx.mainModuleName
+    let cmdState := Elab.Command.mkState headerEnv msgLog ctx.opts
+    return { cmdState with infoState := {
+      enabled := true
+      trees := #[Elab.InfoTree.context ({
+        env     := headerEnv
+        fileMap := ictx.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'
+    }}
+
+  doElab stx cmdState beginPos scope := do
+    let cmdStateRef ← IO.mkRef { cmdState with messages := .empty }
+    let cmdCtx : Elab.Command.Context := { ictx with
+      cmdPos       := beginPos
+      tacticCache? := none
+    }
+    let (output, _) ← IO.FS.withIsolatedStreams (isolateStderr := server.stderrAsMessages.get scope.opts) <| liftM (m := BaseIO) do
+      Elab.Command.catchExceptions
+        (getResetInfoTrees *> Elab.Command.elabCommandTopLevel stx)
+        cmdCtx cmdStateRef
+    let mut postCmdState ← cmdStateRef.get
+    if !output.isEmpty then
+      postCmdState := {
+        postCmdState with
+        messages := postCmdState.messages.add {
+          fileName := ictx.fileName
+          severity := MessageSeverity.information
+          pos      := ictx.fileMap.toPosition beginPos
+          data     := output
+        }
+      }
+    return postCmdState
+
+end Lean
+
+def Lean : Language where
+  mkProcessor := mkIncrementalProcessor Lean.processLean
diff --git a/src/Lean/Server/Snapshots.lean b/src/Lean/Server/Snapshots.lean
index f49e5db56449..996cff4add3d 100644
--- a/src/Lean/Server/Snapshots.lean
+++ b/src/Lean/Server/Snapshots.lean
@@ -92,12 +92,6 @@ def parseNextCmd (inputCtx : Parser.InputContext) (snap : Snapshot) : IO Syntax
     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

From d0fd3d7e2b9f423ba10447f8f8c89db0da7e6b3a Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 13 Oct 2023 14:28:47 +0200
Subject: [PATCH 02/60] feat: switch `Language` to bundled, integrate into
 cmdline driver in full

---
 src/Lean/Elab/Frontend.lean    | 13 +++++++++----
 src/Lean/Language/Basic.lean   | 25 +++++++++++++++++--------
 src/Lean/Language/Lean.lean    | 13 ++++++++++++-
 src/Lean/Server/Snapshots.lean |  3 ++-
 4 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index 99d84d391b54..799d5f05ec99 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -92,12 +92,15 @@ def runFrontend
     (opts : Options)
     (fileName : String)
     (mainModuleName : Name)
-    (trustLevel : UInt32 := 0)
+    -- TODO: do we still want this in the driver?
+    (_trustLevel : UInt32 := 0)
     (ileanFileName? : Option String := none)
     : IO (Environment × Bool) := do
-  let process ← Lean.Language.Lean.mkProcessor { mainModuleName, opts }
   let inputCtx := Parser.mkInputContext input fileName
-  let snaps ← process inputCtx
+  -- TODO: replace with `#lang` processing
+  let lang := Language.Lean
+  let snap ← lang.process { mainModuleName, opts } none inputCtx
+  let snaps := Language.toSnapshotTree snap
   snaps.runAndReport opts
   if let some ileanFileName := ileanFileName? then
     let trees := snaps.getAll.concatMap (match ·.infoTree? with | some t => #[t] | _ => #[])
@@ -106,7 +109,9 @@ def runFrontend
     IO.FS.writeFile ileanFileName $ Json.compress $ toJson ilean
 
   let hasErrors := snaps.getAll.any (·.msgLog.hasErrors)
-  let env := sorry  -- TODO
+  -- 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 := lang.getFinalEnv? snap |>.getD (← mkEmptyEnvironment)
   pure (env, !hasErrors)
 
 end Lean.Elab
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index 1925bb99079a..be3a000ca28a 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -58,6 +58,9 @@ def SnapshotTask.bindIO (t : SnapshotTask α) (act : α → BaseIO (SnapshotTask
     BaseIO (SnapshotTask β) :=
   return { range, task := (← BaseIO.bindTask t.task fun a => (·.task) <$> (act a)) }
 
+def SnapshotTask.get (t : SnapshotTask α) : α :=
+  t.task.get
+
 /--
   Tree of snapshots where each snapshot comes with an array of asynchronous
   further subtrees. Used for asynchronously collecting information about the
@@ -89,7 +92,7 @@ register_builtin_option printMessageEndPos : Bool := {
 partial def SnapshotTree.runAndReport (s : SnapshotTree) (opts : Options) : IO Unit := do
   s.element.msgLog.forM (·.toString (includeEndPos := printMessageEndPos.get opts) >>= IO.println)
   for t in s.children do
-    t.task.get.runAndReport opts
+    t.get.runAndReport opts
 
 /-- Waits on and returns all snapshots in the tree. -/
 partial def SnapshotTree.getAll (s : SnapshotTree) : Array Snapshot :=
@@ -98,7 +101,7 @@ where
   go s : StateM (Array Snapshot) Unit := do
     modify (·.push s.element)
     for t in s.children do
-      go t.task.get
+      go t.get
 
 /-- Reports `IO` exceptions as single snapshot message. -/
 def withFatalExceptions (ex : Snapshot → α) (act : IO α) : BaseIO α := do
@@ -110,12 +113,12 @@ def withFatalExceptions (ex : Snapshot → α) (act : IO α) : BaseIO α := do
 
 def get? (t? : Option (SnapshotTask α)) : BaseIO (Option α) := do
   let some t := t? | return none
-  return if (← IO.hasFinished t.task) then t.task.get else none
+  return if (← IO.hasFinished t.task) then t.get else none
 
 def getOrCancel (t? : Option (SnapshotTask α)) : BaseIO (Option α) := do
   let some t := t? | return none
   if (← IO.hasFinished t.task) then
-    return t.task.get
+    return t.get
   IO.cancel t.task
   return none
 
@@ -138,14 +141,20 @@ structure ProcessingContext where
   opts : Options
 
 structure Language where
-  mkProcessor : ProcessingContext → BaseIO (Parser.InputContext → BaseIO SnapshotTree)
+  InitialSnapshot : Type
+  [instToSnapshotTree : ToSnapshotTree InitialSnapshot]
+  process : ProcessingContext → (old? : Option InitialSnapshot) → Parser.InputContext → BaseIO InitialSnapshot
+  -- TODO: is this the right interface for other languages as well?
+  /-- Gets final environment, if any, that is to be used for persisting, code generation, etc. -/
+  getFinalEnv? : InitialSnapshot → Option Environment
+
+instance (lang : Language) : ToSnapshotTree lang.InitialSnapshot := lang.instToSnapshotTree
 
 /-- Returns a function for processing a language using incremental snapshots. -/
-partial def mkIncrementalProcessor [ToSnapshotTree α]
-  (process : ProcessingContext → Option α → Parser.InputContext → BaseIO α) :
+partial def Language.mkIncrementalProcessor (lang : Language) :
     ProcessingContext → BaseIO (Parser.InputContext → BaseIO SnapshotTree) := fun ctx => do
   let oldRef ← IO.mkRef none
   return fun doc => do
-    let snap ← process ctx (← oldRef.get) doc
+    let snap ← lang.process ctx (← oldRef.get) doc
     oldRef.set (some snap)
     return toSnapshotTree snap
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 5a26ce8e89ef..326c1e1b9bb6 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -273,7 +273,18 @@ where
       }
     return postCmdState
 
+partial def getFinalEnv? (snap : InitialSnapshot) : Option Environment := do
+  let snap ← snap.success?
+  let snap ← snap.next.get.success?
+  goCmd snap.next.get
+where goCmd snap :=
+  if let some next := snap.next? then
+    goCmd next.get
+  else
+    snap.data.sig.get.finished.get.cmdState.env
+
 end Lean
 
 def Lean : Language where
-  mkProcessor := mkIncrementalProcessor Lean.processLean
+  process := Lean.processLean
+  getFinalEnv? := Lean.getFinalEnv?
diff --git a/src/Lean/Server/Snapshots.lean b/src/Lean/Server/Snapshots.lean
index 996cff4add3d..65d84907dd68 100644
--- a/src/Lean/Server/Snapshots.lean
+++ b/src/Lean/Server/Snapshots.lean
@@ -8,6 +8,7 @@ import Init.System.IO
 
 import Lean.Elab.Import
 import Lean.Elab.Command
+import Lean.Language.Lean
 
 import Lean.Widget.InteractiveDiagnostic
 
@@ -115,7 +116,7 @@ def compileNextCmd (inputCtx : Parser.InputContext) (snap : Snapshot) (hasWidget
     fileMap      := inputCtx.fileMap
     tacticCache? := some tacticCacheNew
   }
-  let (output, _) ← IO.FS.withIsolatedStreams (isolateStderr := server.stderrAsMessages.get scope.opts) <| liftM (m := BaseIO) do
+  let (output, _) ← IO.FS.withIsolatedStreams (isolateStderr := Language.Lean.server.stderrAsMessages.get scope.opts) <| liftM (m := BaseIO) do
     Elab.Command.catchExceptions
       (getResetInfoTrees *> Elab.Command.elabCommandTopLevel cmdStx)
       cmdCtx cmdStateRef

From 152d5949b8c2398f5eba7aa934ba790fb7e74cd4 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 13 Oct 2023 15:47:07 +0200
Subject: [PATCH 03/60] fix: do not duplicate output

---
 src/Lean/Language/Basic.lean | 2 +-
 src/Lean/Language/Lean.lean  | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index be3a000ca28a..0897eafbc513 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -90,7 +90,7 @@ register_builtin_option printMessageEndPos : Bool := {
   Runs a tree of snapshots to conclusion and incrementally report messages
   on stdout. Messages are reported in tree preorder. -/
 partial def SnapshotTree.runAndReport (s : SnapshotTree) (opts : Options) : IO Unit := do
-  s.element.msgLog.forM (·.toString (includeEndPos := printMessageEndPos.get opts) >>= IO.println)
+  s.element.msgLog.forM (·.toString (includeEndPos := printMessageEndPos.get opts) >>= IO.print)
   for t in s.children do
     t.get.runAndReport opts
 
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 326c1e1b9bb6..b500a34bbc08 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -190,7 +190,7 @@ where
       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 ictx pmctx parserState cmdState.messages
+      let (stx, parserState, msgLog) := Parser.parseCommand ictx pmctx parserState .empty
 
       -- TODO: move into `parseCommand`; it should add the length of `peekToken`,
       -- which is the full extent the parser considered before stopping
@@ -211,7 +211,7 @@ where
           tacs := #[]
           finished := .pure {
             msgLog := .empty
-            cmdState
+            cmdState := { cmdState with messages := .empty }
           }
         }
 

From a68d1c62140b6585c3c32e6be682470e1c7d1433 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 19 Oct 2023 18:48:51 +0200
Subject: [PATCH 04/60] fix: more output fixes

---
 src/Lean/Elab/Command.lean     | 10 ------
 src/Lean/Language/Lean.lean    | 60 +++++++++++++++++++---------------
 src/Lean/Server/Snapshots.lean |  2 +-
 3 files changed, 35 insertions(+), 37 deletions(-)

diff --git a/src/Lean/Elab/Command.lean b/src/Lean/Elab/Command.lean
index 4cb79e5c5db0..086537f82a28 100644
--- a/src/Lean/Elab/Command.lean
+++ b/src/Lean/Elab/Command.lean
@@ -265,11 +265,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
@@ -317,11 +312,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
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index b500a34bbc08..cd0b20967451 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -19,12 +19,17 @@ def pushOpt (a? : Option α) (as : Array α) : Array α :=
   | some a => as.push a
   | none   => as
 
-register_builtin_option server.stderrAsMessages : Bool := {
-  defValue := true
+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"
 }
 
+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. -/
@@ -203,30 +208,30 @@ where
       -- TODO: do tactic snapshots, reuse old state for them
       let _ ← old?.map (·.data.sig) |> getOrCancel
       let sig ← SnapshotTask.ofIO (stx.getRange?.getD ⟨parserState.pos, parserState.pos⟩) do
-        let cmdState ← doElab stx cmdState beginPos scope
+        let cmdState ← doElab stx cmdState msgLog.hasErrors beginPos scope
         return {
-          msgLog := cmdState.messages
-          infoTree? := some cmdState.infoState.trees[0]!
+          msgLog := .empty
           -- TODO
           tacs := #[]
           finished := .pure {
-            msgLog := .empty
-            cmdState := { cmdState with messages := .empty }
+            msgLog := cmdState.messages
+            infoTree? := some cmdState.infoState.trees[0]!
+            cmdState
           }
         }
 
-      return .mk (data := {
+      let next? ← if Parser.isTerminalCommand stx then pure none
+        -- for now, wait on "command finished" snapshot before parsing next command
+        else some <$> sig.bindIO fun sig => do
+          sig.finished.bindIO fun finished =>
+            parseCmd none parserState finished.cmdState
+      return .mk (next? := next?) {
         msgLog
         stx
         parserState
         stopPos
         sig
-      }) (next? :=
-          if Parser.isTerminalCommand stx then none
-          -- for now, wait on "command finished" snapshot before parsing next command
-          else some (← sig.bindIO fun sig => do
-            sig.finished.bindIO fun finished =>
-              parseCmd none parserState finished.cmdState))
+      }
 
   doImport stx := do
     let (headerEnv, msgLog) ← Elab.processHeader stx ctx.opts .empty ictx
@@ -250,28 +255,31 @@ where
       )].toPArray'
     }}
 
-  doElab stx cmdState beginPos scope := do
+  doElab stx cmdState hasParseError beginPos scope := do
     let cmdStateRef ← IO.mkRef { cmdState with messages := .empty }
     let cmdCtx : Elab.Command.Context := { ictx with
       cmdPos       := beginPos
       tacticCache? := none
     }
-    let (output, _) ← IO.FS.withIsolatedStreams (isolateStderr := server.stderrAsMessages.get scope.opts) <| liftM (m := BaseIO) do
+    let (output, _) ← IO.FS.withIsolatedStreams (isolateStderr := stderrAsMessages.get scope.opts) <| liftM (m := BaseIO) do
       Elab.Command.catchExceptions
         (getResetInfoTrees *> Elab.Command.elabCommandTopLevel stx)
         cmdCtx cmdStateRef
-    let mut postCmdState ← cmdStateRef.get
+    let postCmdState ← cmdStateRef.get
+    let mut msgs := postCmdState.messages
+    -- `stx.hasMissing` should imply `initMsgs.hasErrors`, but the latter should be cheaper to check in general
+    if !showPartialSyntaxErrors.get postCmdState.scopes[0]!.opts && hasParseError && 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)⟩
     if !output.isEmpty then
-      postCmdState := {
-        postCmdState with
-        messages := postCmdState.messages.add {
-          fileName := ictx.fileName
-          severity := MessageSeverity.information
-          pos      := ictx.fileMap.toPosition beginPos
-          data     := output
-        }
+      msgs := msgs.add {
+        fileName := ictx.fileName
+        severity := MessageSeverity.information
+        pos      := ictx.fileMap.toPosition beginPos
+        data     := output
       }
-    return postCmdState
+    return { postCmdState with messages := msgs }
 
 partial def getFinalEnv? (snap : InitialSnapshot) : Option Environment := do
   let snap ← snap.success?
diff --git a/src/Lean/Server/Snapshots.lean b/src/Lean/Server/Snapshots.lean
index 65d84907dd68..95fb8342b236 100644
--- a/src/Lean/Server/Snapshots.lean
+++ b/src/Lean/Server/Snapshots.lean
@@ -116,7 +116,7 @@ def compileNextCmd (inputCtx : Parser.InputContext) (snap : Snapshot) (hasWidget
     fileMap      := inputCtx.fileMap
     tacticCache? := some tacticCacheNew
   }
-  let (output, _) ← IO.FS.withIsolatedStreams (isolateStderr := Language.Lean.server.stderrAsMessages.get scope.opts) <| liftM (m := BaseIO) do
+  let (output, _) ← IO.FS.withIsolatedStreams (isolateStderr := Language.Lean.stderrAsMessages.get? scope.opts |>.getD true) <| liftM (m := BaseIO) do
     Elab.Command.catchExceptions
       (getResetInfoTrees *> Elab.Command.elabCommandTopLevel cmdStx)
       cmdCtx cmdStateRef

From f2f6016ce2e0e1de62f3d3de06081d00dfe41613 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 3 Nov 2023 16:41:52 +0100
Subject: [PATCH 05/60] foo

---
 src/Lean/Language/Basic.lean          |   4 +-
 src/Lean/Language/Lean.lean           |   2 +-
 src/Lean/Server/FileWorker.lean       | 219 ++++++--------------------
 src/Lean/Server/FileWorker/Utils.lean |  26 ++-
 src/Lean/Server/Snapshots.lean        |  82 +---------
 5 files changed, 75 insertions(+), 258 deletions(-)

diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index 0897eafbc513..de82affa9693 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -152,9 +152,9 @@ instance (lang : Language) : ToSnapshotTree lang.InitialSnapshot := lang.instToS
 
 /-- Returns a function for processing a language using incremental snapshots. -/
 partial def Language.mkIncrementalProcessor (lang : Language) :
-    ProcessingContext → BaseIO (Parser.InputContext → BaseIO SnapshotTree) := fun ctx => do
+    ProcessingContext → BaseIO (Parser.InputContext → BaseIO lang.InitialSnapshot) := fun ctx => do
   let oldRef ← IO.mkRef none
   return fun doc => do
     let snap ← lang.process ctx (← oldRef.get) doc
     oldRef.set (some snap)
-    return toSnapshotTree snap
+    return snap
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index cd0b20967451..decb86ddb05e 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -293,6 +293,6 @@ where goCmd snap :=
 
 end Lean
 
-def Lean : Language where
+abbrev Lean : Language where
   process := Lean.processLean
   getFinalEnv? := Lean.getFinalEnv?
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index d9b7ab33f2b3..74cc22b0e0bf 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -14,9 +14,9 @@ import Lean.Data.Json.FromToJson
 
 import Lean.Util.Paths
 import Lean.LoadDynlib
+import Lean.Language.Basic
 
 import Lean.Server.Utils
-import Lean.Server.Snapshots
 import Lean.Server.AsyncList
 import Lean.Server.References
 
@@ -59,8 +59,8 @@ structure WorkerContext where
   hIn              : FS.Stream
   hOut             : FS.Stream
   hLog             : FS.Stream
-  headerTask       : Task (Except Error (Snapshot × SearchPath))
   initParams       : InitializeParams
+  processor        : Parser.InputContext → BaseIO Lean.Language.Lean.InitialSnapshot
   clientHasWidgets : Bool
 
 /-! # Asynchronous snapshot elaboration -/
@@ -73,73 +73,46 @@ section Elab
 
   -- 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
+      (trees : Array Elab.InfoTree) : IO Unit := do
     let references := findModuleRefs m.text trees (localVars := true)
     let param := { version := m.version, references : LeanIleanInfoParams }
     hOut.writeLspNotification { method, param }
 
-  private def publishIleanInfoUpdate : DocumentMeta → FS.Stream → Array Snapshot → IO Unit :=
+  private def publishIleanInfoUpdate : DocumentMeta → FS.Stream → Array Elab.InfoTree → IO Unit :=
     publishIleanInfo "$/lean/ileanInfoUpdate"
 
-  private def publishIleanInfoFinal : DocumentMeta → FS.Stream → Array Snapshot → IO Unit :=
+  private def publishIleanInfoFinal : DocumentMeta → FS.Stream → Array Elab.InfoTree → IO Unit :=
     publishIleanInfo "$/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
+  private partial def reportSnapshots (ctx : WorkerContext) (m : DocumentMeta) (snaps : Language.SnapshotTree)
+      (cancelTk : CancelToken) : IO Unit := do
+    discard <| go snaps #[] #[] fun _ itrees => do
       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
-
-  /-- 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
-      -- Treat 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 })
+      publishIleanInfoFinal m ctx.hOut itrees
+      return .pure <| .ok ()
+  where
+    go node diags itrees cont := do
+      if (← cancelTk.ref.get) then
+        return .pure <| .ok ()
+      let diags := diags ++ (← node.element.msgLog.toList.toArray.mapM (Widget.msgToInteractiveDiagnostic m.text · ctx.clientHasWidgets)).map (·.toDiagnostic)
+      publishDiagnostics m diags ctx.hOut
+      let itrees := match node.element.infoTree? with
+        | some itree => itrees.push itree
+        | none       => itrees
+      publishIleanInfoUpdate m ctx.hOut itrees
+      goSeq cont diags itrees node.children.toList
+    goSeq cont diags itrees
+      | [] => cont diags itrees
+      | t::ts => do
+        publishProgressAtPos m t.range.start ctx.hOut
+        IO.bindTask t.task fun node =>
+          go node diags itrees fun diags itrees =>
+            goSeq cont diags itrees 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 WorkerState where
@@ -150,6 +123,7 @@ structure WorkerState where
   /-- A map of RPC session IDs. We allow asynchronous elab tasks and request handlers
   to modify sessions. A single `Ref` ensures atomic transactions. -/
   rpcSessions     : RBMap UInt64 (IO.Ref RpcSession) compare
+  reportCancelTk  : CancelToken
 
 abbrev WorkerM := ReaderT WorkerContext <| StateRefT WorkerState IO
 
@@ -193,83 +167,30 @@ section Initialization
     | 3 => throwServerError s!"Imports are out of date and must be rebuilt; use the \"Restart File\" command in your editor.\n\n{stdout}"
     | _ => throwServerError s!"`{cmdStr}` failed:\n{stdout}\nstderr:\n{stderr}"
 
-  def compileHeader (m : DocumentMeta) (hOut : FS.Stream) (opts : Options) (hasWidgets : Bool)
-      : IO (Syntax × Task (Except Error (Snapshot × SearchPath))) := do
-    -- parsing should not take long, do synchronously
-    let (headerStx, headerParserState, msgLog) ← Parser.parseHeader m.mkInputContext
-    (headerStx, ·) <$> EIO.asTask do
-      let mut srcSearchPath ← initSrcSearchPath (← getBuildDir)
-      let lakePath ← match (← IO.getEnv "LAKE") with
-        | some path => pure <| System.FilePath.mk path
-        | none =>
-          let lakePath ← match (← IO.getEnv "LEAN_SYSROOT") with
-            | some path => pure <| System.FilePath.mk path / "bin" / "lake"
-            | _         => pure <| (← appDir) / "lake"
-          pure <| lakePath.withExtension System.FilePath.exeExtension
-      let (headerEnv, msgLog) ← try
-        if let some path := System.Uri.fileUriToPath? m.uri then
-          -- NOTE: we assume for now that `lakefile.lean` does not have any non-stdlib deps
-          -- NOTE: lake does not exist in stage 0 (yet?)
-          if path.fileName != "lakefile.lean" && (← System.FilePath.pathExists lakePath) then
-            let pkgSearchPath ← lakeSetupSearchPath lakePath m (Lean.Elab.headerToImports headerStx) hOut
-            srcSearchPath ← initSrcSearchPath (← getBuildDir) pkgSearchPath
-        Elab.processHeader headerStx opts msgLog m.mkInputContext
-      catch e =>  -- should be from `lake print-paths`
-        let msgs := MessageLog.empty.add { fileName := "<ignored>", pos := ⟨0, 0⟩, data := e.toString }
-        pure (← mkEmptyEnvironment, msgs)
-      let mut headerEnv := headerEnv
-      try
-        if let some path := System.Uri.fileUriToPath? m.uri then
-          headerEnv := headerEnv.setMainModule (← moduleNameOfFileName path none)
-      catch _ => pure ()
-      let cmdState := Elab.Command.mkState headerEnv msgLog opts
-      let cmdState := { cmdState with infoState := {
-        enabled := true
-        trees := #[Elab.InfoTree.context ({
-          env     := headerEnv
-          fileMap := m.text
-          ngen    := { namePrefix := `_worker }
-        }) (Elab.InfoTree.node
-            (Elab.Info.ofCommandInfo { elaborator := `header, stx := headerStx })
-            (headerStx[1].getArgs.toList.map (fun importStx =>
-              Elab.InfoTree.node (Elab.Info.ofCommandInfo {
-                elaborator := `import
-                stx := importStx
-              }) #[].toPArray'
-            )).toPArray'
-        )].toPArray'
-      }}
-      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, srcSearchPath)
-
   def initializeWorker (meta : DocumentMeta) (i 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 processor ← Language.Lean.mkIncrementalProcessor { opts, mainModuleName }
+    let initSnap ← processor meta.mkInputContext
     let ctx :=
       { hIn  := i
         hOut := o
         hLog := e
-        headerTask
         initParams
+        processor
         clientHasWidgets
       }
-    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 : EditableDocument := { meta, initSnap }
+    let reportCancelTk ← CancelToken.new
+    reportSnapshots ctx doc.meta (Language.ToSnapshotTree.toSnapshotTree doc.initSnap) reportCancelTk
     return (ctx,
-    { doc             := doc
-      initHeaderStx   := headerStx
+    { doc, reportCancelTk
+      initHeaderStx   := initSnap.stx
       pendingRequests := RBMap.empty
       rpcSessions     := RBMap.empty
     })
@@ -282,49 +203,11 @@ section Updates
   /-- Given the new document, updates editable doc state. -/
   def updateDocument (newMeta : DocumentMeta) : WorkerM Unit := 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
-    let headSnapTask := oldDoc.cmdSnaps.waitHead?
-    let newSnaps ← if initHeaderStx != newHeaderStx then
-      EIO.asTask (ε := ElabTaskError) (prio := .dedicated) do
-        IO.sleep ctx.initParams.editDelay.toUInt32
-        cancelTk.check
-        IO.Process.exit 2
-    else 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 changePos := oldDoc.meta.text.source.firstDiffPos newMeta.text.source
-      -- Ignore exceptions, we are only interested in the successful snapshots
-      let (cmdSnaps, _) ← oldDoc.cmdSnaps.getFinishedPrefix
-      -- 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
-      -- 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)
-    modify fun st => { st with doc := { meta := newMeta, cmdSnaps := AsyncList.delayed newSnaps, cancelTk } }
+    (← get).reportCancelTk.set
+    let initSnap ← ctx.processor newMeta.mkInputContext
+    let reportCancelTk ← CancelToken.new
+    reportSnapshots ctx newMeta (Language.ToSnapshotTree.toSnapshotTree initSnap) reportCancelTk
+    modify fun st => { st with doc := { meta := newMeta, initSnap }, reportCancelTk }
 end Updates
 
 /- Notifications are handled in the main thread. They may change global worker state
@@ -412,8 +295,8 @@ section MessageHandling
       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
+    let t ← IO.bindTask (.pure "") fun _ => do
+     let srcSearchPath := ∅
      let rc : RequestContext :=
        { rpcSessions := st.rpcSessions
          srcSearchPath
@@ -463,8 +346,6 @@ section MainLoop
       handleRequest id method (toJson params)
       mainLoop
     | Message.notification "exit" none =>
-      let doc := st.doc
-      doc.cancelTk.set
       return ()
     | Message.notification method (some params) =>
       handleNotification method (toJson params)
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index 9e637a0482c9..8f9536c98d30 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -4,19 +4,16 @@ Released under Apache 2.0 license as described in the file LICENSE.
 
 Authors: Wojciech Nawrocki, Marc Huisinga
 -/
+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)
@@ -51,11 +48,28 @@ 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
+  -- TODO: generalize to other languages by moving request handlers into `Language`
+  initSnap : Language.Lean.InitialSnapshot
 
 namespace EditableDocument
 
+partial def cmdSnaps (doc : EditableDocument) : AsyncList ElabTaskError Snapshot := Id.run do
+  let some headerParsed := doc.initSnap.success? | return .nil
+  .delayed <| headerParsed.next.task.bind fun headerProcessed => Id.run do
+    let some headerSuccess := headerProcessed.success? | return .pure <| .ok .nil
+    headerSuccess.next.task.bind go
+where go cmdParsed :=
+  cmdParsed.data.sig.task.bind fun sig =>
+    sig.finished.task.map fun finished =>
+      .ok <| .cons {
+        stx := cmdParsed.data.stx
+        mpState := cmdParsed.data.parserState
+        cmdState := finished.cmdState
+        interactiveDiags := .empty  -- TODO
+      } (match cmdParsed.next? with
+        | some next => .delayed <| next.task.bind go
+        | none => .nil)
+
 end EditableDocument
 
 structure RpcSession where
diff --git a/src/Lean/Server/Snapshots.lean b/src/Lean/Server/Snapshots.lean
index 95fb8342b236..f398bb2849da 100644
--- a/src/Lean/Server/Snapshots.lean
+++ b/src/Lean/Server/Snapshots.lean
@@ -26,11 +26,6 @@ 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
@@ -38,7 +33,6 @@ structure Snapshot where
   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
 
@@ -66,10 +60,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,76 +77,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
-
-/-- 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 := Language.Lean.stderrAsMessages.get? scope.opts |>.getD true) <| 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

From 986d1c141878b481d6b028cca35765396bedcabe Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 9 Nov 2023 11:21:53 +0100
Subject: [PATCH 06/60] fix: `run_task`/`deactivate_task` race condition on
 `m_imp->m_closure`

---
 src/runtime/object.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/runtime/object.cpp b/src/runtime/object.cpp
index 1e1276568674..204156f130db 100644
--- a/src/runtime/object.cpp
+++ b/src/runtime/object.cpp
@@ -714,8 +714,9 @@ class task_manager {
             resolve_core(t, v);
         } else {
             // `bind` task has not finished yet, re-add as dependency of nested task
+            object * c = t->m_imp->m_closure;
             lock.unlock();
-            add_dep(lean_to_task(closure_arg_cptr(t->m_imp->m_closure)[0]), t);
+            add_dep(lean_to_task(closure_arg_cptr(c)[0]), t);
             lock.lock();
         }
     }

From d2bf996c2923df67e894cd3c80206a7d9175ba25 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 10 Nov 2023 08:45:36 +0100
Subject: [PATCH 07/60] fix: re-read HTTP header when skipping notification in
 `Ipc.readResponseAs`

---
 src/Lean/Data/JsonRpc.lean |  3 +--
 src/Lean/Data/Lsp/Ipc.lean | 16 ++++++++++++++--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/Lean/Data/JsonRpc.lean b/src/Lean/Data/JsonRpc.lean
index ee731726102c..d0bba7fee330 100644
--- a/src/Lean/Data/JsonRpc.lean
+++ b/src/Lean/Data/JsonRpc.lean
@@ -305,7 +305,7 @@ section
         throw $ userError s!"Expected method '{expectedMethod}', got method '{method}'"
     | _ => throw $ userError s!"Expected JSON-RPC notification, got: '{(toJson m).compress}'"
 
-  partial def readResponseAs (h : FS.Stream) (nBytes : Nat) (expectedID : RequestID) (α) [FromJson α] : IO (Response α) := do
+  def readResponseAs (h : FS.Stream) (nBytes : Nat) (expectedID : RequestID) (α) [FromJson α] : IO (Response α) := do
     let m ← h.readMessage nBytes
     match m with
     | Message.response id result =>
@@ -315,7 +315,6 @@ section
         | Except.error inner => throw $ userError s!"Unexpected result '{result.compress}'\n{inner}"
       else
         throw $ userError s!"Expected id {expectedID}, got id {id}"
-    | Message.notification .. => readResponseAs h nBytes expectedID α
     | _ => throw $ userError s!"Expected JSON-RPC response, got: '{(toJson m).compress}'"
 end
 
diff --git a/src/Lean/Data/Lsp/Ipc.lean b/src/Lean/Data/Lsp/Ipc.lean
index 3a9d18af350e..4f07dbf11852 100644
--- a/src/Lean/Data/Lsp/Ipc.lean
+++ b/src/Lean/Data/Lsp/Ipc.lean
@@ -62,8 +62,20 @@ def readMessage : IpcM JsonRpc.Message := do
 def readRequestAs (expectedMethod : String) (α) [FromJson α] : IpcM (Request α) := do
   (←stdout).readLspRequestAs expectedMethod α
 
-def readResponseAs (expectedID : RequestID) (α) [FromJson α] : IpcM (Response α) := do
-  (←stdout).readLspResponseAs expectedID α
+/-- Reads response, discarding notifications in between (as this is purely for testing). -/
+partial def readResponseAs (expectedID : RequestID) (α) [FromJson α] :
+    IpcM (Response α) := do
+  let m ← (←stdout).readLspMessage
+  match m with
+  | Message.response id result =>
+    if id == expectedID then
+      match fromJson? result with
+      | Except.ok v => pure ⟨expectedID, v⟩
+      | Except.error inner => throw $ userError s!"Unexpected result '{result.compress}'\n{inner}"
+    else
+      throw $ userError s!"Expected id {expectedID}, got id {id}"
+  | .notification .. => readResponseAs expectedID α
+  | _ => throw $ userError s!"Expected JSON-RPC response, got: '{(toJson m).compress}'"
 
 def waitForExit : IpcM UInt32 := do
   (←read).wait

From 23da22e3925f7d74a04913d405dceab62e202593 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 10 Nov 2023 14:26:46 +0100
Subject: [PATCH 08/60] fix: debounce server status reporting

---
 src/Lean/Server/FileWorker.lean               | 52 +++++++++++++------
 .../unterminatedDocComment.lean.expected.out  |  2 -
 2 files changed, 37 insertions(+), 17 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 74cc22b0e0bf..c9d11543e3be 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -84,32 +84,54 @@ section Elab
   private def publishIleanInfoFinal : DocumentMeta → FS.Stream → Array Elab.InfoTree → IO Unit :=
     publishIleanInfo "$/lean/ileanInfoFinal"
 
+  private structure ReportSnapshotsState where
+    hasBlocked := false
+    diagnostics : Array Lsp.Diagnostic := #[]
+    infoTrees : Array Elab.InfoTree := #[]
+
+  /--
+    Reports status of a snapshot tree incrementally to the user: progress,
+    diagnostics, .ilean reference information.
+
+    Debouncing: we only report information
+    * when first blocking, i.e. not before skipping over any unchanged snapshots
+    * 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) (m : DocumentMeta) (snaps : Language.SnapshotTree)
       (cancelTk : CancelToken) : IO Unit := do
-    discard <| go snaps #[] #[] fun _ itrees => do
+    discard <| go snaps { : ReportSnapshotsState } fun st => do
       publishProgressDone m ctx.hOut
+      unless st.hasBlocked do
+        publishDiagnostics m st.diagnostics ctx.hOut
       -- This will overwrite existing ilean info for the file, in case something
       -- went wrong during the incremental updates.
-      publishIleanInfoFinal m ctx.hOut itrees
+      publishIleanInfoFinal m ctx.hOut st.infoTrees
       return .pure <| .ok ()
   where
-    go node diags itrees cont := do
+    go node st cont := do
       if (← cancelTk.ref.get) then
         return .pure <| .ok ()
-      let diags := diags ++ (← node.element.msgLog.toList.toArray.mapM (Widget.msgToInteractiveDiagnostic m.text · ctx.clientHasWidgets)).map (·.toDiagnostic)
-      publishDiagnostics m diags ctx.hOut
-      let itrees := match node.element.infoTree? with
-        | some itree => itrees.push itree
-        | none       => itrees
-      publishIleanInfoUpdate m ctx.hOut itrees
-      goSeq cont diags itrees node.children.toList
-    goSeq cont diags itrees
-      | [] => cont diags itrees
+      let diagnostics := st.diagnostics ++ (← node.element.msgLog.toList.toArray.mapM (Widget.msgToInteractiveDiagnostic m.text · ctx.clientHasWidgets)).map (·.toDiagnostic)
+      if st.hasBlocked && !node.element.msgLog.isEmpty then
+        publishDiagnostics m diagnostics ctx.hOut
+      let infoTrees := match node.element.infoTree? with
+        | some itree => st.infoTrees.push itree
+        | none       => st.infoTrees
+      if st.hasBlocked && node.element.infoTree?.isSome then
+        publishIleanInfoUpdate m ctx.hOut infoTrees
+      goSeq { st with diagnostics, infoTrees } cont node.children.toList
+    goSeq st cont
+      | [] => cont st
       | t::ts => do
-        publishProgressAtPos m t.range.start ctx.hOut
+        let mut st := st
+        unless (← IO.hasFinished t.task) do
+          publishProgressAtPos m t.range.start ctx.hOut
+          if !st.hasBlocked then
+            publishDiagnostics m st.diagnostics ctx.hOut
+            st := { st with hasBlocked := true }
         IO.bindTask t.task fun node =>
-          go node diags itrees fun diags itrees =>
-            goSeq cont diags itrees ts
+          go node st (goSeq · cont ts)
 end Elab
 
 -- Pending requests are tracked so they can be canceled
diff --git a/tests/lean/interactive/unterminatedDocComment.lean.expected.out b/tests/lean/interactive/unterminatedDocComment.lean.expected.out
index 69cbfa37d294..7d543ec0273c 100644
--- a/tests/lean/interactive/unterminatedDocComment.lean.expected.out
+++ b/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": []}

From 7932c0dd4dd71857a8a5da08847ae5f18edb5a06 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 11 Nov 2023 15:36:22 +0100
Subject: [PATCH 09/60] fix: create snapshot for header in legacy adapter

---
 src/Lean/Server/FileWorker.lean       | 5 -----
 src/Lean/Server/FileWorker/Utils.lean | 7 ++++++-
 src/Lean/Server/Snapshots.lean        | 3 ---
 3 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index c9d11543e3be..27ea2de26c90 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -66,11 +66,6 @@ structure WorkerContext where
 /-! # 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)
       (trees : Array Elab.InfoTree) : IO Unit := do
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index 8f9536c98d30..1ba9322f320c 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -57,7 +57,12 @@ partial def cmdSnaps (doc : EditableDocument) : AsyncList ElabTaskError Snapshot
   let some headerParsed := doc.initSnap.success? | return .nil
   .delayed <| headerParsed.next.task.bind fun headerProcessed => Id.run do
     let some headerSuccess := headerProcessed.success? | return .pure <| .ok .nil
-    headerSuccess.next.task.bind go
+    return .pure <| .ok <| .cons {
+      stx := doc.initSnap.stx
+      mpState := headerParsed.parserState
+      cmdState := headerSuccess.cmdState
+      interactiveDiags := .empty  -- TODO
+    } <| .delayed <| headerSuccess.next.task.bind go
 where go cmdParsed :=
   cmdParsed.data.sig.task.bind fun sig =>
     sig.finished.task.map fun finished =>
diff --git a/src/Lean/Server/Snapshots.lean b/src/Lean/Server/Snapshots.lean
index f398bb2849da..1ed26c40c409 100644
--- a/src/Lean/Server/Snapshots.lean
+++ b/src/Lean/Server/Snapshots.lean
@@ -45,9 +45,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

From 00e976cfdfd6b575d131c20d3b7b606670d12894 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 11 Nov 2023 15:59:57 +0100
Subject: [PATCH 10/60] fix: memoize `cmdSnaps` legacy adapter so
 `getFinishedPrefix` works as expected

---
 src/Lean/Server/FileWorker/Utils.lean | 28 +++++++++++++--------------
 1 file changed, 13 insertions(+), 15 deletions(-)

diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index 1ba9322f320c..734bbfcabde9 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -42,23 +42,13 @@ def set (tk : CancelToken) : IO Unit :=
 
 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. -/
-  -- TODO: generalize to other languages by moving request handlers into `Language`
-  initSnap : Language.Lean.InitialSnapshot
-
-namespace EditableDocument
-
-partial def cmdSnaps (doc : EditableDocument) : AsyncList ElabTaskError Snapshot := Id.run do
-  let some headerParsed := doc.initSnap.success? | return .nil
+private partial def mkCmdSnaps (initSnap : Language.Lean.InitialSnapshot) :
+    AsyncList ElabTaskError Snapshot := Id.run do
+  let some headerParsed := initSnap.success? | return .nil
   .delayed <| headerParsed.next.task.bind fun headerProcessed => Id.run do
     let some headerSuccess := headerProcessed.success? | return .pure <| .ok .nil
     return .pure <| .ok <| .cons {
-      stx := doc.initSnap.stx
+      stx := initSnap.stx
       mpState := headerParsed.parserState
       cmdState := headerSuccess.cmdState
       interactiveDiags := .empty  -- TODO
@@ -75,7 +65,15 @@ where go cmdParsed :=
         | some next => .delayed <| next.task.bind go
         | none => .nil)
 
-end EditableDocument
+/-- 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. -/
+  -- TODO: generalize to other languages by moving request handlers into `Language`
+  initSnap : Language.Lean.InitialSnapshot
+  cmdSnaps : AsyncList ElabTaskError Snapshot := mkCmdSnaps initSnap
 
 structure RpcSession where
   objects         : RpcObjectStore

From 7c0ad9c3552e18cbd38aaeeec163c178de667e8a Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 11 Nov 2023 18:09:39 +0100
Subject: [PATCH 11/60] fix: `waitForDiagnostics` must wait on reporter task

---
 src/Lean/Server/FileWorker.lean               | 16 +++++-----
 .../Server/FileWorker/RequestHandling.lean    |  3 +-
 src/Lean/Server/FileWorker/Utils.lean         |  1 +
 .../editAfterError.lean.expected.out          | 30 ++-----------------
 4 files changed, 13 insertions(+), 37 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 27ea2de26c90..31f4d156096c 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -94,8 +94,8 @@ section Elab
     * 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) (m : DocumentMeta) (snaps : Language.SnapshotTree)
-      (cancelTk : CancelToken) : IO Unit := do
-    discard <| go snaps { : ReportSnapshotsState } fun st => do
+      (cancelTk : CancelToken) : IO (Task Unit) := do
+    Task.map (fun _ => ()) <$> go snaps { : ReportSnapshotsState } fun st => do
       publishProgressDone m ctx.hOut
       unless st.hasBlocked do
         publishDiagnostics m st.diagnostics ctx.hOut
@@ -202,9 +202,9 @@ section Initialization
         processor
         clientHasWidgets
       }
-    let doc : EditableDocument := { meta, initSnap }
     let reportCancelTk ← CancelToken.new
-    reportSnapshots ctx doc.meta (Language.ToSnapshotTree.toSnapshotTree doc.initSnap) reportCancelTk
+    let reporter ← reportSnapshots ctx meta (Language.ToSnapshotTree.toSnapshotTree initSnap) reportCancelTk
+    let doc : EditableDocument := { meta, initSnap, reporter }
     return (ctx,
     { doc, reportCancelTk
       initHeaderStx   := initSnap.stx
@@ -218,13 +218,13 @@ section Updates
     modify fun st => { st with pendingRequests := map st.pendingRequests }
 
   /-- Given the new document, updates editable doc state. -/
-  def updateDocument (newMeta : DocumentMeta) : WorkerM Unit := do
+  def updateDocument (meta : DocumentMeta) : WorkerM Unit := do
     let ctx ← read
     (← get).reportCancelTk.set
-    let initSnap ← ctx.processor newMeta.mkInputContext
+    let initSnap ← ctx.processor meta.mkInputContext
     let reportCancelTk ← CancelToken.new
-    reportSnapshots ctx newMeta (Language.ToSnapshotTree.toSnapshotTree initSnap) reportCancelTk
-    modify fun st => { st with doc := { meta := newMeta, initSnap }, reportCancelTk }
+    let reporter ← reportSnapshots ctx meta (Language.ToSnapshotTree.toSnapshotTree initSnap) reportCancelTk
+    modify fun st => { st with doc := { meta, initSnap, reporter }, reportCancelTk }
 end Updates
 
 /- Notifications are handled in the main thread. They may change global worker state
diff --git a/src/Lean/Server/FileWorker/RequestHandling.lean b/src/Lean/Server/FileWorker/RequestHandling.lean
index d04203f62ba5..16b7688b766f 100644
--- a/src/Lean/Server/FileWorker/RequestHandling.lean
+++ b/src/Lean/Server/FileWorker/RequestHandling.lean
@@ -604,8 +604,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 "textDocument/waitForDiagnostics"   WaitForDiagnosticsParams   WaitForDiagnostics      handleWaitForDiagnostics
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index 734bbfcabde9..a93415a66cb7 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -74,6 +74,7 @@ structure EditableDocument where
   -- TODO: generalize to other languages by moving request handlers into `Language`
   initSnap : Language.Lean.InitialSnapshot
   cmdSnaps : AsyncList ElabTaskError Snapshot := mkCmdSnaps initSnap
+  reporter : Task Unit
 
 structure RpcSession where
   objects         : RpcObjectStore
diff --git a/tests/lean/interactive/editAfterError.lean.expected.out b/tests/lean/interactive/editAfterError.lean.expected.out
index 3ad143e86cd2..41193edc5e5f 100644
--- a/tests/lean/interactive/editAfterError.lean.expected.out
+++ b/tests/lean/interactive/editAfterError.lean.expected.out
@@ -4,6 +4,7 @@
    "range":
    {"start": {"line": 1, "character": 10},
     "end": {"line": 1, "character": 11}}}]}
+{"version": 2, "uri": "file://editAfterError.lean", "diagnostics": []}
 {"version": 2,
  "uri": "file://editAfterError.lean",
  "diagnostics":
@@ -13,33 +14,8 @@
    {"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}}}]}
+   {"start": {"line": 0, "character": 7},
+    "end": {"line": 0, "character": 10}}}]}
 {"version": 2,
  "uri": "file://editAfterError.lean",
  "diagnostics":

From 89383cf3c6e33d402f1db1047e5073494d736c88 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Mon, 13 Nov 2023 10:11:32 +0100
Subject: [PATCH 12/60] doc: Language

---
 src/Lean/Language/Basic.lean | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index de82affa9693..474096574416 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -140,9 +140,14 @@ structure ProcessingContext where
   mainModuleName : Name
   opts : Options
 
+/-- Definition of a language processor that can be driven by the cmdline or language server. -/
 structure Language where
+  /--
+    Type of snapshot returned by `process`. It can be converted to a graph of
+    homogeneous snapshot types via `ToSnapshotTree`.  -/
   InitialSnapshot : Type
   [instToSnapshotTree : ToSnapshotTree InitialSnapshot]
+  /-- Processes input into snapshots, potentially reusing information from a previous run. -/
   process : ProcessingContext → (old? : Option InitialSnapshot) → Parser.InputContext → BaseIO InitialSnapshot
   -- TODO: is this the right interface for other languages as well?
   /-- Gets final environment, if any, that is to be used for persisting, code generation, etc. -/

From 58c19d75e13aa963491bc43a12859b9196337cc6 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 25 Nov 2023 12:31:24 +0100
Subject: [PATCH 13/60] chore: rely on auto cancellation of reporter

---
 src/Lean/Server/FileWorker.lean       | 16 ++++++----------
 src/Lean/Server/FileWorker/Utils.lean |  3 +++
 2 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 31f4d156096c..f28e1640f71c 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -94,7 +94,7 @@ section Elab
     * 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) (m : DocumentMeta) (snaps : Language.SnapshotTree)
-      (cancelTk : CancelToken) : IO (Task Unit) := do
+      : IO (Task Unit) := do
     Task.map (fun _ => ()) <$> go snaps { : ReportSnapshotsState } fun st => do
       publishProgressDone m ctx.hOut
       unless st.hasBlocked do
@@ -105,7 +105,7 @@ section Elab
       return .pure <| .ok ()
   where
     go node st cont := do
-      if (← cancelTk.ref.get) then
+      if (← IO.checkCanceled) then
         return .pure <| .ok ()
       let diagnostics := st.diagnostics ++ (← node.element.msgLog.toList.toArray.mapM (Widget.msgToInteractiveDiagnostic m.text · ctx.clientHasWidgets)).map (·.toDiagnostic)
       if st.hasBlocked && !node.element.msgLog.isEmpty then
@@ -140,7 +140,6 @@ structure WorkerState where
   /-- A map of RPC session IDs. We allow asynchronous elab tasks and request handlers
   to modify sessions. A single `Ref` ensures atomic transactions. -/
   rpcSessions     : RBMap UInt64 (IO.Ref RpcSession) compare
-  reportCancelTk  : CancelToken
 
 abbrev WorkerM := ReaderT WorkerContext <| StateRefT WorkerState IO
 
@@ -202,11 +201,10 @@ section Initialization
         processor
         clientHasWidgets
       }
-    let reportCancelTk ← CancelToken.new
-    let reporter ← reportSnapshots ctx meta (Language.ToSnapshotTree.toSnapshotTree initSnap) reportCancelTk
+    let reporter ← reportSnapshots ctx meta (Language.ToSnapshotTree.toSnapshotTree initSnap)
     let doc : EditableDocument := { meta, initSnap, reporter }
     return (ctx,
-    { doc, reportCancelTk
+    { doc
       initHeaderStx   := initSnap.stx
       pendingRequests := RBMap.empty
       rpcSessions     := RBMap.empty
@@ -220,11 +218,9 @@ section Updates
   /-- Given the new document, updates editable doc state. -/
   def updateDocument (meta : DocumentMeta) : WorkerM Unit := do
     let ctx ← read
-    (← get).reportCancelTk.set
     let initSnap ← ctx.processor meta.mkInputContext
-    let reportCancelTk ← CancelToken.new
-    let reporter ← reportSnapshots ctx meta (Language.ToSnapshotTree.toSnapshotTree initSnap) reportCancelTk
-    modify fun st => { st with doc := { meta, initSnap, reporter }, reportCancelTk }
+    let reporter ← reportSnapshots ctx meta (Language.ToSnapshotTree.toSnapshotTree initSnap)
+    modify fun st => { st with doc := { meta, initSnap, reporter } }
 end Updates
 
 /- Notifications are handled in the main thread. They may change global worker state
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index a93415a66cb7..2b4686d596e8 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -74,6 +74,9 @@ structure EditableDocument where
   -- TODO: generalize to other languages by moving request handlers into `Language`
   initSnap : Language.Lean.InitialSnapshot
   cmdSnaps : AsyncList ElabTaskError Snapshot := mkCmdSnaps initSnap
+  /--
+    Task reporting processing status back to client. We store it here for implementing
+    `waitForDiagnostics`. -/
   reporter : Task Unit
 
 structure RpcSession where

From 22aa0dc958726ac04e68f21d46d02704ab3ff069 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Mon, 27 Nov 2023 10:38:59 +0100
Subject: [PATCH 14/60] doc: more

---
 src/Lean/Language/Basic.lean | 116 +++++++++++++++++++++--------------
 src/Lean/Language/Lean.lean  |  19 ++++--
 src/Lean/Server/README.md    |  23 +++++--
 3 files changed, 100 insertions(+), 58 deletions(-)

diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index 474096574416..8809ea4c32e3 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -2,21 +2,26 @@
 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
 -/
 
 import Lean.Message
 import Lean.Parser.Types
 
+set_option linter.missingDocs true
+
 namespace Lean.Language
 
 /--
-  The base class of all snapshots: all the generic information the language server
-  needs about a snapshot. -/
+  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. -/
+    The messages produced by this step. The union of message logs of all finished snapshots is
+    reported to the user. -/
   msgLog : MessageLog
   /-- General elaboration metadata produced by this step. -/
   infoTree? : Option Elab.InfoTree := none
@@ -24,71 +29,87 @@ structure Snapshot where
 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.
+-- 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 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. -/
 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 : α → β) : SnapshotTask β :=
   { t with task := t.task.map f }
 
 /--
-  Chain 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) :
-    BaseIO (SnapshotTask β) :=
+  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) : BaseIO (SnapshotTask β) :=
   return { range, task := (← BaseIO.bindTask t.task fun a => (·.task) <$> (act a)) }
 
+/-- Synchronously waits on the result of the task. -/
 def SnapshotTask.get (t : SnapshotTask α) : α :=
   t.task.get
 
 /--
-  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. -/
+  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. -/
+  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. -/
+  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.msgLog.forM (·.toString (includeEndPos := printMessageEndPos.get opts) >>= IO.print)
   for t in s.children do
@@ -103,19 +124,11 @@ where
     for t in s.children do
       go t.get
 
-/-- Reports `IO` exceptions as single snapshot message. -/
-def withFatalExceptions (ex : Snapshot → α) (act : IO α) : BaseIO α := do
-  match (← act.toBaseIO) with
-  | .error e => return ex {
-    msgLog := MessageLog.empty.add { fileName := "TODO", pos := ⟨0, 0⟩, data := e.toString }
-  }
-  | .ok a => return a
-
-def get? (t? : Option (SnapshotTask α)) : BaseIO (Option α) := do
-  let some t := t? | return none
-  return if (← IO.hasFinished t.task) then t.get else none
-
-def getOrCancel (t? : Option (SnapshotTask α)) : BaseIO (Option α) := do
+/--
+  Returns the value of a snapshot task if it has already finished and otherwise cancels the task.
+  Tasks will be canceled implicitly when the last reference to them is dropped but being explicit
+  here doesn't hurt. -/
+def getOrCancel? (t? : Option (SnapshotTask α)) : BaseIO (Option α) := do
   let some t := t? | return none
   if (← IO.hasFinished t.task) then
     return t.get
@@ -123,12 +136,9 @@ def getOrCancel (t? : Option (SnapshotTask α)) : BaseIO (Option α) := do
   return none
 
 /--
-  Checks whether reparsing `stx` in `newSource` should result in `stx` again.
-  `stopPos` is the furthest position the parser has looked at for producing
-  `stx`, which usually is the extent of the subsequent token (if any) plus one. -/
--- This is a fundamentally different approach from the "go up one snapshot from
--- the point of change" in the old implementation which was failing in edge
--- cases and is not generalizable to a DAG of snapshots.
+  Checks whether reparsing `stx` in `newSource` should result in `stx` again. `stopPos` is the
+  furthest position the parser could have looked at for producing `stx`, which usually is the extent
+  of the subsequent token (if any) plus one. -/
 def unchangedParse (stx : Syntax) (stopPos : String.Pos) (newSource : String) : Bool :=
   if let .original start .. := stx.getHeadInfo then
     let oldSubstr := { start with stopPos }
@@ -137,27 +147,39 @@ def unchangedParse (stx : Syntax) (stopPos : String.Pos) (newSource : String) :
 
 /-- Metadata that does not change during the lifetime of the language processing process. -/
 structure ProcessingContext 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
 
+end Language
+open Language
+
 /-- Definition of a language processor that can be driven by the cmdline or language server. -/
 structure Language where
   /--
-    Type of snapshot returned by `process`. It can be converted to a graph of
-    homogeneous snapshot types via `ToSnapshotTree`.  -/
+    Type of snapshot returned by `process`. It can be converted to a graph of homogeneous snapshot
+    types via `ToSnapshotTree`.  -/
   InitialSnapshot : Type
+  /-- Instance for transforming the initial snapshot into a snapshot tree for reporting. -/
   [instToSnapshotTree : ToSnapshotTree InitialSnapshot]
-  /-- Processes input into snapshots, potentially reusing information from a previous run. -/
-  process : ProcessingContext → (old? : Option InitialSnapshot) → Parser.InputContext → BaseIO InitialSnapshot
+  /--
+    Processes input into snapshots, potentially reusing information from a previous run.
+    Constructing the initial snapshot is assumed to be cheap enough that it can be done
+    synchronously, which simplifies use of this function. -/
+  process : ProcessingContext → (old? : Option InitialSnapshot) → Parser.InputContext →
+    BaseIO InitialSnapshot
   -- TODO: is this the right interface for other languages as well?
   /-- Gets final environment, if any, that is to be used for persisting, code generation, etc. -/
   getFinalEnv? : InitialSnapshot → Option Environment
 
 instance (lang : Language) : ToSnapshotTree lang.InitialSnapshot := lang.instToSnapshotTree
 
-/-- Returns a function for processing a language using incremental snapshots. -/
-partial def Language.mkIncrementalProcessor (lang : Language) :
-    ProcessingContext → BaseIO (Parser.InputContext → BaseIO lang.InitialSnapshot) := fun ctx => do
+/--
+  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 (lang : Language) (ctx: ProcessingContext) :
+    BaseIO (Parser.InputContext → BaseIO lang.InitialSnapshot) := do
   let oldRef ← IO.mkRef none
   return fun doc => do
     let snap ← lang.process ctx (← oldRef.get) doc
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index decb86ddb05e..1a71bb56d785 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -109,6 +109,15 @@ instance : ToSnapshotTree HeaderParsedSnapshot where
 
 abbrev InitialSnapshot := HeaderParsedSnapshot
 
+/--
+  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 : IO α) : BaseIO α := do
+  match (← act.toBaseIO) with
+  | .error e => return ex {
+    msgLog := MessageLog.empty.add { fileName := "<input>", pos := ⟨0, 0⟩, data := e.toString } }
+  | .ok a => return a
+
 /-- Entry point of the Lean language processor. -/
 -- As a general note, for each processing function we pass in the previous
 -- state, if any, in order to reuse still-valid state. Thus the logic of
@@ -134,7 +143,7 @@ where
         -- ...go immediately to next snapshot
         return (← unchanged old success)
 
-    withFatalExceptions ({ · with stx := .missing, success? := none }) do
+    withHeaderExceptions ({ · with stx := .missing, success? := none }) do
       let (stx, parserState, msgLog) ← Parser.parseHeader ictx
       -- TODO: move into `parseHeader`; it should add the length of `peekToken`,
       -- which is the full extent the parser considered before stopping
@@ -164,9 +173,9 @@ where
         return .pure old
 
     SnapshotTask.ofIO ⟨0, ictx.input.endPos⟩ do
-    withFatalExceptions ({ · with success? := none }) do
+    withHeaderExceptions ({ · with success? := none }) do
     -- discard existing continuation if any, there is nothing to reuse
-    let _ ← old?.bind (·.success?) |>.map (·.next) |> getOrCancel
+    let _ ← old?.bind (·.success?) |>.map (·.next) |> getOrCancel?
     -- function copied from current implementation, see below
     -- TODO: we should do this at most once per process
     let cmdState ← doImport stx
@@ -183,7 +192,7 @@ where
         return .mk (data := old.data)
           (next? := (← old.data.sig.bindIO fun sig =>
             sig.finished.bindIO fun finished => do
-              parseCmd (← getOrCancel oldNext) old.data.parserState finished.cmdState))
+              parseCmd (← getOrCancel? oldNext) old.data.parserState finished.cmdState))
       else return old  -- terminal command, we're done!
 
     -- fast path, do not even start new task for this snapshot
@@ -206,7 +215,7 @@ where
 
       -- signature elaboration task; for now, does full elaboration
       -- TODO: do tactic snapshots, reuse old state for them
-      let _ ← old?.map (·.data.sig) |> getOrCancel
+      let _ ← old?.map (·.data.sig) |> getOrCancel?
       let sig ← SnapshotTask.ofIO (stx.getRange?.getD ⟨parserState.pos, parserState.pos⟩) do
         let cmdState ← doElab stx cmdState msgLog.hasErrors beginPos scope
         return {
diff --git a/src/Lean/Server/README.md b/src/Lean/Server/README.md
index b6d1a0d01363..e3d90aad2ee6 100644
--- a/src/Lean/Server/README.md
+++ b/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,12 +47,19 @@ 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 generic `Lean.Language` interface.
+It currently has a single implementation `Lean.Language.Lean` for the main Lean 4 language but other languages will be loadable via a `#lang` file annotation in the future.
+The interface is shared with the cmdline driver that is used for building Lean files but the core concept of *snapshots* produced by a `Language` is mostly of interest to the worker: snapshots are how processing data is saved between versions of a file such as to make processing on file edits incremental.
+How exactly incrementality is implemented is left completely to the `Language`; 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 `Language.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 `Language` given the new file content and asynchronously wait on all these nodes and report the processing status stored therein to the client (`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.

From 0ce8fee589c3135a428d28a1bf2bd9b913b3c64a Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Wed, 29 Nov 2023 16:26:37 +0100
Subject: [PATCH 15/60] import at most once per worker process

---
 src/Lean/Language/Lean.lean | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 61df5535de31..f126c751478b 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -121,6 +121,8 @@ private def withHeaderExceptions (ex : Snapshot → α) (act : IO α) : BaseIO 
   | .error e => return ex { msgLog := msglogOfHeaderError e.toString }
   | .ok a => return a
 
+builtin_initialize importsLoadedRef : IO.Ref Bool ← IO.mkRef false
+
 /-- Entry point of the Lean language processor. -/
 -- As a general note, for each processing function we pass in the previous
 -- state, if any, in order to reuse still-valid state. Thus the logic of
@@ -179,7 +181,6 @@ where
     withHeaderExceptions ({ · with success? := none }) do
     -- discard existing continuation if any, there is nothing to reuse
     let _ ← old?.bind (·.success?) |>.map (·.next) |> getOrCancel?
-    -- TODO: we should do this at most once per process
     let fileSetupResult ← if let some handler := ctx.fileSetupHandler? then
       handler (Elab.headerToImports stx)
     else
@@ -198,6 +199,15 @@ where
       }
     | _ => pure ()
 
+    if (← importsLoadedRef.modifyGet ((·, true))) 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 { msgLog := .empty, success? := none }
+
     let opts := ctx.opts.mergeBy (fun _ _ fileOpt => fileOpt) fileSetupResult.fileOptions
     let mut (headerEnv, msgLog) ← Elab.processHeader (leakEnv := true) stx opts .empty ictx
     headerEnv := headerEnv.setMainModule ctx.mainModuleName

From 4c9b70bd03591ce1058b24a758952a8afd361165 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Wed, 29 Nov 2023 16:47:01 +0100
Subject: [PATCH 16/60] fix test outputs

---
 tests/lean/Process.lean.expected.out                    | 4 ++--
 tests/lean/dbgMacros.lean.expected.out                  | 4 ++--
 tests/lean/interactive/editAfterError.lean.expected.out | 2 +-
 3 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tests/lean/Process.lean.expected.out b/tests/lean/Process.lean.expected.out
index fbeb40bc4ba3..dbe051b1ade5 100644
--- a/tests/lean/Process.lean.expected.out
+++ b/tests/lean/Process.lean.expected.out
@@ -1,6 +1,5 @@
-hi!
-<stdin>:1:0: warning: using 'exit' to interrupt Lean
 1
+hi!
 0
 "ho!\n"
 "hu!\n"
@@ -9,5 +8,6 @@ flush of broken pipe failed
 100000
 0
 0
+<stdin>:1:0: warning: using 'exit' to interrupt Lean
 0
 0
diff --git a/tests/lean/dbgMacros.lean.expected.out b/tests/lean/dbgMacros.lean.expected.out
index 9b70d1b86045..54d6f8756c15 100644
--- a/tests/lean/dbgMacros.lean.expected.out
+++ b/tests/lean/dbgMacros.lean.expected.out
@@ -1,10 +1,10 @@
 PANIC at f dbgMacros:2:14: unexpected zero
-PANIC at g dbgMacros:10:14: unreachable code has been reached
-PANIC at h dbgMacros:16:0: assertion violation: x != 0
 0
 9
+PANIC at g dbgMacros:10:14: unreachable code has been reached
 0
 0
+PANIC at h dbgMacros:16:0: assertion violation: x != 0
 0
 f2, x: 10
 11
diff --git a/tests/lean/interactive/editAfterError.lean.expected.out b/tests/lean/interactive/editAfterError.lean.expected.out
index 68aef365e18e..76f7246bd0ea 100644
--- a/tests/lean/interactive/editAfterError.lean.expected.out
+++ b/tests/lean/interactive/editAfterError.lean.expected.out
@@ -4,7 +4,7 @@
    "range":
    {"start": {"line": 1, "character": 10},
     "end": {"line": 1, "character": 11}}}]}
-{"version": 2, "uri": "file://editAfterError.lean", "diagnostics": []}
+{"version": 2, "uri": "file:///editAfterError.lean", "diagnostics": []}
 {"version": 2,
  "uri": "file:///editAfterError.lean",
  "diagnostics":

From 0282d1c249f6664682cdbe902c249459e5f4bfdd Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 30 Nov 2023 15:18:26 +0100
Subject: [PATCH 17/60] File worker: synchronize output messages

---
 src/Init/Data/Channel.lean      |   2 +-
 src/Lean/Data/JsonRpc.lean      |   3 +
 src/Lean/Server/FileWorker.lean | 119 ++++++++++++++++++++------------
 src/Lean/Server/Requests.lean   |   1 -
 src/Lean/Server/Utils.lean      |  45 ++++++------
 src/Lean/Server/Watchdog.lean   |   6 +-
 6 files changed, 102 insertions(+), 74 deletions(-)

diff --git a/src/Init/Data/Channel.lean b/src/Init/Data/Channel.lean
index e109c3e40fc7..087e36052bdf 100644
--- a/src/Init/Data/Channel.lean
+++ b/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
diff --git a/src/Lean/Data/JsonRpc.lean b/src/Lean/Data/JsonRpc.lean
index d0bba7fee330..ea9eb9e631f2 100644
--- a/src/Lean/Data/JsonRpc.lean
+++ b/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⟩
 
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 68231bd36c17..69e0cc2c9427 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -57,8 +57,9 @@ 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
   hLog             : FS.Stream
   initParams       : InitializeParams
   processor        : Parser.InputContext → BaseIO Lean.Language.Lean.InitialSnapshot
@@ -68,17 +69,19 @@ structure WorkerContext where
 
 section Elab
   -- Placed here instead of Lean.Server.Utils because of an import loop
-  private def publishIleanInfo (method : String) (m : DocumentMeta) (hOut : FS.Stream)
-      (trees : Array Elab.InfoTree) : IO Unit := do
+  private def mkIleanInfoNotification (method : String) (m : DocumentMeta)
+      (trees : Array Elab.InfoTree) : JsonRpc.Notification Lsp.LeanIleanInfoParams :=
     let references := findModuleRefs m.text trees (localVars := true)
-    let param := { version := m.version, references : LeanIleanInfoParams }
-    hOut.writeLspNotification { method, param }
+    let param := { version := m.version, references }
+    { method, param }
 
-  private def publishIleanInfoUpdate : DocumentMeta → FS.Stream → Array Elab.InfoTree → IO Unit :=
-    publishIleanInfo "$/lean/ileanInfoUpdate"
+  private def mkIleanInfoUpdateNotification : DocumentMeta → Array Elab.InfoTree →
+      JsonRpc.Notification Lsp.LeanIleanInfoParams :=
+    mkIleanInfoNotification "$/lean/ileanInfoUpdate"
 
-  private def publishIleanInfoFinal : DocumentMeta → FS.Stream → Array Elab.InfoTree → IO Unit :=
-    publishIleanInfo "$/lean/ileanInfoFinal"
+  private def mkIleanInfoFinalNotification : DocumentMeta → Array Elab.InfoTree →
+      JsonRpc.Notification Lsp.LeanIleanInfoParams :=
+    mkIleanInfoNotification "$/lean/ileanInfoFinal"
 
   private structure ReportSnapshotsState where
     hasBlocked := false
@@ -97,12 +100,12 @@ section Elab
   private partial def reportSnapshots (ctx : WorkerContext) (m : DocumentMeta) (snaps : Language.SnapshotTree)
       : IO (Task Unit) := do
     Task.map (fun _ => ()) <$> go snaps { : ReportSnapshotsState } fun st => do
-      publishProgressDone m ctx.hOut
+      ctx.chanOut.send <| mkFileProgressDoneNotification m
       unless st.hasBlocked do
-        publishDiagnostics m st.diagnostics ctx.hOut
+        ctx.chanOut.send <| mkPublishDiagnosticsNotification m st.diagnostics
       -- This will overwrite existing ilean info for the file, in case something
       -- went wrong during the incremental updates.
-      publishIleanInfoFinal m ctx.hOut st.infoTrees
+      ctx.chanOut.send <| mkIleanInfoFinalNotification m st.infoTrees
       return .pure <| .ok ()
   where
     go node st cont := do
@@ -110,21 +113,21 @@ section Elab
         return .pure <| .ok ()
       let diagnostics := st.diagnostics ++ (← node.element.msgLog.toList.toArray.mapM (Widget.msgToInteractiveDiagnostic m.text · ctx.clientHasWidgets)).map (·.toDiagnostic)
       if st.hasBlocked && !node.element.msgLog.isEmpty then
-        publishDiagnostics m diagnostics ctx.hOut
+        ctx.chanOut.send <| mkPublishDiagnosticsNotification m diagnostics
       let infoTrees := match node.element.infoTree? with
         | some itree => st.infoTrees.push itree
         | none       => st.infoTrees
       if st.hasBlocked && node.element.infoTree?.isSome then
-        publishIleanInfoUpdate m ctx.hOut infoTrees
+        ctx.chanOut.send <| mkIleanInfoUpdateNotification m infoTrees
       goSeq { st with diagnostics, infoTrees } cont node.children.toList
     goSeq st cont
       | [] => cont st
       | t::ts => do
         let mut st := st
         unless (← IO.hasFinished t.task) do
-          publishProgressAtPos m t.range.start ctx.hOut
+          ctx.chanOut.send <| mkFileProgressAtPosNotification m t.range.start
           if !st.hasBlocked then
-            publishDiagnostics m st.diagnostics ctx.hOut
+            ctx.chanOut.send <| mkPublishDiagnosticsNotification m st.diagnostics
             st := { st with hasBlocked := true }
         IO.bindTask t.task fun node =>
           go node st (goSeq · cont ts)
@@ -154,7 +157,7 @@ abbrev WorkerM := ReaderT WorkerContext <| StateRefT WorkerState IO
 
 /- Worker initialization sequence. -/
 section Initialization
-  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 mut mainModuleName := Name.anonymous
@@ -162,6 +165,7 @@ section Initialization
       if let some path := System.Uri.fileUriToPath? meta.uri then
         mainModuleName ← moduleNameOfFileName path none
     catch _ => pure ()
+    let chanOut ← mkLspOutputChannel
     let processor ← Language.Lean.mkIncrementalProcessor {
       opts, mainModuleName
       fileSetupHandler? := some fun imports =>
@@ -171,18 +175,16 @@ section Initialization
             severity? := DiagnosticSeverity.information
             message   := stderrLine
           }
-          -- TODO: should use a channel to the reporting thread instead of using `o` directly
-          publishDiagnostics meta #[progressDiagnostic] o
+          chanOut.send <| mkPublishDiagnosticsNotification meta #[progressDiagnostic]
     }
     let initSnap ← processor meta.mkInputContext
-    let ctx :=
-      { hIn  := i
-        hOut := o
-        hLog := e
-        initParams
-        processor
-        clientHasWidgets
-      }
+    let ctx := {
+      chanOut
+      hLog := e
+      initParams
+      processor
+      clientHasWidgets
+    }
     let reporter ← reportSnapshots ctx meta (Language.ToSnapshotTree.toSnapshotTree initSnap)
     let doc : EditableDocument := { meta, initSnap, reporter }
     return (ctx,
@@ -193,6 +195,32 @@ section Initialization
       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 : IO (IO.Channel JsonRpc.Message) := do
+      let chanOut ← IO.Channel.new
+      -- most recent document version seen in notifications
+      let maxVersion ← IO.mkRef 0
+      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 Nat := do
+          let doc ← match msg with
+            | .notification "textDocument/publishDiagnostics" (some <| .obj params) => some params
+            | .notification "$/lean/fileProgress" (some <| .obj params) =>
+              params.find compare "textDocument" |>.bind (·.getObj?.toOption)
+            | _ => none
+          doc.find compare "version" |>.bind (·.getNat?.toOption)
+        if let some version := version? then
+          if version < (← maxVersion.get) then
+            return
+          else
+            maxVersion.set version
+        o.writeLspMessage msg |>.catchExceptions (fun _ => pure ())
+      return chanOut
 end Initialization
 
 section Updates
@@ -285,7 +313,7 @@ section MessageHandling
       let availableImports ← ImportCompletion.collectAvailableImports
       let lastRequestTimestampMs ← IO.monoMsNow
       let completions := ImportCompletion.find text st.currHeaderStx params availableImports
-      ctx.hOut.writeLspResponse ⟨id, completions⟩
+      ctx.chanOut.send <| .response id (toJson completions)
       pure { availableImports, lastRequestTimestampMs : AvailableImportsCache }
 
     | some task => IO.mapTask (t := task) fun result => do
@@ -295,7 +323,7 @@ section MessageHandling
         availableImports ← ImportCompletion.collectAvailableImports
       lastRequestTimestampMs := timestampNowMs
       let completions := ImportCompletion.find text st.currHeaderStx params availableImports
-      ctx.hOut.writeLspResponse ⟨id, completions⟩
+      ctx.chanOut.send <| .response id (toJson completions)
       pure { availableImports, lastRequestTimestampMs : AvailableImportsCache }
 
   def handleRequest (id : RequestID) (method : String) (params : Json)
@@ -307,12 +335,9 @@ section MessageHandling
       try
         let ps ← parseParams RpcConnectParams params
         let resp ← handleRpcConnect ps
-        ctx.hOut.writeLspResponse ⟨id, resp⟩
+        ctx.chanOut.send <| .response id (toJson resp)
       catch e =>
-        ctx.hOut.writeLspResponseError
-          { id
-            code := ErrorCode.internalError
-            message := toString e }
+        ctx.chanOut.send <| .responseError id .internalError (toString e) none
       return
 
     if method == "textDocument/completion" then
@@ -333,26 +358,25 @@ section MessageHandling
          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
@@ -399,12 +423,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]
diff --git a/src/Lean/Server/Requests.lean b/src/Lean/Server/Requests.lean
index 901fdfb9f219..8ca1229424ac 100644
--- a/src/Lean/Server/Requests.lean
+++ b/src/Lean/Server/Requests.lean
@@ -63,7 +63,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 α)
diff --git a/src/Lean/Server/Utils.lean b/src/Lean/Server/Utils.lean
index f8d7e66a5736..ebf85dbdc7b0 100644
--- a/src/Lean/Server/Utils.lean
+++ b/src/Lean/Server/Utils.lean
@@ -112,36 +112,35 @@ 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?    := m.version
-      diagnostics := diagnostics
-      : PublishDiagnosticsParams
-    }
+def mkPublishDiagnosticsNotification (m : DocumentMeta) (diagnostics : Array Lsp.Diagnostic) :
+    JsonRpc.Notification Lsp.PublishDiagnosticsParams where
+  method := "textDocument/publishDiagnostics"
+  param  := {
+    uri         := m.uri
+    version?    := 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
-    }
+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
+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
+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
 
diff --git a/src/Lean/Server/Watchdog.lean b/src/Lean/Server/Watchdog.lean
index f09f51d3463c..f527cacb3f51 100644
--- a/src/Lean/Server/Watchdog.lean
+++ b/src/Lean/Server/Watchdog.lean
@@ -221,7 +221,7 @@ section ServerM
             ++ " or the file was closed, 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
@@ -229,7 +229,7 @@ section ServerM
           -- Worker crashed
           fw.errorPendingRequests o (if exitCode = 1 then ErrorCode.workerExited else ErrorCode.workerCrashed)
             s!"Server process for {fw.doc.uri} crashed, {if exitCode = 1 then "see stderr for exception" else "likely due to a stack overflow or a bug"}."
-          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
@@ -238,7 +238,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

From 8fff2e2907f640eec524d214ecad42150dcd7559 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 1 Dec 2023 17:08:50 +0100
Subject: [PATCH 18/60] fix: incremental parsing

---
 src/Init/Data/Option/Basic.lean |  6 +++++
 src/Lean/Language/Basic.lean    | 14 +++-------
 src/Lean/Language/Lean.lean     | 47 +++++++++++++++------------------
 3 files changed, 31 insertions(+), 36 deletions(-)

diff --git a/src/Init/Data/Option/Basic.lean b/src/Init/Data/Option/Basic.lean
index 21359e6dfaaf..c07201c8e697 100644
--- a/src/Init/Data/Option/Basic.lean
+++ b/src/Init/Data/Option/Basic.lean
@@ -34,6 +34,12 @@ def toMonad [Monad m] [Alternative m] : Option α → m α
   | none,   _ => none
   | some a, b => b a
 
+@[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)
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index 373da45daaaf..ca8f0217bf69 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -99,6 +99,10 @@ def SnapshotTask.bindIO (t : SnapshotTask α) (act : α → BaseIO (SnapshotTask
 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.
@@ -161,16 +165,6 @@ def getOrCancel? (t? : Option (SnapshotTask α)) : BaseIO (Option α) := do
   IO.cancel t.task
   return none
 
-/--
-  Checks whether reparsing `stx` in `newSource` should result in `stx` again. `stopPos` is the
-  furthest position the parser could have looked at for producing `stx`, which usually is the extent
-  of the subsequent token (if any) plus one. -/
-def unchangedParse (stx : Syntax) (stopPos : String.Pos) (newSource : String) : Bool :=
-  if let .original start .. := stx.getHeadInfo then
-    let oldSubstr := { start with stopPos }
-    oldSubstr == { oldSubstr with str := newSource }
-  else false
-
 /-- Metadata that does not change during the lifetime of the language processing process. -/
 structure ProcessingContext where
   /-- Module name of the file being processed. -/
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index f126c751478b..dca3dc4a63a0 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -67,8 +67,6 @@ instance : ToSnapshotTree CommandSignatureProcessedSnapshot where
 structure CommandParsedSnapshotData extends Snapshot where
   stx : Syntax
   parserState : Parser.ModuleParserState
-  /-- Furthest position the parser has looked at; used for incremental parsing. -/
-  stopPos : String.Pos
   sig : SnapshotTask CommandSignatureProcessedSnapshot
 deriving Nonempty
 inductive CommandParsedSnapshot where
@@ -97,12 +95,11 @@ instance : ToSnapshotTree HeaderProcessedSnapshot where
 
 structure HeaderParsedSucessfully where
   parserState : Parser.ModuleParserState
-  /-- Furthest position the parser has looked at; used for incremental parsing. -/
-  stopPos : String.Pos
   next : SnapshotTask HeaderProcessedSnapshot
 
 /-- State after the module header has been parsed. -/
 structure HeaderParsedSnapshot extends Snapshot where
+  ictx : Parser.InputContext
   stx : Syntax
   success? : Option HeaderParsedSucessfully
 instance : ToSnapshotTree HeaderParsedSnapshot where
@@ -135,6 +132,7 @@ partial def processLean
     BaseIO InitialSnapshot :=
   parseHeader old?
 where
+  firstDiffPos? := old?.map (·.ictx.input.firstDiffPos ictx.input)
   parseHeader (old? : Option HeaderParsedSnapshot) := do
     let unchanged old success :=
       -- when header syntax is unchanged, reuse import processing task as is
@@ -143,29 +141,27 @@ where
 
     -- fast path: if we have parsed the header sucessfully...
     if let some old@{ success? := some success, .. } := old? then
-      -- ...and the header syntax appears unchanged...
-      if unchangedParse old.stx success.stopPos ictx.input then
+      -- ...and the edit location is after the next command...
+      if let some nextCom ← (← success.next.get?).bind (·.success?) |>.bindM (·.next.get?) then
+        if firstDiffPos?.any (nextCom.data.parserState.pos < ·) then
         -- ...go immediately to next snapshot
         return (← unchanged old success)
+        -- (if there is no next command, we always reparse)
 
-    withHeaderExceptions ({ · with stx := .missing, success? := none }) do
+    withHeaderExceptions ({ · with ictx, stx := .missing, success? := none }) do
       let (stx, parserState, msgLog) ← Parser.parseHeader ictx
-      -- TODO: move into `parseHeader`; it should add the length of `peekToken`,
-      -- which is the full extent the parser considered before stopping
-      let stopPos := parserState.pos + (⟨1⟩ : String.Pos)
-
       if msgLog.hasErrors then
-        return { stx, msgLog, success? := none }
+        return { ictx, stx, msgLog, success? := none }
 
       -- semi-fast path: go to next snapshot if syntax tree is unchanged
-      if let some old@{ success? := some success, .. } := old? then
-        if old.stx == stx then
-          return (← unchanged old success)
-        success.next.cancel
+      -- TODO: would need adjusting all contained positions
+      --if let some old@{ success? := some success, .. } := old? then
+      --  if old.stx == stx then
+      --    return (← unchanged old success)
+      --  success.next.cancel
 
-      return { stx, msgLog, success? := some {
+      return { ictx, stx, msgLog, success? := some {
         parserState
-        stopPos
         next := (← processHeader none stx parserState) } }
 
   processHeader (old? : Option HeaderProcessedSnapshot) (stx : Syntax) (parserState : Parser.ModuleParserState) := do
@@ -251,8 +247,9 @@ where
 
     -- fast path, do not even start new task for this snapshot
     if let some old := old? then
-      if unchangedParse old.data.stx old.data.stopPos ictx.input then
-        return .pure (← unchanged old)
+      if let some nextCom ← old.next?.bindM (·.get?) then
+        if firstDiffPos?.any (nextCom.data.parserState.pos < ·) then
+          return .pure (← unchanged old)
 
     SnapshotTask.ofIO ⟨parserState.pos, ictx.input.endPos⟩ do
       let beginPos := parserState.pos
@@ -260,12 +257,10 @@ where
       let pmctx := { env := cmdState.env, options := scope.opts, currNamespace := scope.currNamespace, openDecls := scope.openDecls }
       let (stx, parserState, msgLog) := Parser.parseCommand ictx pmctx parserState .empty
 
-      -- TODO: move into `parseCommand`; it should add the length of `peekToken`,
-      -- which is the full extent the parser considered before stopping
-      let stopPos := parserState.pos + (⟨1⟩ : String.Pos)
-      if let some old := old? then
-        if old.data.stx == stx then
-          return (← unchanged old)
+      -- TODO: would need adjusting all contained positions
+      --if let some old := old? then
+      --  if old.data.stx == stx then
+      --    return (← unchanged old)
 
       -- signature elaboration task; for now, does full elaboration
       -- TODO: do tactic snapshots, reuse old state for them

From ca75ce2a770635e408baab3cbbc8c41cf921772a Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 1 Dec 2023 19:15:38 +0100
Subject: [PATCH 19/60] polish Lean.Language.Lean

---
 src/Lean/Language/Lean.lean               | 393 +++++++++++++---------
 src/Lean/Server/FileWorker/SetupFile.lean |   5 -
 2 files changed, 238 insertions(+), 160 deletions(-)

diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index dca3dc4a63a0..1fb4df2731ae 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -2,6 +2,9 @@
 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
 -/
 
@@ -10,21 +13,69 @@ 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
 
-def pushOpt (a? : Option α) (as : Array α) : Array α :=
+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)"
@@ -34,77 +85,86 @@ register_builtin_option showPartialSyntaxErrors : Bool := {
 
 /-- 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 execution of a single synchronous tactic step. -/
-inductive TacticEvaluatedSnapshot where
-  | mk (toSnapshot : Snapshot) (next : Array (SnapshotTask TacticEvaluatedSnapshot))
-/-- Potential, potentially parallel, follow-up tactic executions. -/
--- In the first, non-parallel version, each task will depend on its predecessor
-abbrev TacticEvaluatedSnapshot.next : TacticEvaluatedSnapshot → Array (SnapshotTask TacticEvaluatedSnapshot)
-  | mk _ next => next
-partial instance : ToSnapshotTree TacticEvaluatedSnapshot where
-  toSnapshotTree := go where
-    go := fun ⟨s, next⟩ => ⟨s, next.map (·.map go)⟩
-
 /--
-  State after processing a command's signature and before executing its tactic
-  body, if any. Other commands should immediately proceed to `finished`. -/
+  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
-  /-- Potential, potentially parallel, follow-up tactic executions. -/
-  tacs : Array (SnapshotTask TacticEvaluatedSnapshot)
   /-- State after processing is finished. -/
-  -- If we make proofs completely opaque, this will not have to depend on `tacs`
   finished : SnapshotTask CommandFinishedSnapshot
 deriving Nonempty
 instance : ToSnapshotTree CommandSignatureProcessedSnapshot where
-  toSnapshotTree s := ⟨s.toSnapshot, s.tacs.map (·.map toSnapshotTree) |>.push (s.finished.map toSnapshotTree)⟩
+  toSnapshotTree s := ⟨s.toSnapshot, #[s.finished.map 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. -/
   sig : 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) (next? : 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)
+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.sig.map toSnapshotTree] |> pushOpt (s.next?.map (·.map go))⟩
+    go s := ⟨s.data.toSnapshot,
+      #[s.data.sig.map toSnapshotTree] |> pushOpt (s.next?.map (·.map go))⟩
 
+/-- State after successful importing. -/
 structure HeaderProcessedSucessfully where
+  /-- The resulting initial elaboration state. -/
   cmdState : Command.State
+  /-- The search path communicated by `lake setup-file`, if any. -/
   srcSearchPath : SearchPath
+  /-- First command task (there is always at least a terminal command). -/
   next : SnapshotTask CommandParsedSnapshot
 
 /-- State after the module header has been processed including imports. -/
 structure HeaderProcessedSnapshot extends Snapshot where
+  /-- State after successful importing. -/
   success? : Option HeaderProcessedSucessfully
 instance : ToSnapshotTree HeaderProcessedSnapshot where
   toSnapshotTree s := ⟨s.toSnapshot, #[] |> pushOpt (s.success?.map (·.next.map toSnapshotTree))⟩
 
+/-- State after successfully parsing the module header. -/
 structure HeaderParsedSucessfully where
+  /-- Resulting parser state. -/
   parserState : Parser.ModuleParserState
+  /-- Header processing task. -/
   next : 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. -/
   success? : Option HeaderParsedSucessfully
 instance : ToSnapshotTree HeaderParsedSnapshot where
   toSnapshotTree s := ⟨s.toSnapshot, #[] |> pushOpt (s.success?.map (·.next.map toSnapshotTree))⟩
 
+/-- Initial snapshot of the Lean language processor: a "header parsed" snapshot. -/
 abbrev InitialSnapshot := HeaderParsedSnapshot
 
 private def msglogOfHeaderError (data : MessageData) : MessageLog :=
@@ -118,123 +178,141 @@ private def withHeaderExceptions (ex : Snapshot → α) (act : IO α) : BaseIO 
   | .error e => return ex { msgLog := msglogOfHeaderError e.toString }
   | .ok a => return a
 
-builtin_initialize importsLoadedRef : IO.Ref Bool ← IO.mkRef false
+/-- 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
 
 /-- Entry point of the Lean language processor. -/
--- As a general note, for each processing function we pass in the previous
--- state, if any, in order to reuse still-valid state. Thus the logic of
--- what snapshots to reuse is completely moved out of the server into the
--- language-specific processor. As there is no cheap way to check whether
--- the `Environment` is unchanged, we must make sure to pass `none` as all
--- follow-up "previous states" when we do change it.
+/-
+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 (try to remember to) cancel tasks that we cannot reuse proactively even though they are
+  auto-canceled when the last reference to them is dropped; but it is not always easy to ensure that
+  the last drop happens precisely where we expect it.
+* 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 processLean
-  (ctx : ProcessingContext) (old? : Option InitialSnapshot) (ictx : Parser.InputContext) :
+    (ctx : ProcessingContext) (old? : Option InitialSnapshot) (ictx : Parser.InputContext) :
     BaseIO InitialSnapshot :=
   parseHeader old?
 where
+  -- compute position of syntactic change once
   firstDiffPos? := old?.map (·.ictx.input.firstDiffPos ictx.input)
+
   parseHeader (old? : Option HeaderParsedSnapshot) := do
     let unchanged old success :=
-      -- when header syntax is unchanged, reuse import processing task as is
+      -- when header syntax is unchanged, reuse import processing task as is and continue with
+      -- parsing
       return { old with success? := some { success with
-        next := (← success.next.bindIO (processHeader · old.stx success.parserState)) } }
+        next := (← success.next.bindIO fun processed => do
+          if let some procSuccess := processed.success? then
+            let oldCmd? ← getOrCancel? procSuccess.next
+            return .pure { processed with success? := some { procSuccess with
+              next := (← parseCmd oldCmd? success.parserState procSuccess.cmdState) } }
+          else
+            return .pure processed) } }
 
     -- fast path: if we have parsed the header sucessfully...
     if let some old@{ success? := some success, .. } := old? then
       -- ...and the edit location is after the next command...
       if let some nextCom ← (← success.next.get?).bind (·.success?) |>.bindM (·.next.get?) then
         if firstDiffPos?.any (nextCom.data.parserState.pos < ·) then
-        -- ...go immediately to next snapshot
-        return (← unchanged old success)
-        -- (if there is no next command, we always reparse)
+          -- ...go immediately to next snapshot
+          return (← unchanged old success)
 
     withHeaderExceptions ({ · with ictx, stx := .missing, success? := none }) do
       let (stx, parserState, msgLog) ← Parser.parseHeader ictx
       if msgLog.hasErrors then
         return { ictx, stx, msgLog, success? := none }
 
-      -- semi-fast path: go to next snapshot if syntax tree is unchanged
-      -- TODO: would need adjusting all contained positions
-      --if let some old@{ success? := some success, .. } := old? then
-      --  if old.stx == stx then
-      --    return (← unchanged old success)
-      --  success.next.cancel
-
-      return { ictx, stx, msgLog, success? := some {
-        parserState
-        next := (← processHeader none stx parserState) } }
+      -- 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
+      if let some old@{ success? := some success, .. } := old? then
+        if firstDiffPos?.any (parserState.pos < ·) && old.stx == stx then
+          return (← unchanged old success)
+        success.next.cancel
 
-  processHeader (old? : Option HeaderProcessedSnapshot) (stx : Syntax) (parserState : Parser.ModuleParserState) := do
-    -- fast path, do not even start new task for this snapshot
-    if let some old := old? then
-      if let some success := old.success? then
-        return .pure { old with success? := some { success with
-          next := (← success.next.bindIO (parseCmd · parserState success.cmdState)) } }
-      else
-        return .pure old
-
-    SnapshotTask.ofIO ⟨0, ictx.input.endPos⟩ do
-    withHeaderExceptions ({ · with success? := none }) do
-    -- discard existing continuation if any, there is nothing to reuse
-    let _ ← old?.bind (·.success?) |>.map (·.next) |> getOrCancel?
-    let fileSetupResult ← if let some handler := ctx.fileSetupHandler? then
-      handler (Elab.headerToImports stx)
-    else
-      pure { kind := .success, srcSearchPath := ∅, fileOptions := .empty }
-    match fileSetupResult.kind with
-    | .importsOutOfDate =>
       return {
-        msgLog := msglogOfHeaderError
-          "Imports are out of date and must be rebuilt; use the \"Restart File\" command in your editor."
-        success? := none
+        ictx, stx, msgLog
+        success? := some {
+          parserState
+          next := (← processHeader stx parserState)
+        }
       }
-    | .error msg =>
+
+  processHeader (stx : Syntax) (parserState : Parser.ModuleParserState) :=
+    SnapshotTask.ofIO ⟨0, ictx.input.endPos⟩ <|
+    withHeaderExceptions ({ · with success? := none }) do
+      -- use `lake setup-file` if in language server
+      let fileSetupResult ← if let some handler := ctx.fileSetupHandler? then
+        handler (Elab.headerToImports stx)
+      else
+        pure { kind := .success, srcSearchPath := ∅, fileOptions := .empty }
+      match fileSetupResult.kind with
+      | .importsOutOfDate =>
+        return {
+          msgLog := msglogOfHeaderError
+            "Imports are out of date and must be rebuilt; use the \"Restart File\" command in your editor."
+          success? := none
+        }
+      | .error msg =>
+        return {
+          msgLog := msglogOfHeaderError msg
+          success? := none
+        }
+      | _ => pure ()
+
+      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 { msgLog := .empty, success? := none }
+
+      -- override context options with file options
+      let opts := ctx.opts.mergeBy (fun _ _ fileOpt => fileOpt) fileSetupResult.fileOptions
+
+      -- 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 ictx
+
+      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 ({
+          env     := headerEnv
+          fileMap := ictx.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 {
-        msgLog := msglogOfHeaderError msg
-        success? := none
-      }
-    | _ => pure ()
-
-    if (← importsLoadedRef.modifyGet ((·, true))) 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 { msgLog := .empty, success? := none }
-
-    let opts := ctx.opts.mergeBy (fun _ _ fileOpt => fileOpt) fileSetupResult.fileOptions
-    let mut (headerEnv, msgLog) ← Elab.processHeader (leakEnv := true) stx opts .empty ictx
-    headerEnv := headerEnv.setMainModule ctx.mainModuleName
-    let cmdState := Elab.Command.mkState headerEnv msgLog opts
-    let cmdState := { cmdState with infoState := {
-      enabled := true
-      trees := #[Elab.InfoTree.context ({
-        env     := headerEnv
-        fileMap := ictx.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 {
-      msgLog := cmdState.messages
-      infoTree? := cmdState.infoState.trees[0]!
-      success? := some {
-        cmdState
-        srcSearchPath := fileSetupResult.srcSearchPath
-        next := (← parseCmd none parserState cmdState)
+        msgLog := cmdState.messages
+        infoTree? := cmdState.infoState.trees[0]!
+        success? := some {
+          cmdState
+          srcSearchPath := fileSetupResult.srcSearchPath
+          next := (← parseCmd none parserState cmdState)
+        }
       }
-    }
 
-  parseCmd (old? : Option CommandParsedSnapshot) (parserState : Parser.ModuleParserState) (cmdState : Command.State) :
+  parseCmd (old? : Option CommandParsedSnapshot) (parserState : Parser.ModuleParserState)
+      (cmdState : Command.State) :
       BaseIO (SnapshotTask CommandParsedSnapshot) := do
     let unchanged old : BaseIO CommandParsedSnapshot :=
       -- when syntax is unchanged, reuse command processing task as is
@@ -254,30 +332,19 @@ where
     SnapshotTask.ofIO ⟨parserState.pos, ictx.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 pmctx := {
+        env := cmdState.env, options := scope.opts, currNamespace := scope.currNamespace
+        openDecls := scope.openDecls
+      }
       let (stx, parserState, msgLog) := Parser.parseCommand ictx pmctx parserState .empty
 
-      -- TODO: would need adjusting all contained positions
-      --if let some old := old? then
-      --  if old.data.stx == stx then
-      --    return (← unchanged old)
+      -- semi-fast path
+      if let some old := old? then
+        if firstDiffPos?.any (parserState.pos < ·) && old.data.stx == stx then
+          return (← unchanged old)
 
-      -- signature elaboration task; for now, does full elaboration
-      -- TODO: do tactic snapshots, reuse old state for them
       let _ ← old?.map (·.data.sig) |> getOrCancel?
-      let sig ← SnapshotTask.ofIO (stx.getRange?.getD ⟨parserState.pos, parserState.pos⟩) do
-        let cmdState ← doElab stx cmdState msgLog.hasErrors beginPos scope
-        return {
-          msgLog := .empty
-          -- TODO
-          tacs := #[]
-          finished := .pure {
-            msgLog := cmdState.messages
-            infoTree? := some cmdState.infoState.trees[0]!
-            cmdState
-          }
-        }
-
+      let sig ← processCmdSignature stx cmdState msgLog.hasErrors beginPos
       let next? ← if Parser.isTerminalCommand stx then pure none
         -- for now, wait on "command finished" snapshot before parsing next command
         else some <$> sig.bindIO fun sig => do
@@ -287,37 +354,52 @@ where
         msgLog
         stx
         parserState
-        stopPos
         sig
       }
 
-  doElab stx cmdState hasParseError beginPos scope := do
-    let cmdStateRef ← IO.mkRef { cmdState with messages := .empty }
-    let cmdCtx : Elab.Command.Context := { ictx with
-      cmdPos       := beginPos
-      tacticCache? := none
-    }
-    let (output, _) ← IO.FS.withIsolatedStreams (isolateStderr := stderrAsMessages.get scope.opts) <| liftM (m := BaseIO) do
-      Elab.Command.catchExceptions
-        (getResetInfoTrees *> Elab.Command.elabCommandTopLevel stx)
-        cmdCtx cmdStateRef
-    let postCmdState ← cmdStateRef.get
-    let mut msgs := postCmdState.messages
-    -- `stx.hasMissing` should imply `initMsgs.hasErrors`, but the latter should be cheaper to check in general
-    if !showPartialSyntaxErrors.get postCmdState.scopes[0]!.opts && hasParseError && 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)⟩
-    if !output.isEmpty then
-      msgs := msgs.add {
-        fileName := ictx.fileName
-        severity := MessageSeverity.information
-        pos      := ictx.fileMap.toPosition beginPos
-        data     := output
+  processCmdSignature (stx : Syntax) (cmdState : Command.State) (hasParseError : Bool)
+      (beginPos : String.Pos) :=
+    -- 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 := { ictx 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 := ictx.fileName
+          severity := MessageSeverity.information
+          pos      := ictx.fileMap.toPosition beginPos
+          data     := output
+        }
+      let cmdState := { cmdState with messages }
+      return {
+        msgLog := .empty
+        finished := .pure {
+          msgLog := cmdState.messages
+          infoTree? := some cmdState.infoState.trees[0]!
+          cmdState
+        }
       }
-    return { postCmdState with messages := msgs }
 
-partial def getFinalEnv? (snap : InitialSnapshot) : Option Environment := do
+private partial def getFinalEnv? (snap : InitialSnapshot) : Option Environment := do
   let snap ← snap.success?
   let snap ← snap.next.get.success?
   goCmd snap.next.get
@@ -329,6 +411,7 @@ where goCmd snap :=
 
 end Lean
 
+/-- The Lean language processor. -/
 abbrev Lean : Language where
   process := Lean.processLean
   getFinalEnv? := Lean.getFinalEnv?
diff --git a/src/Lean/Server/FileWorker/SetupFile.lean b/src/Lean/Server/FileWorker/SetupFile.lean
index f7ebca7df6af..27b0d78947e8 100644
--- a/src/Lean/Server/FileWorker/SetupFile.lean
+++ b/src/Lean/Server/FileWorker/SetupFile.lean
@@ -77,11 +77,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. -/

From 44cb1cc2dc3c13673f0d4533449e330a64755147 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 2 Dec 2023 11:06:01 +0100
Subject: [PATCH 20/60] better naming

---
 src/Lean/Language/Lean.lean           | 26 ++++++++++++++------------
 src/Lean/Server/FileWorker.lean       |  2 +-
 src/Lean/Server/FileWorker/Utils.lean |  2 +-
 3 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 1fb4df2731ae..35e6588ec8e1 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -151,7 +151,7 @@ structure HeaderParsedSucessfully where
   /-- Resulting parser state. -/
   parserState : Parser.ModuleParserState
   /-- Header processing task. -/
-  next : SnapshotTask HeaderProcessedSnapshot
+  processed : SnapshotTask HeaderProcessedSnapshot
 
 /-- State after the module header has been parsed. -/
 structure HeaderParsedSnapshot extends Snapshot where
@@ -162,7 +162,8 @@ structure HeaderParsedSnapshot extends Snapshot where
   /-- State after successful parsing. -/
   success? : Option HeaderParsedSucessfully
 instance : ToSnapshotTree HeaderParsedSnapshot where
-  toSnapshotTree s := ⟨s.toSnapshot, #[] |> pushOpt (s.success?.map (·.next.map toSnapshotTree))⟩
+  toSnapshotTree s := ⟨s.toSnapshot,
+    #[] |> pushOpt (s.success?.map (·.processed.map toSnapshotTree))⟩
 
 /-- Initial snapshot of the Lean language processor: a "header parsed" snapshot. -/
 abbrev InitialSnapshot := HeaderParsedSnapshot
@@ -208,7 +209,7 @@ where
       -- when header syntax is unchanged, reuse import processing task as is and continue with
       -- parsing
       return { old with success? := some { success with
-        next := (← success.next.bindIO fun processed => do
+        processed := (← success.processed.bindIO fun processed => do
           if let some procSuccess := processed.success? then
             let oldCmd? ← getOrCancel? procSuccess.next
             return .pure { processed with success? := some { procSuccess with
@@ -216,13 +217,14 @@ where
           else
             return .pure processed) } }
 
-    -- fast path: if we have parsed the header sucessfully...
+    -- fast path: if we have parsed the header successfully...
     if let some old@{ success? := some success, .. } := old? then
-      -- ...and the edit location is after the next command...
-      if let some nextCom ← (← success.next.get?).bind (·.success?) |>.bindM (·.next.get?) then
-        if firstDiffPos?.any (nextCom.data.parserState.pos < ·) then
-          -- ...go immediately to next snapshot
-          return (← unchanged old success)
+      if let some processed ← success.processed.get? then
+        -- ...and the edit location is after the next command...
+        if let some nextCom ← processed.success? |>.bindM (·.next.get?) then
+          if firstDiffPos?.any (nextCom.data.parserState.pos < ·) then
+            -- ...go immediately to next snapshot
+            return (← unchanged old success)
 
     withHeaderExceptions ({ · with ictx, stx := .missing, success? := none }) do
       let (stx, parserState, msgLog) ← Parser.parseHeader ictx
@@ -235,13 +237,13 @@ where
       if let some old@{ success? := some success, .. } := old? then
         if firstDiffPos?.any (parserState.pos < ·) && old.stx == stx then
           return (← unchanged old success)
-        success.next.cancel
+        success.processed.cancel
 
       return {
         ictx, stx, msgLog
         success? := some {
           parserState
-          next := (← processHeader stx parserState)
+          processed := (← processHeader stx parserState)
         }
       }
 
@@ -401,7 +403,7 @@ where
 
 private partial def getFinalEnv? (snap : InitialSnapshot) : Option Environment := do
   let snap ← snap.success?
-  let snap ← snap.next.get.success?
+  let snap ← snap.processed.get.success?
   goCmd snap.next.get
 where goCmd snap :=
   if let some next := snap.next? then
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 69e0cc2c9427..90ed75063b16 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -350,7 +350,7 @@ section MessageHandling
     -- we assume that any other request requires at least the the search path
     -- TODO: move into language-specific request handling
     let srcSearchPathTask :=
-      st.doc.initSnap.success?.map (·.next.task.map (·.success?.map (·.srcSearchPath) |>.getD ∅))
+      st.doc.initSnap.success?.map (·.processed.task.map (·.success?.map (·.srcSearchPath) |>.getD ∅))
       |>.getD (.pure ∅)
     let t ← IO.bindTask srcSearchPathTask fun srcSearchPath => do
      let rc : RequestContext :=
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index bea7f1069a28..7e00a2e8dc5c 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -45,7 +45,7 @@ end CancelToken
 private partial def mkCmdSnaps (initSnap : Language.Lean.InitialSnapshot) :
     AsyncList ElabTaskError Snapshot := Id.run do
   let some headerParsed := initSnap.success? | return .nil
-  .delayed <| headerParsed.next.task.bind fun headerProcessed => Id.run do
+  .delayed <| headerParsed.processed.task.bind fun headerProcessed => Id.run do
     let some headerSuccess := headerProcessed.success? | return .pure <| .ok .nil
     return .pure <| .ok <| .cons {
       stx := initSnap.stx

From 057184b70980415962e2ea809ca97b28b1bae1c6 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 2 Dec 2023 11:56:37 +0100
Subject: [PATCH 21/60] feat: synchronous execution of task continuations

---
 src/Init/Core.lean      |  7 +++++--
 src/Init/System/IO.lean | 42 +++++++++++++++++++++++++----------------
 src/include/lean/lean.h | 12 ++++++------
 src/runtime/io.cpp      | 14 ++++++++------
 src/runtime/object.cpp  | 10 ++++++----
 src/runtime/object.h    |  4 ++--
 6 files changed, 53 insertions(+), 36 deletions(-)

diff --git a/src/Init/Core.lean b/src/Init/Core.lean
index f9899767cf41..47d23616f3a8 100644
--- a/src/Init/Core.lean
+++ b/src/Init/Core.lean
@@ -411,9 +411,10 @@ set_option linter.unusedVariables.funArgs false in
 be available and then calls `f` on the result.
 
 `prio`, if provided, is the priority of the task.
+If `sync` is set to true, `f` is executed on the current thread if `x` has already finished.
 -/
 @[noinline, extern "lean_task_map"]
-protected def map {α : Type u} {β : Type v} (f : α → β) (x : Task α) (prio := Priority.default) : Task β :=
+protected def map (f : α → β) (x : Task α) (prio := Priority.default) (sync := false) : Task β :=
   ⟨f x.get⟩
 
 set_option linter.unusedVariables.funArgs false in
@@ -424,9 +425,11 @@ for the value of `x` to be available and then calls `f` on the result,
 resulting in a new task which is then run for a result.
 
 `prio`, if provided, is the priority of the task.
+If `sync` is set to true, `f` is executed on the current thread if `x` has already finished.
 -/
 @[noinline, extern "lean_task_bind"]
-protected def bind {α : Type u} {β : Type v} (x : Task α) (f : α → Task β) (prio := Priority.default) : Task β :=
+protected def bind (x : Task α) (f : α → Task β) (prio := Priority.default) (sync := false) :
+    Task β :=
   ⟨(f x.get).get⟩
 
 end Task
diff --git a/src/Init/System/IO.lean b/src/Init/System/IO.lean
index 85ed35dbb246..31a0f5b2a849 100644
--- a/src/Init/System/IO.lean
+++ b/src/Init/System/IO.lean
@@ -117,20 +117,23 @@ opaque asTask (act : BaseIO α) (prio := Task.Priority.default) : BaseIO (Task 
 
 /-- See `BaseIO.asTask`. -/
 @[extern "lean_io_map_task"]
-opaque mapTask (f : α → BaseIO β) (t : Task α) (prio := Task.Priority.default) : BaseIO (Task β) :=
+opaque mapTask (f : α → BaseIO β) (t : Task α) (prio := Task.Priority.default) (sync := false) :
+    BaseIO (Task β) :=
   Task.pure <$> f t.get
 
 /-- See `BaseIO.asTask`. -/
 @[extern "lean_io_bind_task"]
-opaque bindTask (t : Task α) (f : α → BaseIO (Task β)) (prio := Task.Priority.default) : BaseIO (Task β) :=
+opaque bindTask (t : Task α) (f : α → BaseIO (Task β)) (prio := Task.Priority.default)
+    (sync := false) : BaseIO (Task β) :=
   f t.get
 
-def mapTasks (f : List α → BaseIO β) (tasks : List (Task α)) (prio := Task.Priority.default) : BaseIO (Task β) :=
+def mapTasks (f : List α → BaseIO β) (tasks : List (Task α)) (prio := Task.Priority.default)
+    (sync := false) : BaseIO (Task β) :=
   go tasks []
 where
   go
     | t::ts, as =>
-      BaseIO.bindTask t (fun a => go ts (a :: as)) prio
+      BaseIO.bindTask t (fun a => go ts (a :: as)) prio sync
     | [], as => f as.reverse |>.asTask prio
 
 end BaseIO
@@ -142,16 +145,20 @@ namespace EIO
   act.toBaseIO.asTask prio
 
 /-- `EIO` specialization of `BaseIO.mapTask`. -/
-@[inline] def mapTask (f : α → EIO ε β) (t : Task α) (prio := Task.Priority.default) : BaseIO (Task (Except ε β)) :=
-  BaseIO.mapTask (fun a => f a |>.toBaseIO) t prio
+@[inline] def mapTask (f : α → EIO ε β) (t : Task α) (prio := Task.Priority.default)
+    (sync := false) : BaseIO (Task (Except ε β)) :=
+  BaseIO.mapTask (fun a => f a |>.toBaseIO) t prio sync
 
 /-- `EIO` specialization of `BaseIO.bindTask`. -/
-@[inline] def bindTask (t : Task α) (f : α → EIO ε (Task (Except ε β))) (prio := Task.Priority.default) : BaseIO (Task (Except ε β)) :=
-  BaseIO.bindTask t (fun a => f a |>.catchExceptions fun e => return Task.pure <| Except.error e) prio
+@[inline] def bindTask (t : Task α) (f : α → EIO ε (Task (Except ε β)))
+    (prio := Task.Priority.default) (sync := false) : BaseIO (Task (Except ε β)) :=
+  BaseIO.bindTask t (fun a => f a |>.catchExceptions fun e => return Task.pure <| Except.error e)
+    prio sync
 
 /-- `EIO` specialization of `BaseIO.mapTasks`. -/
-@[inline] def mapTasks (f : List α → EIO ε β) (tasks : List (Task α)) (prio := Task.Priority.default) : BaseIO (Task (Except ε β)) :=
-  BaseIO.mapTasks (fun as => f as |>.toBaseIO) tasks prio
+@[inline] def mapTasks (f : List α → EIO ε β) (tasks : List (Task α))
+    (prio := Task.Priority.default) (sync := false) : BaseIO (Task (Except ε β)) :=
+  BaseIO.mapTasks (fun as => f as |>.toBaseIO) tasks prio sync
 
 end EIO
 
@@ -184,16 +191,19 @@ def sleep (ms : UInt32) : BaseIO Unit :=
   EIO.asTask act prio
 
 /-- `IO` specialization of `EIO.mapTask`. -/
-@[inline] def mapTask (f : α → IO β) (t : Task α) (prio := Task.Priority.default) : BaseIO (Task (Except IO.Error β)) :=
-  EIO.mapTask f t prio
+@[inline] def mapTask (f : α → IO β) (t : Task α) (prio := Task.Priority.default) (sync := false) :
+    BaseIO (Task (Except IO.Error β)) :=
+  EIO.mapTask f t prio sync
 
 /-- `IO` specialization of `EIO.bindTask`. -/
-@[inline] def bindTask (t : Task α) (f : α → IO (Task (Except IO.Error β))) (prio := Task.Priority.default) : BaseIO (Task (Except IO.Error β)) :=
-  EIO.bindTask t f prio
+@[inline] def bindTask (t : Task α) (f : α → IO (Task (Except IO.Error β)))
+    (prio := Task.Priority.default) (sync := false) : BaseIO (Task (Except IO.Error β)) :=
+  EIO.bindTask t f prio sync
 
 /-- `IO` specialization of `EIO.mapTasks`. -/
-@[inline] def mapTasks (f : List α → IO β) (tasks : List (Task α)) (prio := Task.Priority.default) : BaseIO (Task (Except IO.Error β)) :=
-  EIO.mapTasks f tasks prio
+@[inline] def mapTasks (f : List α → IO β) (tasks : List (Task α)) (prio := Task.Priority.default)
+    (sync := false) : BaseIO (Task (Except IO.Error β)) :=
+  EIO.mapTasks f tasks prio sync
 
 /-- Check if the task's cancellation flag has been set by calling `IO.cancel` or dropping the last reference to the task. -/
 @[extern "lean_io_check_canceled"] opaque checkCanceled : BaseIO Bool
diff --git a/src/include/lean/lean.h b/src/include/lean/lean.h
index 4d339e7f7ada..08946c2557ae 100644
--- a/src/include/lean/lean.h
+++ b/src/include/lean/lean.h
@@ -1092,12 +1092,12 @@ LEAN_SHARED lean_obj_res lean_task_spawn_core(lean_obj_arg c, unsigned prio, boo
 static inline lean_obj_res lean_task_spawn(lean_obj_arg c, lean_obj_arg prio) { return lean_task_spawn_core(c, lean_unbox(prio), false); }
 /* Convert a value `a : A` into `Task A` */
 LEAN_SHARED lean_obj_res lean_task_pure(lean_obj_arg a);
-LEAN_SHARED lean_obj_res lean_task_bind_core(lean_obj_arg x, lean_obj_arg f, unsigned prio, bool keep_alive);
-/* Task.bind (x : Task A) (f : A -> Task B) (prio : Nat) : Task B */
-static inline lean_obj_res lean_task_bind(lean_obj_arg x, lean_obj_arg f, lean_obj_arg prio) { return lean_task_bind_core(x, f, lean_unbox(prio), false); }
-LEAN_SHARED lean_obj_res lean_task_map_core(lean_obj_arg f, lean_obj_arg t, unsigned prio, bool keep_alive);
-/* Task.map (f : A -> B) (t : Task A) (prio : Nat) : Task B */
-static inline lean_obj_res lean_task_map(lean_obj_arg f, lean_obj_arg t, lean_obj_arg prio) { return lean_task_map_core(f, t, lean_unbox(prio), false); }
+LEAN_SHARED lean_obj_res lean_task_bind_core(lean_obj_arg x, lean_obj_arg f, unsigned prio, bool sync, bool keep_alive);
+/* Task.bind (x : Task A) (f : A -> Task B) (prio : Nat) (sync : Bool) : Task B */
+static inline lean_obj_res lean_task_bind(lean_obj_arg x, lean_obj_arg f, lean_obj_arg prio, uint8_t sync) { return lean_task_bind_core(x, f, lean_unbox(prio), sync, false); }
+LEAN_SHARED lean_obj_res lean_task_map_core(lean_obj_arg f, lean_obj_arg t, unsigned prio, bool sync, bool keep_alive);
+/* Task.map (f : A -> B) (t : Task A) (prio : Nat) (sync : Bool) : Task B */
+static inline lean_obj_res lean_task_map(lean_obj_arg f, lean_obj_arg t, lean_obj_arg prio, uint8_t sync) { return lean_task_map_core(f, t, lean_unbox(prio), sync, false); }
 LEAN_SHARED b_lean_obj_res lean_task_get(b_lean_obj_arg t);
 /* Primitive for implementing Task.get : Task A -> A */
 static inline lean_obj_res lean_task_get_own(lean_obj_arg t) {
diff --git a/src/runtime/io.cpp b/src/runtime/io.cpp
index 207e33b8c494..85fa4820d6d1 100644
--- a/src/runtime/io.cpp
+++ b/src/runtime/io.cpp
@@ -1011,19 +1011,21 @@ static obj_res lean_io_bind_task_fn(obj_arg f, obj_arg a) {
     return object_ref(io_result_get_value(r.raw()), true).steal();
 }
 
-/*  mapTask {α : Type u} {β : Type} (f : α → BaseIO β) (t : Task α) (prio : Nat) : BaseIO (Task β) */
-extern "C" LEAN_EXPORT obj_res lean_io_map_task(obj_arg f, obj_arg t, obj_arg prio, obj_arg) {
+/*  mapTask (f : α → BaseIO β) (t : Task α) (prio : Nat) (sync : Bool) : BaseIO (Task β) */
+extern "C" LEAN_EXPORT obj_res lean_io_map_task(obj_arg f, obj_arg t, obj_arg prio, uint8 sync,
+        obj_arg) {
     object * c = lean_alloc_closure((void*)lean_io_bind_task_fn, 2, 1);
     lean_closure_set(c, 0, f);
-    object * t2 = lean_task_map_core(c, t, lean_unbox(prio), /* keep_alive */ true);
+    object * t2 = lean_task_map_core(c, t, lean_unbox(prio), sync, /* keep_alive */ true);
     return io_result_mk_ok(t2);
 }
 
-/*  bindTask {α : Type u} {β : Type} (t : Task α) (f : α → BaseIO (Task β)) (prio : Nat) : BaseIO (Task β) */
-extern "C" LEAN_EXPORT obj_res lean_io_bind_task(obj_arg t, obj_arg f, obj_arg prio, obj_arg) {
+/*  bindTask (t : Task α) (f : α → BaseIO (Task β)) (prio : Nat) (sync : Bool) : BaseIO (Task β) */
+extern "C" LEAN_EXPORT obj_res lean_io_bind_task(obj_arg t, obj_arg f, obj_arg prio, uint8 sync,
+        obj_arg) {
     object * c = lean_alloc_closure((void*)lean_io_bind_task_fn, 2, 1);
     lean_closure_set(c, 0, f);
-    object * t2 = lean_task_bind_core(t, c, lean_unbox(prio), /* keep_alive */ true);
+    object * t2 = lean_task_bind_core(t, c, lean_unbox(prio), sync, /* keep_alive */ true);
     return io_result_mk_ok(t2);
 }
 
diff --git a/src/runtime/object.cpp b/src/runtime/object.cpp
index 204156f130db..8227e340a53b 100644
--- a/src/runtime/object.cpp
+++ b/src/runtime/object.cpp
@@ -954,8 +954,9 @@ static obj_res task_map_fn(obj_arg f, obj_arg t, obj_arg) {
     return lean_apply_1(f, v);
 }
 
-extern "C" LEAN_EXPORT obj_res lean_task_map_core(obj_arg f, obj_arg t, unsigned prio, bool keep_alive) {
-    if (!g_task_manager) {
+extern "C" LEAN_EXPORT obj_res lean_task_map_core(obj_arg f, obj_arg t, unsigned prio,
+      bool sync, bool keep_alive) {
+    if (!g_task_manager || (sync && lean_to_task(t)->m_value)) {
         return lean_task_pure(apply_1(f, lean_task_get_own(t)));
     } else {
         lean_task_object * new_task = alloc_task(mk_closure_3_2(task_map_fn, f, t), prio, keep_alive);
@@ -996,8 +997,9 @@ static obj_res task_bind_fn1(obj_arg x, obj_arg f, obj_arg) {
     return nullptr; /* notify queue that task did not finish yet. */
 }
 
-extern "C" LEAN_EXPORT obj_res lean_task_bind_core(obj_arg x, obj_arg f, unsigned prio, bool keep_alive) {
-    if (!g_task_manager) {
+extern "C" LEAN_EXPORT obj_res lean_task_bind_core(obj_arg x, obj_arg f, unsigned prio,
+      bool sync, bool keep_alive) {
+    if (!g_task_manager || (sync && lean_to_task(x)->m_value)) {
         return apply_1(f, lean_task_get_own(x));
     } else {
         lean_task_object * new_task = alloc_task(mk_closure_3_2(task_bind_fn1, x, f), prio, keep_alive);
diff --git a/src/runtime/object.h b/src/runtime/object.h
index 9a057383e3dd..f93ccd621e8c 100644
--- a/src/runtime/object.h
+++ b/src/runtime/object.h
@@ -283,8 +283,8 @@ class scoped_task_manager {
 
 inline obj_res task_spawn(obj_arg c, unsigned prio = 0, bool keep_alive = false) { return lean_task_spawn_core(c, prio, keep_alive); }
 inline obj_res task_pure(obj_arg a) { return lean_task_pure(a); }
-inline obj_res task_bind(obj_arg x, obj_arg f, unsigned prio = 0, bool keep_alive = false) { return lean_task_bind_core(x, f, prio, keep_alive); }
-inline obj_res task_map(obj_arg f, obj_arg t, unsigned prio = 0, bool keep_alive = false) { return lean_task_map_core(f, t, prio, keep_alive); }
+inline obj_res task_bind(obj_arg x, obj_arg f, unsigned prio = 0, bool sync = false, bool keep_alive = false) { return lean_task_bind_core(x, f, prio, sync, keep_alive); }
+inline obj_res task_map(obj_arg f, obj_arg t, unsigned prio = 0, bool sync = false, bool keep_alive = false) { return lean_task_map_core(f, t, prio, sync, keep_alive); }
 inline b_obj_res task_get(b_obj_arg t) { return lean_task_get(t); }
 
 inline bool io_check_canceled_core() { return lean_io_check_canceled_core(); }

From 0688e2219fd34ca99afa5b54fc74d97113bc2e68 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 2 Dec 2023 12:18:14 +0100
Subject: [PATCH 22/60] synchronous fast forwarding in Lean processor

---
 src/Lean/Language/Basic.lean |  8 ++++----
 src/Lean/Language/Lean.lean  | 16 +++++++++-------
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index ca8f0217bf69..dcd36a0b2c27 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -85,15 +85,15 @@ 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 : α → β) : SnapshotTask β :=
-  { t with task := t.task.map f }
+def SnapshotTask.map (t : SnapshotTask α) (f : α → β) (sync := false) : SnapshotTask β :=
+  { t with 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.bindIO (t : SnapshotTask α) (act : α → BaseIO (SnapshotTask β))
-    (range : String.Range := t.range) : BaseIO (SnapshotTask β) :=
-  return { range, task := (← BaseIO.bindTask t.task fun a => (·.task) <$> (act a)) }
+    (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 α) : α :=
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 35e6588ec8e1..df11895d169d 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -100,7 +100,7 @@ structure CommandSignatureProcessedSnapshot extends Snapshot where
   finished : SnapshotTask CommandFinishedSnapshot
 deriving Nonempty
 instance : ToSnapshotTree CommandSignatureProcessedSnapshot where
-  toSnapshotTree s := ⟨s.toSnapshot, #[s.finished.map toSnapshotTree]⟩
+  toSnapshotTree s := ⟨s.toSnapshot, #[s.finished.map (sync := true) toSnapshotTree]⟩
 
 /-- State after a command has been parsed. -/
 structure CommandParsedSnapshotData extends Snapshot where
@@ -128,7 +128,8 @@ abbrev CommandParsedSnapshot.next? : CommandParsedSnapshot →
 partial instance : ToSnapshotTree CommandParsedSnapshot where
   toSnapshotTree := go where
     go s := ⟨s.data.toSnapshot,
-      #[s.data.sig.map toSnapshotTree] |> pushOpt (s.next?.map (·.map go))⟩
+      #[s.data.sig.map (sync := true) toSnapshotTree] |>
+        pushOpt (s.next?.map (·.map (sync := true) go))⟩
 
 /-- State after successful importing. -/
 structure HeaderProcessedSucessfully where
@@ -144,7 +145,8 @@ structure HeaderProcessedSnapshot extends Snapshot where
   /-- State after successful importing. -/
   success? : Option HeaderProcessedSucessfully
 instance : ToSnapshotTree HeaderProcessedSnapshot where
-  toSnapshotTree s := ⟨s.toSnapshot, #[] |> pushOpt (s.success?.map (·.next.map toSnapshotTree))⟩
+  toSnapshotTree s := ⟨s.toSnapshot, #[] |>
+    pushOpt (s.success?.map (·.next.map (sync := true) toSnapshotTree))⟩
 
 /-- State after successfully parsing the module header. -/
 structure HeaderParsedSucessfully where
@@ -163,7 +165,7 @@ structure HeaderParsedSnapshot extends Snapshot where
   success? : Option HeaderParsedSucessfully
 instance : ToSnapshotTree HeaderParsedSnapshot where
   toSnapshotTree s := ⟨s.toSnapshot,
-    #[] |> pushOpt (s.success?.map (·.processed.map toSnapshotTree))⟩
+    #[] |> pushOpt (s.success?.map (·.processed.map (sync := true) toSnapshotTree))⟩
 
 /-- Initial snapshot of the Lean language processor: a "header parsed" snapshot. -/
 abbrev InitialSnapshot := HeaderParsedSnapshot
@@ -209,7 +211,7 @@ where
       -- when header syntax is unchanged, reuse import processing task as is and continue with
       -- parsing
       return { old with success? := some { success with
-        processed := (← success.processed.bindIO fun processed => do
+        processed := (← success.processed.bindIO (sync := true) fun processed => do
           if let some procSuccess := processed.success? then
             let oldCmd? ← getOrCancel? procSuccess.next
             return .pure { processed with success? := some { procSuccess with
@@ -320,8 +322,8 @@ where
       -- when syntax is unchanged, reuse command processing task as is
       if let some oldNext := old.next? then
         return .mk (data := old.data)
-          (next? := (← old.data.sig.bindIO fun sig =>
-            sig.finished.bindIO fun finished => do
+          (next? := (← old.data.sig.bindIO (sync := true) fun sig =>
+            sig.finished.bindIO (sync := true) fun finished => do
               parseCmd (← getOrCancel? oldNext) old.data.parserState finished.cmdState))
       else return old  -- terminal command, we're done!
 

From d0e9409196b0a08f21f4a65a1d7acc9a47cb4059 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 2 Dec 2023 14:58:13 +0100
Subject: [PATCH 23/60] Bring back trust level

---
 src/Lean/Elab/Frontend.lean  | 5 +++--
 src/Lean/Language/Basic.lean | 2 ++
 src/Lean/Language/Lean.lean  | 3 ++-
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index d8113bc7f243..c9f8b1bdf77f 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -93,13 +93,14 @@ def runFrontend
     (fileName : String)
     (mainModuleName : Name)
     -- TODO: do we still want this in the driver?
-    (_trustLevel : UInt32 := 0)
+    (trustLevel : UInt32 := 0)
     (ileanFileName? : Option String := none)
     : IO (Environment × Bool) := do
   let inputCtx := Parser.mkInputContext input fileName
   -- TODO: replace with `#lang` processing
   let lang := Language.Lean
-  let snap ← lang.process { mainModuleName, opts, fileSetupHandler? := none } none inputCtx
+  let ctx := { mainModuleName, opts, trustLevel, fileSetupHandler? := none }
+  let snap ← lang.process ctx none inputCtx
   let snaps := Language.toSnapshotTree snap
   snaps.runAndReport opts
   if let some ileanFileName := ileanFileName? then
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index dcd36a0b2c27..e5d9b3a3e46d 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -171,6 +171,8 @@ structure ProcessingContext where
   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
   /--
     Callback available in server mode for building imports and retrieving per-library options using
     `lake setup-file`. -/
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index df11895d169d..a4b939477828 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -210,7 +210,7 @@ where
     let unchanged old success :=
       -- when header syntax is unchanged, reuse import processing task as is and continue with
       -- parsing
-      return { old with success? := some { success with
+      return { old with ictx, success? := some { success with
         processed := (← success.processed.bindIO (sync := true) fun processed => do
           if let some procSuccess := processed.success? then
             let oldCmd? ← getOrCancel? procSuccess.next
@@ -286,6 +286,7 @@ where
 
       -- 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 ictx
+        ctx.trustLevel
 
       let headerEnv := headerEnv.setMainModule ctx.mainModuleName
       let cmdState := Elab.Command.mkState headerEnv msgLog opts

From 3930e400f5856068781ad89dfd931349847b4f28 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 2 Dec 2023 15:25:36 +0100
Subject: [PATCH 24/60] Try a shortcut

---
 src/Lean/Language/Lean.lean     | 40 +++++++++++++++++++--------------
 src/Lean/Server/FileWorker.lean | 37 +++++++++++++++---------------
 2 files changed, 41 insertions(+), 36 deletions(-)

diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index a4b939477828..45a782c9a093 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -167,6 +167,11 @@ instance : ToSnapshotTree HeaderParsedSnapshot where
   toSnapshotTree s := ⟨s.toSnapshot,
     #[] |> pushOpt (s.success?.map (·.processed.map (sync := true) toSnapshotTree))⟩
 
+/-- Shortcut accessor to the final header state, if successful. -/
+def HeaderParsedSnapshot.processedSuccessfully (snap : HeaderParsedSnapshot) :
+    SnapshotTask (Option HeaderProcessedSucessfully) :=
+  snap.success?.bind (·.processed.map (sync := true) (·.success?)) |>.getD (.pure none)
+
 /-- Initial snapshot of the Lean language processor: a "header parsed" snapshot. -/
 abbrev InitialSnapshot := HeaderParsedSnapshot
 
@@ -207,26 +212,28 @@ where
   firstDiffPos? := old?.map (·.ictx.input.firstDiffPos ictx.input)
 
   parseHeader (old? : Option HeaderParsedSnapshot) := do
-    let unchanged old success :=
+    let unchanged old :=
       -- when header syntax is unchanged, reuse import processing task as is and continue with
-      -- parsing
-      return { old with ictx, success? := some { success with
-        processed := (← success.processed.bindIO (sync := true) fun processed => do
-          if let some procSuccess := processed.success? then
-            let oldCmd? ← getOrCancel? procSuccess.next
-            return .pure { processed with success? := some { procSuccess with
-              next := (← parseCmd oldCmd? success.parserState procSuccess.cmdState) } }
-          else
-            return .pure processed) } }
+      -- parsing the first command
+      if let some success := old.success? then
+        return { old with ictx, success? := some { success with
+          processed := (← success.processed.bindIO (sync := true) fun processed => do
+            if let some procSuccess := processed.success? then
+              let oldCmd? ← getOrCancel? procSuccess.next
+              return .pure { processed with success? := some { procSuccess with
+                next := (← parseCmd oldCmd? success.parserState procSuccess.cmdState) } }
+            else
+              return .pure processed) } }
+      else return old
 
     -- fast path: if we have parsed the header successfully...
-    if let some old@{ success? := some success, .. } := old? then
-      if let some processed ← success.processed.get? then
+    if let some old := old? then
+      if let some (some processed) ← old.processedSuccessfully.get? then
         -- ...and the edit location is after the next command...
-        if let some nextCom ← processed.success? |>.bindM (·.next.get?) then
+        if let some nextCom ← processed.next.get? then
           if firstDiffPos?.any (nextCom.data.parserState.pos < ·) then
             -- ...go immediately to next snapshot
-            return (← unchanged old success)
+            return (← unchanged old)
 
     withHeaderExceptions ({ · with ictx, stx := .missing, success? := none }) do
       let (stx, parserState, msgLog) ← Parser.parseHeader ictx
@@ -236,10 +243,9 @@ where
       -- 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
-      if let some old@{ success? := some success, .. } := old? then
+      if let some old := old? then
         if firstDiffPos?.any (parserState.pos < ·) && old.stx == stx then
-          return (← unchanged old success)
-        success.processed.cancel
+          return (← unchanged old)
 
       return {
         ictx, stx, msgLog
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 90ed75063b16..5df2220d1759 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -350,25 +350,24 @@ section MessageHandling
     -- we assume that any other request requires at least the the search path
     -- TODO: move into language-specific request handling
     let srcSearchPathTask :=
-      st.doc.initSnap.success?.map (·.processed.task.map (·.success?.map (·.srcSearchPath) |>.getD ∅))
-      |>.getD (.pure ∅)
-    let t ← IO.bindTask srcSearchPathTask fun srcSearchPath => do
-     let rc : RequestContext :=
-       { rpcSessions := st.rpcSessions
-         srcSearchPath
-         doc := st.doc
-         hLog := ctx.hLog
-         initParams := ctx.initParams }
-     let t? ← EIO.toIO' <| handleLspRequest method params rc
-     let t₁ ← match t? with
-       | Except.error e =>
-         IO.asTask do
-           ctx.chanOut.send <| e.toLspResponseError id
-       | Except.ok t => (IO.mapTask · t) fun
-         | Except.ok resp =>
-           ctx.chanOut.send <| .response id (toJson resp)
-         | Except.error e =>
-           ctx.chanOut.send <| e.toLspResponseError id
+      st.doc.initSnap.processedSuccessfully.map (·.map (·.srcSearchPath) |>.getD ∅)
+    let t ← IO.bindTask srcSearchPathTask.task fun srcSearchPath => do
+      let rc : RequestContext :=
+        { rpcSessions := st.rpcSessions
+          srcSearchPath
+          doc := st.doc
+          hLog := ctx.hLog
+          initParams := ctx.initParams }
+      let t? ← EIO.toIO' <| handleLspRequest method params rc
+      let t₁ ← match t? with
+        | Except.error e =>
+          IO.asTask do
+            ctx.chanOut.send <| e.toLspResponseError id
+        | Except.ok t => (IO.mapTask · t) fun
+          | Except.ok resp =>
+            ctx.chanOut.send <| .response id (toJson resp)
+          | Except.error e =>
+            ctx.chanOut.send <| e.toLspResponseError id
     queueRequest id t
 end MessageHandling
 

From 3b14598abf548434db512c30e3a96361de2ddb91 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 2 Dec 2023 15:28:44 +0100
Subject: [PATCH 25/60] Remove redundant fields

---
 src/Lean/Server/FileWorker.lean | 13 +++----------
 1 file changed, 3 insertions(+), 10 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 5df2220d1759..b62dc08bbd37 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -142,11 +142,6 @@ 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
   importCachingTask? : Option (Task (Except Error AvailableImportsCache))
   pendingRequests    : PendingRequestMap
   /-- A map of RPC session IDs. We allow asynchronous elab tasks and request handlers
@@ -189,8 +184,6 @@ section Initialization
     let doc : EditableDocument := { meta, initSnap, reporter }
     return (ctx,
     { doc
-      initHeaderStx      := initSnap.stx
-      currHeaderStx      := initSnap.stx
       pendingRequests    := RBMap.empty
       rpcSessions        := RBMap.empty
       importCachingTask? := none
@@ -312,7 +305,7 @@ section MessageHandling
     | none => IO.asTask do
       let availableImports ← ImportCompletion.collectAvailableImports
       let lastRequestTimestampMs ← IO.monoMsNow
-      let completions := ImportCompletion.find text st.currHeaderStx params availableImports
+      let completions := ImportCompletion.find text st.doc.initSnap.stx params availableImports
       ctx.chanOut.send <| .response id (toJson completions)
       pure { availableImports, lastRequestTimestampMs : AvailableImportsCache }
 
@@ -322,7 +315,7 @@ section MessageHandling
       if timestampNowMs - lastRequestTimestampMs >= 10000 then
         availableImports ← ImportCompletion.collectAvailableImports
       lastRequestTimestampMs := timestampNowMs
-      let completions := ImportCompletion.find text st.currHeaderStx params availableImports
+      let completions := ImportCompletion.find text st.doc.initSnap.stx params availableImports
       ctx.chanOut.send <| .response id (toJson completions)
       pure { availableImports, lastRequestTimestampMs : AvailableImportsCache }
 
@@ -342,7 +335,7 @@ section MessageHandling
 
     if method == "textDocument/completion" then
       let params ← parseParams CompletionParams params
-      if ImportCompletion.isImportCompletionRequest st.doc.meta.text st.currHeaderStx params then
+      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

From daf6fde5b8699898d5d95e52605bbfe415ca8ea5 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 2 Dec 2023 15:51:41 +0100
Subject: [PATCH 26/60] doc

---
 src/Lean/Server/FileWorker/Utils.lean | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index 7e00a2e8dc5c..c57898fa6df5 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -42,6 +42,7 @@ def set (tk : CancelToken) : IO Unit :=
 
 end CancelToken
 
+-- TEMP: translate from new heterogeneous snapshot tree to old homegenous async list
 private partial def mkCmdSnaps (initSnap : Language.Lean.InitialSnapshot) :
     AsyncList ElabTaskError Snapshot := Id.run do
   let some headerParsed := initSnap.success? | return .nil

From c6dd2b1027cf2f6c96ef94190eec5b40416000e4 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 2 Dec 2023 17:03:01 +0100
Subject: [PATCH 27/60] code owner

---
 CODEOWNERS | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CODEOWNERS b/CODEOWNERS
index 7ef90703b00d..20400858ab22 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -14,6 +14,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

From f89de7f27c5cdeaf1795c1581ca98b64d3f1b794 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 2 Dec 2023 17:07:40 +0100
Subject: [PATCH 28/60] doc

---
 src/Lean/Language/Lean.lean | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 45a782c9a093..13322ab68d4d 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -214,7 +214,7 @@ where
   parseHeader (old? : Option HeaderParsedSnapshot) := do
     let unchanged old :=
       -- when header syntax is unchanged, reuse import processing task as is and continue with
-      -- parsing the first command
+      -- parsing the first command, synchronously if possible
       if let some success := old.success? then
         return { old with ictx, success? := some { success with
           processed := (← success.processed.bindIO (sync := true) fun processed => do
@@ -229,13 +229,14 @@ where
     -- fast path: if we have parsed the header successfully...
     if let some old := old? then
       if let some (some processed) ← old.processedSuccessfully.get? then
-        -- ...and the edit location is after the next command...
+        -- ...and the edit location is after the next command (see note [Incremental Parsing])...
         if let some nextCom ← processed.next.get? then
           if firstDiffPos?.any (nextCom.data.parserState.pos < ·) then
             -- ...go immediately to next snapshot
             return (← unchanged old)
 
     withHeaderExceptions ({ · with ictx, stx := .missing, success? := 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, msgLog, success? := none }

From 55506231609483f2ff2dbd34d995e49d512e24bb Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Tue, 12 Dec 2023 15:43:52 +0100
Subject: [PATCH 29/60] Fix explicit cancellation

---
 src/Init/Data/Option/Basic.lean       |  4 ++
 src/Lean/Language/Basic.lean          | 25 +++++------
 src/Lean/Language/Lean.lean           | 61 ++++++++++++++++++++-------
 src/Lean/Server/FileWorker/Utils.lean |  2 +-
 4 files changed, 62 insertions(+), 30 deletions(-)

diff --git a/src/Init/Data/Option/Basic.lean b/src/Init/Data/Option/Basic.lean
index c07201c8e697..84d3ed7e272c 100644
--- a/src/Init/Data/Option/Basic.lean
+++ b/src/Init/Data/Option/Basic.lean
@@ -46,6 +46,10 @@ def toMonad [Monad m] [Alternative m] : Option α → m α
   else
     return none
 
+@[inline] protected def forM [Monad m] (f : α → m Unit) : Option α → m Unit
+  | some a => f a
+  | none => pure ()
+
 theorem map_id : (Option.map id : Option α → Option α) = id :=
   funext (fun o => match o with | none => rfl | some _ => rfl)
 
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index e5d9b3a3e46d..a2445aa2136e 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -80,13 +80,21 @@ def SnapshotTask.pure (a : α) : SnapshotTask α where
 
 /--
   Explicitly cancels a tasks. Like with basic `Tasks`s, cancellation happens implicitly when the
-  last reference to the task is dropped. -/
+  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 : α → β) (sync := false) : SnapshotTask β :=
-  { t with task := t.task.map (sync := sync) f }
+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
@@ -154,17 +162,6 @@ where
     for t in s.children do
       go t.get
 
-/--
-  Returns the value of a snapshot task if it has already finished and otherwise cancels the task.
-  Tasks will be canceled implicitly when the last reference to them is dropped but being explicit
-  here doesn't hurt. -/
-def getOrCancel? (t? : Option (SnapshotTask α)) : BaseIO (Option α) := do
-  let some t := t? | return none
-  if (← IO.hasFinished t.task) then
-    return t.get
-  IO.cancel t.task
-  return none
-
 /-- Metadata that does not change during the lifetime of the language processing process. -/
 structure ProcessingContext where
   /-- Module name of the file being processed. -/
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 13322ab68d4d..c7644af4448a 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -131,6 +131,16 @@ partial instance : ToSnapshotTree CommandParsedSnapshot where
       #[s.data.sig.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.sig.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 HeaderProcessedSucessfully where
   /-- The resulting initial elaboration state. -/
@@ -215,15 +225,17 @@ where
     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 success := old.success? then
-        return { old with ictx, success? := some { success with
-          processed := (← success.processed.bindIO (sync := true) fun processed => do
-            if let some procSuccess := processed.success? then
-              let oldCmd? ← getOrCancel? procSuccess.next
-              return .pure { processed with success? := some { procSuccess with
-                next := (← parseCmd oldCmd? success.parserState procSuccess.cmdState) } }
+      if let some oldSuccess := old.success? then
+        return { old with ictx, success? := some { oldSuccess with
+          processed := (← oldSuccess.processed.bindIO (sync := true) fun oldProcessed => do
+            if let some oldProcSuccess := oldProcessed.success? then
+              -- also wait on old command parse snapshot as parsing is cheap and may allow for
+              -- elaboration reuse
+              oldProcSuccess.next.bindIO (sync := true) fun oldCmd =>
+                return .pure { oldProcessed with success? := some { oldProcSuccess with
+                  next := (← parseCmd oldCmd oldSuccess.parserState oldProcSuccess.cmdState) } }
             else
-              return .pure processed) } }
+              return .pure oldProcessed) } }
       else return old
 
     -- fast path: if we have parsed the header successfully...
@@ -247,6 +259,12 @@ where
       if let some old := old? then
         if firstDiffPos?.any (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.success? then
+          oldSuccess.processed.cancel
+          let _ ← BaseIO.mapTask (t := oldSuccess.processed.task) fun processed => do
+            if let some oldProcSuccess := processed.success? then
+              let _ ← BaseIO.mapTask (·.cancel) oldProcSuccess.next.task
 
       return {
         ictx, stx, msgLog
@@ -326,13 +344,26 @@ where
   parseCmd (old? : Option CommandParsedSnapshot) (parserState : Parser.ModuleParserState)
       (cmdState : Command.State) :
       BaseIO (SnapshotTask CommandParsedSnapshot) := do
+    -- 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 (next? := none) {
+        msgLog := .empty, stx := .missing, parserState
+        sig := .pure { msgLog := .empty, finished := .pure { msgLog := .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)
-          (next? := (← old.data.sig.bindIO (sync := true) fun sig =>
-            sig.finished.bindIO (sync := true) fun finished => do
-              parseCmd (← getOrCancel? oldNext) old.data.parserState finished.cmdState))
+          (next? := (← old.data.sig.bindIO (sync := true) fun oldSig =>
+            oldSig.finished.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))
       else return old  -- terminal command, we're done!
 
     -- fast path, do not even start new task for this snapshot
@@ -354,14 +385,14 @@ where
       if let some old := old? then
         if firstDiffPos?.any (parserState.pos < ·) && old.data.stx == stx then
           return (← unchanged old)
+        -- on first change, make sure to cancel all further old tasks
+        old.cancel
 
-      let _ ← old?.map (·.data.sig) |> getOrCancel?
       let sig ← processCmdSignature stx cmdState msgLog.hasErrors beginPos
       let next? ← if Parser.isTerminalCommand stx then pure none
         -- for now, wait on "command finished" snapshot before parsing next command
-        else some <$> sig.bindIO fun sig => do
-          sig.finished.bindIO fun finished =>
-            parseCmd none parserState finished.cmdState
+        else some <$> (sig.bind (·.finished)).bindIO fun finished =>
+          parseCmd none parserState finished.cmdState
       return .mk (next? := next?) {
         msgLog
         stx
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index c57898fa6df5..a7689d418483 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -42,7 +42,7 @@ def set (tk : CancelToken) : IO Unit :=
 
 end CancelToken
 
--- TEMP: translate from new heterogeneous snapshot tree to old homegenous async list
+-- TEMP: translate from new heterogeneous snapshot tree to old homogeneous async list
 private partial def mkCmdSnaps (initSnap : Language.Lean.InitialSnapshot) :
     AsyncList ElabTaskError Snapshot := Id.run do
   let some headerParsed := initSnap.success? | return .nil

From 91e7e28f6e882d069b1b35ca56c41769daef92f1 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Wed, 13 Dec 2023 11:40:56 +0100
Subject: [PATCH 30/60] Adjust comments

---
 src/Lean/Elab/Frontend.lean | 1 -
 src/Lean/Language/Lean.lean | 8 +++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index c9f8b1bdf77f..89514f78c248 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -92,7 +92,6 @@ def runFrontend
     (opts : Options)
     (fileName : String)
     (mainModuleName : Name)
-    -- TODO: do we still want this in the driver?
     (trustLevel : UInt32 := 0)
     (ileanFileName? : Option String := none)
     : IO (Environment × Bool) := do
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index c7644af4448a..30e1ff19f096 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -206,9 +206,8 @@ General notes:
   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 (try to remember to) cancel tasks that we cannot reuse proactively even though they are
-  auto-canceled when the last reference to them is dropped; but it is not always easy to ensure that
-  the last drop happens precisely where we expect it.
+* 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.
@@ -256,6 +255,9 @@ where
       -- 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 firstDiffPos?.any (parserState.pos < ·) && old.stx == stx then
           return (← unchanged old)

From 33a9647b260da31fd0763c44bef31b9ca281d392 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Wed, 13 Dec 2023 16:39:06 +0100
Subject: [PATCH 31/60] Cache interactive diagnostics properly

---
 src/Lean/Elab/Frontend.lean                   |  10 +-
 src/Lean/Language/Basic.lean                  |  53 +++++--
 src/Lean/Language/Lean.lean                   | 136 ++++++++++++------
 src/Lean/Server/FileWorker.lean               |   6 +-
 src/Lean/Server/FileWorker/Utils.lean         |   7 +-
 .../Server/FileWorker/WidgetRequests.lean     |   2 +-
 src/Lean/Server/Snapshots.lean                |   2 +-
 src/Lean/Widget/InteractiveDiagnostic.lean    |  10 +-
 8 files changed, 157 insertions(+), 69 deletions(-)

diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index 89514f78c248..e4dd403895d2 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -98,8 +98,12 @@ def runFrontend
   let inputCtx := Parser.mkInputContext input fileName
   -- TODO: replace with `#lang` processing
   let lang := Language.Lean
-  let ctx := { mainModuleName, opts, trustLevel, fileSetupHandler? := none }
-  let snap ← lang.process ctx none inputCtx
+  let ctx := { inputCtx with
+    mainModuleName, opts, trustLevel
+    fileSetupHandler? := none
+    clientHasWidgets := false
+  }
+  let snap ← lang.process none ctx
   let snaps := Language.toSnapshotTree snap
   snaps.runAndReport opts
   if let some ileanFileName := ileanFileName? then
@@ -108,7 +112,7 @@ def runFrontend
     let ilean := { module := mainModuleName, references : Lean.Server.Ilean }
     IO.FS.writeFile ileanFileName $ Json.compress $ toJson ilean
 
-  let hasErrors := snaps.getAll.any (·.msgLog.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 := lang.getFinalEnv? snap |>.getD (← mkEmptyEnvironment)
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index a2445aa2136e..53fe44d325df 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -10,6 +10,7 @@ Authors: Sebastian Ullrich
 
 import Lean.Message
 import Lean.Parser.Types
+import Lean.Widget.InteractiveDiagnostic
 
 set_option linter.missingDocs true
 
@@ -41,6 +42,21 @@ end Lean.Server.FileWorker
 
 namespace Lean.Language
 
+/-- `MessageLog` with caching of interactive diagnostics. -/
+structure Snapshot.Diagnostics where
+  /-- Non-interactive message log. -/
+  msgLog : MessageLog
+  /-- 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 : Array Widget.InteractiveDiagnostic
+deriving Inhabited
+
+/-- The empty set of diagnostics. -/
+def Snapshot.Diagnostics.empty : Snapshot.Diagnostics where
+  msgLog := .empty
+  interactiveDiags := #[]
+
 /--
   The base class of all snapshots: all the generic information the language server needs about a
   snapshot. -/
@@ -48,7 +64,7 @@ structure Snapshot where
   /--
     The messages produced by this step. The union of message logs of all finished snapshots is
     reported to the user. -/
-  msgLog : MessageLog
+  diagnostics : Snapshot.Diagnostics
   /-- General elaboration metadata produced by this step. -/
   infoTree? : Option Elab.InfoTree := none
   -- (`InfoTree` is quite Lean-specific at this point, but we want to make it more generic)
@@ -149,7 +165,8 @@ register_builtin_option printMessageEndPos : Bool := {
   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.msgLog.forM (·.toString (includeEndPos := printMessageEndPos.get opts) >>= IO.print)
+  s.element.diagnostics.msgLog.forM
+    (·.toString (includeEndPos := printMessageEndPos.get opts) >>= IO.print)
   for t in s.children do
     t.get.runAndReport opts
 
@@ -163,18 +180,39 @@ where
       go t.get
 
 /-- Metadata that does not change during the lifetime of the language processing process. -/
-structure ProcessingContext where
+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
+  /-- Whether to create interactive diagnostics. -/
+  clientHasWidgets : Bool
   /--
     Callback available in server mode for building imports and retrieving per-library options using
     `lake setup-file`. -/
   fileSetupHandler? : Option (Array Import → IO Server.FileWorker.FileSetupResult)
 
+/-- Context of an input processing invocation. -/
+structure ProcessingContext extends ModuleProcessingContext, Parser.InputContext
+
+/-- Monad holding all relevant data for processing. -/
+abbrev ProcessingM := ReaderT ProcessingContext BaseIO
+
+/--
+Creates snapshot message log from non-interactive message log, caching derived interactive
+diagnostics.
+-/
+def Snapshot.Diagnostics.ofMessageLog (msgLog : Lean.MessageLog) :
+    ProcessingM Snapshot.Diagnostics :=
+  return {
+    msgLog
+    interactiveDiags := (← msgLog.toList.toArray.mapM fun msg => do
+      let ctx ← read
+      Widget.msgToInteractiveDiagnostic ctx.fileMap msg ctx.clientHasWidgets)
+  }
+
 end Language
 open Language
 
@@ -190,8 +228,7 @@ structure Language where
     Processes input into snapshots, potentially reusing information from a previous run.
     Constructing the initial snapshot is assumed to be cheap enough that it can be done
     synchronously, which simplifies use of this function. -/
-  process : ProcessingContext → (old? : Option InitialSnapshot) → Parser.InputContext →
-    BaseIO InitialSnapshot
+  process : (old? : Option InitialSnapshot) → ProcessingM InitialSnapshot
   -- TODO: is this the right interface for other languages as well?
   /-- Gets final environment, if any, that is to be used for persisting, code generation, etc. -/
   getFinalEnv? : InitialSnapshot → Option Environment
@@ -201,10 +238,10 @@ instance (lang : Language) : ToSnapshotTree lang.InitialSnapshot := lang.instToS
 /--
   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 (lang : Language) (ctx: ProcessingContext) :
+partial def Language.mkIncrementalProcessor (lang : Language) (ctx : ModuleProcessingContext) :
     BaseIO (Parser.InputContext → BaseIO lang.InitialSnapshot) := do
   let oldRef ← IO.mkRef none
-  return fun doc => do
-    let snap ← lang.process ctx (← oldRef.get) doc
+  return fun ictx => do
+    let snap ← lang.process (← oldRef.get) { ctx, ictx with }
     oldRef.set (some snap)
     return snap
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 30e1ff19f096..05fe9f20a611 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -185,15 +185,40 @@ def HeaderParsedSnapshot.processedSuccessfully (snap : HeaderParsedSnapshot) :
 /-- Initial snapshot of the Lean language processor: a "header parsed" snapshot. -/
 abbrev InitialSnapshot := HeaderParsedSnapshot
 
-private def msglogOfHeaderError (data : MessageData) : MessageLog :=
-  MessageLog.empty.add { fileName := "<input>", pos := ⟨0, 0⟩, data }
+/-- 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 ProcessingM LeanProcessingM 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 < ·)
+
+private def diagnosticsOfHeaderError (msg : String) : ProcessingM Snapshot.Diagnostics :=
+  let msgLog := MessageLog.empty.add { fileName := "<input>", pos := ⟨0, 0⟩, 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. -/
-private def withHeaderExceptions (ex : Snapshot → α) (act : IO α) : BaseIO α := do
-  match (← act.toBaseIO) with
-  | .error e => return ex { msgLog := msglogOfHeaderError e.toString }
+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
 
 /-- Makes sure we load imports at most once per process as they cannot be unloaded. -/
@@ -212,15 +237,15 @@ General notes:
   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 processLean
-    (ctx : ProcessingContext) (old? : Option InitialSnapshot) (ictx : Parser.InputContext) :
-    BaseIO InitialSnapshot :=
-  parseHeader old?
-where
+partial def processLean (old? : Option InitialSnapshot) : ProcessingM InitialSnapshot := do
   -- compute position of syntactic change once
-  firstDiffPos? := old?.map (·.ictx.input.firstDiffPos ictx.input)
-
-  parseHeader (old? : Option HeaderParsedSnapshot) := do
+  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
@@ -232,7 +257,7 @@ where
               -- elaboration reuse
               oldProcSuccess.next.bindIO (sync := true) fun oldCmd =>
                 return .pure { oldProcessed with success? := some { oldProcSuccess with
-                  next := (← parseCmd oldCmd oldSuccess.parserState oldProcSuccess.cmdState) } }
+                  next := (← parseCmd oldCmd oldSuccess.parserState oldProcSuccess.cmdState ctx) } }
             else
               return .pure oldProcessed) } }
       else return old
@@ -242,7 +267,7 @@ where
       if let some (some processed) ← old.processedSuccessfully.get? then
         -- ...and the edit location is after the next command (see note [Incremental Parsing])...
         if let some nextCom ← processed.next.get? then
-          if firstDiffPos?.any (nextCom.data.parserState.pos < ·) then
+          if (← isBeforeEditPos nextCom.data.parserState.pos) then
             -- ...go immediately to next snapshot
             return (← unchanged old)
 
@@ -250,7 +275,11 @@ where
       -- parsing the header should be cheap enough to do synchronously
       let (stx, parserState, msgLog) ← Parser.parseHeader ictx
       if msgLog.hasErrors then
-        return { ictx, stx, msgLog, success? := none }
+        return {
+          ictx, stx
+          diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
+          success? := none
+        }
 
       -- semi-fast path: go to next snapshot if syntax tree is unchanged AND we're still in front
       -- of the edit location
@@ -259,7 +288,7 @@ where
       -- 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 firstDiffPos?.any (parserState.pos < ·) && old.stx == stx 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.success? then
@@ -269,16 +298,20 @@ where
               let _ ← BaseIO.mapTask (·.cancel) oldProcSuccess.next.task
 
       return {
-        ictx, stx, msgLog
+        ictx, stx
+        diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
         success? := some {
           parserState
           processed := (← processHeader stx parserState)
         }
       }
 
-  processHeader (stx : Syntax) (parserState : Parser.ModuleParserState) :=
-    SnapshotTask.ofIO ⟨0, ictx.input.endPos⟩ <|
-    withHeaderExceptions ({ · with success? := none }) do
+  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 success? := none }) do
       -- use `lake setup-file` if in language server
       let fileSetupResult ← if let some handler := ctx.fileSetupHandler? then
         handler (Elab.headerToImports stx)
@@ -287,13 +320,14 @@ where
       match fileSetupResult.kind with
       | .importsOutOfDate =>
         return {
-          msgLog := msglogOfHeaderError
-            "Imports are out of date and must be rebuilt; use the \"Restart File\" command in your editor."
+          diagnostics := (← diagnosticsOfHeaderError
+            "Imports are out of date and must be rebuilt; \
+             use the \"Restart File\" command in your editor.")
           success? := none
         }
       | .error msg =>
         return {
-          msgLog := msglogOfHeaderError msg
+          diagnostics := (← diagnosticsOfHeaderError msg)
           success? := none
         }
       | _ => pure ()
@@ -306,14 +340,14 @@ where
         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 { msgLog := .empty, success? := none }
+        return { diagnostics := .empty, success? := none }
 
       -- override context options with file options
       let opts := ctx.opts.mergeBy (fun _ _ fileOpt => fileOpt) fileSetupResult.fileOptions
 
       -- 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 ictx
-        ctx.trustLevel
+      let (headerEnv, msgLog) ← Elab.processHeader (leakEnv := true) stx opts .empty
+        ctx.toInputContext ctx.trustLevel
 
       let headerEnv := headerEnv.setMainModule ctx.mainModuleName
       let cmdState := Elab.Command.mkState headerEnv msgLog opts
@@ -321,7 +355,7 @@ where
         enabled := true
         trees := #[Elab.InfoTree.context ({
           env     := headerEnv
-          fileMap := ictx.fileMap
+          fileMap := ctx.fileMap
           ngen    := { namePrefix := `_import }
         }) (Elab.InfoTree.node
             (Elab.Info.ofCommandInfo { elaborator := `header, stx })
@@ -334,7 +368,7 @@ where
         )].toPArray'
       }}
       return {
-        msgLog := cmdState.messages
+        diagnostics := (← Snapshot.Diagnostics.ofMessageLog cmdState.messages)
         infoTree? := cmdState.infoState.trees[0]!
         success? := some {
           cmdState
@@ -344,8 +378,9 @@ where
       }
 
   parseCmd (old? : Option CommandParsedSnapshot) (parserState : Parser.ModuleParserState)
-      (cmdState : Command.State) :
-      BaseIO (SnapshotTask CommandParsedSnapshot) := do
+      (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
@@ -353,8 +388,10 @@ where
       -- (as no-one should look at this result in that case) but anything containing `Environment`
       -- is not `Inhabited`
       return .pure <| .mk (next? := none) {
-        msgLog := .empty, stx := .missing, parserState
-        sig := .pure { msgLog := .empty, finished := .pure { msgLog := .empty, cmdState } } }
+        diagnostics := .empty, stx := .missing, parserState
+        sig := .pure {
+          diagnostics := .empty
+          finished := .pure { diagnostics := .empty, cmdState } } }
 
     let unchanged old : BaseIO CommandParsedSnapshot :=
       -- when syntax is unchanged, reuse command processing task as is
@@ -365,51 +402,55 @@ where
               -- 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))
+                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 firstDiffPos?.any (nextCom.data.parserState.pos < ·) then
+        if (← isBeforeEditPos nextCom.data.parserState.pos) then
           return .pure (← unchanged old)
 
-    SnapshotTask.ofIO ⟨parserState.pos, ictx.input.endPos⟩ do
+    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 ictx pmctx parserState .empty
+      let (stx, parserState, msgLog) := Parser.parseCommand ctx.toInputContext pmctx parserState
+        .empty
 
       -- semi-fast path
       if let some old := old? then
-        if firstDiffPos?.any (parserState.pos < ·) && old.data.stx == stx 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 sig ← processCmdSignature stx cmdState msgLog.hasErrors beginPos
+      let sig ← 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 <$> (sig.bind (·.finished)).bindIO fun finished =>
-          parseCmd none parserState finished.cmdState
+          parseCmd none parserState finished.cmdState ctx
       return .mk (next? := next?) {
-        msgLog
+        diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog ctx.toProcessingContext)
         stx
         parserState
         sig
       }
 
   processCmdSignature (stx : Syntax) (cmdState : Command.State) (hasParseError : Bool)
-      (beginPos : String.Pos) :=
+      (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 := { ictx with cmdPos := beginPos, tacticCache? := none }
+      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
@@ -429,16 +470,17 @@ where
             tag == `Tactic.unsolvedGoals || (`_traceMsg).isSuffixOf tag)⟩
       if !output.isEmpty then
         messages := messages.add {
-          fileName := ictx.fileName
+          fileName := ctx.fileName
           severity := MessageSeverity.information
-          pos      := ictx.fileMap.toPosition beginPos
+          pos      := ctx.fileMap.toPosition beginPos
           data     := output
         }
       let cmdState := { cmdState with messages }
       return {
-        msgLog := .empty
+        diagnostics := .empty
         finished := .pure {
-          msgLog := cmdState.messages
+          diagnostics :=
+            (← Snapshot.Diagnostics.ofMessageLog cmdState.messages ctx.toProcessingContext)
           infoTree? := some cmdState.infoState.trees[0]!
           cmdState
         }
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index b62dc08bbd37..2f4ee3889431 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -111,8 +111,9 @@ section Elab
     go node st cont := do
       if (← IO.checkCanceled) then
         return .pure <| .ok ()
-      let diagnostics := st.diagnostics ++ (← node.element.msgLog.toList.toArray.mapM (Widget.msgToInteractiveDiagnostic m.text · ctx.clientHasWidgets)).map (·.toDiagnostic)
-      if st.hasBlocked && !node.element.msgLog.isEmpty then
+      let diagnostics := st.diagnostics ++
+        node.element.diagnostics.interactiveDiags.map (·.toDiagnostic)
+      if st.hasBlocked && !node.element.diagnostics.msgLog.isEmpty then
         ctx.chanOut.send <| mkPublishDiagnosticsNotification m diagnostics
       let infoTrees := match node.element.infoTree? with
         | some itree => st.infoTrees.push itree
@@ -163,6 +164,7 @@ section Initialization
     let chanOut ← mkLspOutputChannel
     let processor ← Language.Lean.mkIncrementalProcessor {
       opts, mainModuleName
+      clientHasWidgets
       fileSetupHandler? := some fun imports =>
         setupFile meta imports fun stderrLine =>
           let progressDiagnostic := {
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index a7689d418483..00317dbb396a 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -47,12 +47,14 @@ private partial def mkCmdSnaps (initSnap : Language.Lean.InitialSnapshot) :
     AsyncList ElabTaskError Snapshot := Id.run do
   let some headerParsed := initSnap.success? | return .nil
   .delayed <| headerParsed.processed.task.bind fun headerProcessed => Id.run do
+    -- NOTE: this throws away interactive diagnostics of header errors but these are not interactive
+    -- anyway
     let some headerSuccess := headerProcessed.success? | return .pure <| .ok .nil
     return .pure <| .ok <| .cons {
       stx := initSnap.stx
       mpState := headerParsed.parserState
       cmdState := headerSuccess.cmdState
-      interactiveDiags := .empty  -- TODO
+      interactiveDiags := headerProcessed.diagnostics.interactiveDiags
     } <| .delayed <| headerSuccess.next.task.bind go
 where go cmdParsed :=
   cmdParsed.data.sig.task.bind fun sig =>
@@ -61,7 +63,8 @@ where go cmdParsed :=
         stx := cmdParsed.data.stx
         mpState := cmdParsed.data.parserState
         cmdState := finished.cmdState
-        interactiveDiags := .empty  -- TODO
+        interactiveDiags :=
+          cmdParsed.data.diagnostics.interactiveDiags ++ sig.diagnostics.interactiveDiags
       } (match cmdParsed.next? with
         | some next => .delayed <| next.task.bind go
         | none => .nil)
diff --git a/src/Lean/Server/FileWorker/WidgetRequests.lean b/src/Lean/Server/FileWorker/WidgetRequests.lean
index 2c7fa5788469..d7eb9376bab0 100644
--- a/src/Lean/Server/FileWorker/WidgetRequests.lean
+++ b/src/Lean/Server/FileWorker/WidgetRequests.lean
@@ -101,7 +101,7 @@ def getInteractiveDiagnostics (params : GetInteractiveDiagnosticsParams) : Reque
   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 =>
+      snap.interactiveDiags.filter fun diag =>
         params.lineRange?.all fun ⟨s, e⟩ =>
           -- does [s,e) intersect [diag.fullRange.start.line,diag.fullRange.end.line)?
           s ≤ diag.fullRange.start.line ∧ diag.fullRange.start.line < e ∨
diff --git a/src/Lean/Server/Snapshots.lean b/src/Lean/Server/Snapshots.lean
index 1ed26c40c409..d33fda462cc0 100644
--- a/src/Lean/Server/Snapshots.lean
+++ b/src/Lean/Server/Snapshots.lean
@@ -32,7 +32,7 @@ structure Snapshot where
   /-- 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
+  interactiveDiags : Array Widget.InteractiveDiagnostic
 
 namespace Snapshot
 
diff --git a/src/Lean/Widget/InteractiveDiagnostic.lean b/src/Lean/Widget/InteractiveDiagnostic.lean
index b2cead4ec431..0abdc9660175 100644
--- a/src/Lean/Widget/InteractiveDiagnostic.lean
+++ b/src/Lean/Widget/InteractiveDiagnostic.lean
@@ -167,7 +167,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
@@ -189,10 +190,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

From 63a853e70544c76505849937dafe92c85763dd43 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Wed, 13 Dec 2023 19:03:08 +0100
Subject: [PATCH 32/60] fix: clear progress notifications from `lake` in the
 end

---
 src/Lean/Language/Basic.lean    | 2 +-
 src/Lean/Server/FileWorker.lean | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index 53fe44d325df..d64aa7075338 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -228,7 +228,7 @@ structure Language where
     Processes input into snapshots, potentially reusing information from a previous run.
     Constructing the initial snapshot is assumed to be cheap enough that it can be done
     synchronously, which simplifies use of this function. -/
-  process : (old? : Option InitialSnapshot) → ProcessingM InitialSnapshot
+  process (old? : Option InitialSnapshot) : ProcessingM InitialSnapshot
   -- TODO: is this the right interface for other languages as well?
   /-- Gets final environment, if any, that is to be used for persisting, code generation, etc. -/
   getFinalEnv? : InitialSnapshot → Option Environment
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 2f4ee3889431..ff8773d5db33 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -165,14 +165,17 @@ section Initialization
     let processor ← Language.Lean.mkIncrementalProcessor {
       opts, mainModuleName
       clientHasWidgets
-      fileSetupHandler? := some fun imports =>
-        setupFile meta imports fun stderrLine =>
+      fileSetupHandler? := some fun imports => do
+        let result ← setupFile meta imports fun stderrLine => do
           let progressDiagnostic := {
             range     := ⟨⟨0, 0⟩, ⟨0, 0⟩⟩
             severity? := DiagnosticSeverity.information
             message   := stderrLine
           }
           chanOut.send <| mkPublishDiagnosticsNotification meta #[progressDiagnostic]
+        -- clear progress notifications in the end
+        chanOut.send <| mkPublishDiagnosticsNotification meta #[]
+        return result
     }
     let initSnap ← processor meta.mkInputContext
     let ctx := {

From 9c49e4970fd8e6bc654625fd6c664cd11bbb8ef1 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 14 Dec 2023 12:42:54 +0100
Subject: [PATCH 33/60] fix: header errors should be fatal

---
 src/Lean/Language/Lean.lean | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 05fe9f20a611..8ada8c31b154 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -348,6 +348,9 @@ where
       -- 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, success? := none }
 
       let headerEnv := headerEnv.setMainModule ctx.mainModuleName
       let cmdState := Elab.Command.mkState headerEnv msgLog opts
@@ -368,7 +371,7 @@ where
         )].toPArray'
       }}
       return {
-        diagnostics := (← Snapshot.Diagnostics.ofMessageLog cmdState.messages)
+        diagnostics
         infoTree? := cmdState.infoState.trees[0]!
         success? := some {
           cmdState

From 67c0599df16cf48c6cc36e1a9b36ec1b736848eb Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 14 Dec 2023 13:47:30 +0100
Subject: [PATCH 34/60] Header diagnostics should cover the full file

---
 src/Lean/Language/Lean.lean     | 9 +++++++--
 src/Lean/Server/FileWorker.lean | 7 ++++---
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 8ada8c31b154..cad50dbc0f8e 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -208,8 +208,13 @@ compared to it.
 def isBeforeEditPos (pos : String.Pos) : LeanProcessingM Bool := do
   return (← read).firstDiffPos?.any (pos < ·)
 
-private def diagnosticsOfHeaderError (msg : String) : ProcessingM Snapshot.Diagnostics :=
-  let msgLog := MessageLog.empty.add { fileName := "<input>", pos := ⟨0, 0⟩, data := msg }
+private def diagnosticsOfHeaderError (msg : String) : ProcessingM Snapshot.Diagnostics := do
+  let msgLog := MessageLog.empty.add {
+    fileName := "<input>"
+    pos := ⟨0, 0⟩
+    endPos := dbgTraceVal <| (← read).fileMap.toPosition (← read).fileMap.source.endPos
+    data := msg
+  }
   Snapshot.Diagnostics.ofMessageLog msgLog
 
 /--
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index ff8773d5db33..ad400bd14297 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -168,9 +168,10 @@ section Initialization
       fileSetupHandler? := some fun imports => do
         let result ← setupFile meta imports fun stderrLine => do
           let progressDiagnostic := {
-            range     := ⟨⟨0, 0⟩, ⟨0, 0⟩⟩
-            severity? := DiagnosticSeverity.information
-            message   := stderrLine
+            range      := ⟨⟨0, 0⟩, ⟨0, 0⟩⟩
+            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

From e0e7faf86c750aba9aec07a00e941565db4ef68d Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 14 Dec 2023 14:04:23 +0100
Subject: [PATCH 35/60] server reporting delay

---
 src/Lean/Server/FileWorker.lean | 52 ++++++++++++++++++++++++---------
 1 file changed, 38 insertions(+), 14 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index ad400bd14297..8633c3fbcb39 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -60,10 +60,16 @@ structure WorkerContext where
   /-- 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.
+  -/
+  maxDocVersion    : IO.Ref Nat
   hLog             : FS.Stream
   initParams       : InitializeParams
   processor        : Parser.InputContext → BaseIO Lean.Language.Lean.InitialSnapshot
   clientHasWidgets : Bool
+  opts             : Options
 
 /-! # Asynchronous snapshot elaboration -/
 
@@ -88,29 +94,43 @@ section Elab
     diagnostics : Array Lsp.Diagnostic := #[]
     infoTrees : Array Elab.InfoTree := #[]
 
+  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"
+  }
+
   /--
     Reports status of a snapshot tree incrementally to the user: progress,
     diagnostics, .ilean reference information.
 
     Debouncing: we only report information
-    * when first blocking, i.e. not before skipping over any unchanged snapshots
+    * 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) (m : DocumentMeta) (snaps : Language.SnapshotTree)
-      : IO (Task Unit) := do
-    Task.map (fun _ => ()) <$> go snaps { : ReportSnapshotsState } fun st => do
+  private partial def reportSnapshots (ctx : WorkerContext) (m : DocumentMeta)
+      (snaps : Language.SnapshotTree) : BaseIO (Task Unit) := do
+    let t ← BaseIO.asTask do
+      IO.sleep (server.reportDelayMs.get ctx.opts).toUInt32
+    BaseIO.bindTask t fun _ =>
+      start
+  where
+    start := go snaps { : ReportSnapshotsState } fun st => do
+      -- callback at the end of reporting
       ctx.chanOut.send <| mkFileProgressDoneNotification m
       unless st.hasBlocked do
         ctx.chanOut.send <| mkPublishDiagnosticsNotification m st.diagnostics
       -- This will overwrite existing ilean info for the file, in case something
       -- went wrong during the incremental updates.
       ctx.chanOut.send <| mkIleanInfoFinalNotification m st.infoTrees
-      return .pure <| .ok ()
-  where
+      return .pure ()
     go node st cont := do
       if (← IO.checkCanceled) then
-        return .pure <| .ok ()
+        return .pure ()
       let diagnostics := st.diagnostics ++
         node.element.diagnostics.interactiveDiags.map (·.toDiagnostic)
       if st.hasBlocked && !node.element.diagnostics.msgLog.isEmpty then
@@ -130,7 +150,7 @@ section Elab
           if !st.hasBlocked then
             ctx.chanOut.send <| mkPublishDiagnosticsNotification m st.diagnostics
             st := { st with hasBlocked := true }
-        IO.bindTask t.task fun node =>
+        BaseIO.bindTask t.task fun node =>
           go node st (goSeq · cont ts)
 end Elab
 
@@ -161,7 +181,8 @@ section Initialization
       if let some path := System.Uri.fileUriToPath? meta.uri then
         mainModuleName ← moduleNameOfFileName path none
     catch _ => pure ()
-    let chanOut ← mkLspOutputChannel
+    let maxDocVersion ← IO.mkRef 0
+    let chanOut ← mkLspOutputChannel maxDocVersion
     let processor ← Language.Lean.mkIncrementalProcessor {
       opts, mainModuleName
       clientHasWidgets
@@ -185,6 +206,8 @@ section Initialization
       initParams
       processor
       clientHasWidgets
+      maxDocVersion
+      opts
     }
     let reporter ← reportSnapshots ctx meta (Language.ToSnapshotTree.toSnapshotTree initSnap)
     let doc : EditableDocument := { meta, initSnap, reporter }
@@ -199,10 +222,9 @@ section Initialization
         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 : IO (IO.Channel JsonRpc.Message) := do
+    mkLspOutputChannel maxDocVersion : IO (IO.Channel JsonRpc.Message) := do
       let chanOut ← IO.Channel.new
       -- most recent document version seen in notifications
-      let maxVersion ← IO.mkRef 0
       let _ ← chanOut.forAsync (prio := .dedicated) fun msg => do
         -- discard outdated notifications; note that in contrast to responses, notifications can
         -- always be silently discarded
@@ -214,10 +236,10 @@ section Initialization
             | _ => none
           doc.find compare "version" |>.bind (·.getNat?.toOption)
         if let some version := version? then
-          if version < (← maxVersion.get) then
+          if version < (← maxDocVersion.get) then
             return
-          else
-            maxVersion.set version
+          -- 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
@@ -232,6 +254,8 @@ section Updates
     let initSnap ← ctx.processor meta.mkInputContext
     let reporter ← reportSnapshots ctx meta (Language.ToSnapshotTree.toSnapshotTree initSnap)
     modify fun st => { st with doc := { meta, initSnap, reporter } }
+    -- we assume versions are monotonous
+    ctx.maxDocVersion.set meta.version
 end Updates
 
 /- Notifications are handled in the main thread. They may change global worker state

From cef0d59404626a3ea3481813c8e17f6708c4512c Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 14 Dec 2023 14:09:20 +0100
Subject: [PATCH 36/60] fix: header errors should be fatal

---
 src/Lean/Language/Basic.lean    |  6 +++++-
 src/Lean/Language/Lean.lean     |  3 +++
 src/Lean/Server/FileWorker.lean | 10 ++++++++--
 3 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index d64aa7075338..00e1619fd577 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -67,7 +67,11 @@ structure Snapshot where
   diagnostics : Snapshot.Diagnostics
   /-- General elaboration metadata produced by this step. -/
   infoTree? : Option Elab.InfoTree := none
-  -- (`InfoTree` is quite Lean-specific at this point, but we want to make it more generic)
+  /--
+  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`). -/
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index cad50dbc0f8e..a241dd3db76f 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -154,6 +154,7 @@ structure HeaderProcessedSucessfully where
 structure HeaderProcessedSnapshot extends Snapshot where
   /-- State after successful importing. -/
   success? : Option HeaderProcessedSucessfully
+  isFatal := success?.isNone
 instance : ToSnapshotTree HeaderProcessedSnapshot where
   toSnapshotTree s := ⟨s.toSnapshot, #[] |>
     pushOpt (s.success?.map (·.next.map (sync := true) toSnapshotTree))⟩
@@ -173,6 +174,8 @@ structure HeaderParsedSnapshot extends Snapshot where
   stx : Syntax
   /-- State after successful parsing. -/
   success? : Option HeaderParsedSucessfully
+  isFatal := success?.isNone
+
 instance : ToSnapshotTree HeaderParsedSnapshot where
   toSnapshotTree s := ⟨s.toSnapshot,
     #[] |> pushOpt (s.success?.map (·.processed.map (sync := true) toSnapshotTree))⟩
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 8633c3fbcb39..1b488b0ba69d 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -93,6 +93,7 @@ section Elab
     hasBlocked := false
     diagnostics : Array Lsp.Diagnostic := #[]
     infoTrees : Array Elab.InfoTree := #[]
+    isFatal := false
 
   register_builtin_option server.reportDelayMs : Nat := {
     defValue := 200
@@ -121,7 +122,10 @@ section Elab
   where
     start := go snaps { : ReportSnapshotsState } fun st => do
       -- callback at the end of reporting
-      ctx.chanOut.send <| mkFileProgressDoneNotification m
+      if st.isFatal then
+        ctx.chanOut.send <| mkFileProgressAtPosNotification m 0 .fatalError
+      else
+        ctx.chanOut.send <| mkFileProgressDoneNotification m
       unless st.hasBlocked do
         ctx.chanOut.send <| mkPublishDiagnosticsNotification m st.diagnostics
       -- This will overwrite existing ilean info for the file, in case something
@@ -140,7 +144,9 @@ section Elab
         | none       => st.infoTrees
       if st.hasBlocked && node.element.infoTree?.isSome then
         ctx.chanOut.send <| mkIleanInfoUpdateNotification m infoTrees
-      goSeq { st with diagnostics, infoTrees } cont node.children.toList
+      -- we assume that only the last node in the tree sets `isFatal`
+      let st := { st with diagnostics, infoTrees, isFatal := node.element.isFatal }
+      goSeq st cont node.children.toList
     goSeq st cont
       | [] => cont st
       | t::ts => do

From dc97d753c1ca3471d6ad8c90ce06812d9648e9dd Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 14 Dec 2023 15:46:31 +0100
Subject: [PATCH 37/60] fix nondeterministic test output from new debouncing

---
 src/Lean/Data/Lsp/Ipc.lean                    | 20 +++++++++------
 .../editAfterError.lean.expected.out          | 12 ---------
 tests/lean/interactive/run.lean               |  5 ++--
 tests/lean/server/diags.lean                  | 25 ++++++++-----------
 tests/lean/server/edits.lean                  | 25 ++++++++-----------
 tests/lean/server/init_exit_worker.lean       |  4 +--
 6 files changed, 39 insertions(+), 52 deletions(-)

diff --git a/src/Lean/Data/Lsp/Ipc.lean b/src/Lean/Data/Lsp/Ipc.lean
index c6580ede7137..b398e71ad109 100644
--- a/src/Lean/Data/Lsp/Ipc.lean
+++ b/src/Lean/Data/Lsp/Ipc.lean
@@ -83,15 +83,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.
+ -/
 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
@@ -99,10 +106,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 {
diff --git a/tests/lean/interactive/editAfterError.lean.expected.out b/tests/lean/interactive/editAfterError.lean.expected.out
index 76f7246bd0ea..7206de2b27de 100644
--- a/tests/lean/interactive/editAfterError.lean.expected.out
+++ b/tests/lean/interactive/editAfterError.lean.expected.out
@@ -4,18 +4,6 @@
    "range":
    {"start": {"line": 1, "character": 10},
     "end": {"line": 1, "character": 11}}}]}
-{"version": 2, "uri": "file:///editAfterError.lean", "diagnostics": []}
-{"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}}}]}
 {"version": 2,
  "uri": "file:///editAfterError.lean",
  "diagnostics":
diff --git a/tests/lean/interactive/run.lean b/tests/lean/interactive/run.lean
index 5b56d44f203d..4b5fba9bfb20 100644
--- a/tests/lean/interactive/run.lean
+++ b/tests/lean/interactive/run.lean
@@ -103,9 +103,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 := {
diff --git a/tests/lean/server/diags.lean b/tests/lean/server/diags.lean
index 7a17722756b7..562131665798 100644
--- a/tests/lean/server/diags.lean
+++ b/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."
-    else
-      let diag := diags.getLast!
-      FS.writeFile "content_diag.json.produced" (toString <| toJson (diag : JsonRpc.Message))
+    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
-        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
diff --git a/tests/lean/server/edits.lean b/tests/lean/server/edits.lean
index 377e2306a205..5245ab44c91b 100644
--- a/tests/lean/server/edits.lean
+++ b/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."
-    else
-      let diag := diags.getLast!
-      FS.writeFile "edits_diag.json.produced" (toString <| toJson (diag : JsonRpc.Message))
+    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
-        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
diff --git a/tests/lean/server/init_exit_worker.lean b/tests/lean/server/init_exit_worker.lean
index 1032ffc90873..c4bafcbc21ba 100644
--- a/tests/lean/server/init_exit_worker.lean
+++ b/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

From 893ea994f47814b7816ebf82e346ce296df76b9f Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 14 Dec 2023 18:14:22 +0100
Subject: [PATCH 38/60] trigger CI


From 078cef09992e6a17f2c79657c8336a2c48e28944 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 15 Dec 2023 09:21:03 +0100
Subject: [PATCH 39/60] trigger CI again


From 561f2607cb1c9ddf777d0b00ed2c5928e4d598c5 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 15 Dec 2023 12:45:58 +0100
Subject: [PATCH 40/60] remove unused Option.forM, already in std4

---
 src/Init/Data/Option/Basic.lean | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/src/Init/Data/Option/Basic.lean b/src/Init/Data/Option/Basic.lean
index 84d3ed7e272c..c07201c8e697 100644
--- a/src/Init/Data/Option/Basic.lean
+++ b/src/Init/Data/Option/Basic.lean
@@ -46,10 +46,6 @@ def toMonad [Monad m] [Alternative m] : Option α → m α
   else
     return none
 
-@[inline] protected def forM [Monad m] (f : α → m Unit) : Option α → m Unit
-  | some a => f a
-  | none => pure ()
-
 theorem map_id : (Option.map id : Option α → Option α) = id :=
   funext (fun o => match o with | none => rfl | some _ => rfl)
 

From 6ec85ec80241ae5815afbd611b481150603f9948 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Mon, 18 Dec 2023 17:31:36 +0100
Subject: [PATCH 41/60] perf: keep alive old single-threaded cmdline driver for
 a little longer

---
 src/Lean/Elab/Frontend.lean            | 39 ++++++++++++++++++++++++++
 tests/lean/Process.lean.expected.out   |  4 +--
 tests/lean/dbgMacros.lean.expected.out |  4 +--
 3 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index e4dd403895d2..08af54818a6d 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -38,7 +38,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 }
@@ -98,6 +110,32 @@ def runFrontend
   let inputCtx := Parser.mkInputContext input fileName
   -- TODO: replace with `#lang` processing
   let lang := Language.Lean
+  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 }
+
+    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)
+      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
     fileSetupHandler? := none
@@ -118,4 +156,5 @@ def runFrontend
   let env := lang.getFinalEnv? snap |>.getD (← mkEmptyEnvironment)
   pure (env, !hasErrors)
 
+
 end Lean.Elab
diff --git a/tests/lean/Process.lean.expected.out b/tests/lean/Process.lean.expected.out
index dbe051b1ade5..fbeb40bc4ba3 100644
--- a/tests/lean/Process.lean.expected.out
+++ b/tests/lean/Process.lean.expected.out
@@ -1,5 +1,6 @@
-1
 hi!
+<stdin>:1:0: warning: using 'exit' to interrupt Lean
+1
 0
 "ho!\n"
 "hu!\n"
@@ -8,6 +9,5 @@ flush of broken pipe failed
 100000
 0
 0
-<stdin>:1:0: warning: using 'exit' to interrupt Lean
 0
 0
diff --git a/tests/lean/dbgMacros.lean.expected.out b/tests/lean/dbgMacros.lean.expected.out
index 54d6f8756c15..9b70d1b86045 100644
--- a/tests/lean/dbgMacros.lean.expected.out
+++ b/tests/lean/dbgMacros.lean.expected.out
@@ -1,10 +1,10 @@
 PANIC at f dbgMacros:2:14: unexpected zero
+PANIC at g dbgMacros:10:14: unreachable code has been reached
+PANIC at h dbgMacros:16:0: assertion violation: x != 0
 0
 9
-PANIC at g dbgMacros:10:14: unreachable code has been reached
 0
 0
-PANIC at h dbgMacros:16:0: assertion violation: x != 0
 0
 f2, x: 10
 11

From b50c9e2f777469f443cc45d70c071faa1a2a4872 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 21 Dec 2023 15:57:21 +0100
Subject: [PATCH 42/60] fix: don't report the same info tree twice

---
 src/Lean/Server/FileWorker.lean | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 1b488b0ba69d..1bd45b318f50 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -142,10 +142,11 @@ section Elab
       let infoTrees := match node.element.infoTree? with
         | some itree => st.infoTrees.push itree
         | none       => st.infoTrees
+      let mut st := { st with diagnostics, infoTrees, isFatal := node.element.isFatal }
       if st.hasBlocked && node.element.infoTree?.isSome then
         ctx.chanOut.send <| mkIleanInfoUpdateNotification m infoTrees
+        st := { st with infoTrees := #[] }
       -- we assume that only the last node in the tree sets `isFatal`
-      let st := { st with diagnostics, infoTrees, isFatal := node.element.isFatal }
       goSeq st cont node.children.toList
     goSeq st cont
       | [] => cont st

From 04800453d4ee623fa43d4233bbc06d3e0d39b1dd Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 21 Dec 2023 16:01:54 +0100
Subject: [PATCH 43/60] refactor: previous commit

---
 src/Lean/Server/FileWorker.lean | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 1bd45b318f50..90b6218a3d6f 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -139,14 +139,15 @@ section Elab
         node.element.diagnostics.interactiveDiags.map (·.toDiagnostic)
       if st.hasBlocked && !node.element.diagnostics.msgLog.isEmpty then
         ctx.chanOut.send <| mkPublishDiagnosticsNotification m diagnostics
-      let infoTrees := match node.element.infoTree? with
-        | some itree => st.infoTrees.push itree
-        | none       => st.infoTrees
-      let mut st := { st with diagnostics, infoTrees, isFatal := node.element.isFatal }
-      if st.hasBlocked && node.element.infoTree?.isSome then
-        ctx.chanOut.send <| mkIleanInfoUpdateNotification m infoTrees
-        st := { st with infoTrees := #[] }
       -- we assume that only the last node in the tree sets `isFatal`
+      let mut st := { st with diagnostics, isFatal := node.element.isFatal }
+
+      if let some itree := node.element.infoTree? then
+        let mut infoTrees := st.infoTrees.push itree
+        if st.hasBlocked then
+          ctx.chanOut.send <| mkIleanInfoUpdateNotification m infoTrees
+          infoTrees := #[]
+        st := { st with infoTrees }
       goSeq st cont node.children.toList
     goSeq st cont
       | [] => cont st

From 29b6612944cae62307730e717643081db4101092 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 21 Dec 2023 17:25:53 +0100
Subject: [PATCH 44/60] fix: report all info trees at `ileanInfoFinal`

---
 src/Lean/Server/FileWorker.lean | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 90b6218a3d6f..b8f6de62fd54 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -89,10 +89,17 @@ section Elab
       JsonRpc.Notification Lsp.LeanIleanInfoParams :=
     mkIleanInfoNotification "$/lean/ileanInfoFinal"
 
+  /-- 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 diagnostics encountered so far. -/
     diagnostics : Array Lsp.Diagnostic := #[]
-    infoTrees : Array Elab.InfoTree := #[]
+    /-- 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 should finish with a fatal progress notification. -/
     isFatal := false
 
   register_builtin_option server.reportDelayMs : Nat := {
@@ -130,7 +137,7 @@ section Elab
         ctx.chanOut.send <| mkPublishDiagnosticsNotification m st.diagnostics
       -- This will overwrite existing ilean info for the file, in case something
       -- went wrong during the incremental updates.
-      ctx.chanOut.send <| mkIleanInfoFinalNotification m st.infoTrees
+      ctx.chanOut.send <| mkIleanInfoFinalNotification m st.allInfoTrees
       return .pure ()
     go node st cont := do
       if (← IO.checkCanceled) then
@@ -143,11 +150,11 @@ section Elab
       let mut st := { st with diagnostics, isFatal := node.element.isFatal }
 
       if let some itree := node.element.infoTree? then
-        let mut infoTrees := st.infoTrees.push itree
+        let mut newInfoTrees := st.newInfoTrees.push itree
         if st.hasBlocked then
-          ctx.chanOut.send <| mkIleanInfoUpdateNotification m infoTrees
-          infoTrees := #[]
-        st := { st with infoTrees }
+          ctx.chanOut.send <| mkIleanInfoUpdateNotification m newInfoTrees
+          newInfoTrees := #[]
+        st := { st with newInfoTrees, allInfoTrees := st.allInfoTrees.push itree }
       goSeq st cont node.children.toList
     goSeq st cont
       | [] => cont st

From 5de7e2ceff573fa171d19d0b6e2a408ee4972821 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Sat, 23 Dec 2023 12:20:27 +0100
Subject: [PATCH 45/60] refactor: move interactive diagnostics caching from
 `Snapshot` to file worker

Avoids many problematic dependencies this early in the elab pipeline
---
 src/Lean/Elab/Frontend.lean                   |   2 +-
 src/Lean/Language/Basic.lean                  |  27 ++--
 src/Lean/Server/FileWorker.lean               | 135 ++++++++++++------
 src/Lean/Server/FileWorker/Utils.lean         |  34 +++--
 .../Server/FileWorker/WidgetRequests.lean     |  22 ---
 src/Lean/Server/README.md                     |  17 ++-
 src/Lean/Server/Rpc/RequestHandling.lean      |   4 +
 src/Lean/Server/Snapshots.lean                |   4 -
 8 files changed, 147 insertions(+), 98 deletions(-)

diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index 08af54818a6d..37244e8d7475 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -139,7 +139,7 @@ def runFrontend
   let ctx := { inputCtx with
     mainModuleName, opts, trustLevel
     fileSetupHandler? := none
-    clientHasWidgets := false
+    nextDiagsIdRef := (← IO.mkRef 0)
   }
   let snap ← lang.process none ctx
   let snaps := Language.toSnapshotTree snap
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index 00e1619fd577..7a51263c574d 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -10,7 +10,6 @@ Authors: Sebastian Ullrich
 
 import Lean.Message
 import Lean.Parser.Types
-import Lean.Widget.InteractiveDiagnostic
 
 set_option linter.missingDocs true
 
@@ -46,16 +45,18 @@ namespace Lean.Language
 structure Snapshot.Diagnostics where
   /-- Non-interactive message log. -/
   msgLog : MessageLog
-  /-- 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 : Array Widget.InteractiveDiagnostic
+  /--
+  Unique ID used by the file worker for caching diagnostics per message log. If `none`, no caching
+  is done, which should only be used for messages not containing any interactive elements.
+  -/
+  id? : Option Nat
 deriving Inhabited
 
 /-- The empty set of diagnostics. -/
 def Snapshot.Diagnostics.empty : Snapshot.Diagnostics where
   msgLog := .empty
-  interactiveDiags := #[]
+  -- nothing to cache
+  id? := none
 
 /--
   The base class of all snapshots: all the generic information the language server needs about a
@@ -191,8 +192,8 @@ structure ModuleProcessingContext where
   opts : Options
   /-- Kernel trust level. -/
   trustLevel : UInt32 := 0
-  /-- Whether to create interactive diagnostics. -/
-  clientHasWidgets : Bool
+  /-- Next ID to be used for `Snapshot.Diagnostics.id?`. -/
+  nextDiagsIdRef : IO.Ref Nat
   /--
     Callback available in server mode for building imports and retrieving per-library options using
     `lake setup-file`. -/
@@ -209,13 +210,9 @@ Creates snapshot message log from non-interactive message log, caching derived i
 diagnostics.
 -/
 def Snapshot.Diagnostics.ofMessageLog (msgLog : Lean.MessageLog) :
-    ProcessingM Snapshot.Diagnostics :=
-  return {
-    msgLog
-    interactiveDiags := (← msgLog.toList.toArray.mapM fun msg => do
-      let ctx ← read
-      Widget.msgToInteractiveDiagnostic ctx.fileMap msg ctx.clientHasWidgets)
-  }
+    ProcessingM Snapshot.Diagnostics := do
+  let id ← (← read).nextDiagsIdRef.modifyGet fun id => (id, id + 1)
+  return { msgLog, id? := some id }
 
 end Language
 open Language
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index b8f6de62fd54..f9d6a664095c 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -64,7 +64,7 @@ structure WorkerContext where
   Latest document version received by the client, used for filtering out notifications from
   previous versions.
   -/
-  maxDocVersion    : IO.Ref Nat
+  maxDocVersionRef : IO.Ref Nat
   hLog             : FS.Stream
   initParams       : InitializeParams
   processor        : Parser.InputContext → BaseIO Lean.Language.Lean.InitialSnapshot
@@ -93,8 +93,6 @@ section Elab
   private structure ReportSnapshotsState where
     /-- Whether we have waited for a snapshot to finish at least once (see debouncing below). -/
     hasBlocked := false
-    /-- All diagnostics encountered so far. -/
-    diagnostics : Array Lsp.Diagnostic := #[]
     /-- All info trees encountered so far. -/
     allInfoTrees : Array Elab.InfoTree := #[]
     /-- New info trees encountered since we last sent a .ilean update notification. -/
@@ -120,39 +118,51 @@ section Elab
     * 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) (m : DocumentMeta)
-      (snaps : Language.SnapshotTree) : BaseIO (Task Unit) := do
+  private partial def reportSnapshots (ctx : WorkerContext) (doc : EditableDocumentCore)
+      (prevDiagnosticsCache : DiagnosticsCache) : BaseIO (Task Unit) := do
     let t ← BaseIO.asTask do
       IO.sleep (server.reportDelayMs.get ctx.opts).toUInt32
     BaseIO.bindTask t fun _ =>
       start
   where
-    start := go snaps { : ReportSnapshotsState } fun st => do
+    start := go (Language.toSnapshotTree doc.initSnap) { : ReportSnapshotsState } fun st => do
       -- callback at the end of reporting
       if st.isFatal then
-        ctx.chanOut.send <| mkFileProgressAtPosNotification m 0 .fatalError
+        ctx.chanOut.send <| mkFileProgressAtPosNotification doc.meta 0 .fatalError
       else
-        ctx.chanOut.send <| mkFileProgressDoneNotification m
+        ctx.chanOut.send <| mkFileProgressDoneNotification doc.meta
       unless st.hasBlocked do
-        ctx.chanOut.send <| mkPublishDiagnosticsNotification m st.diagnostics
+        publishDiagnostics
       -- This will overwrite existing ilean info for the file, in case something
       -- went wrong during the incremental updates.
-      ctx.chanOut.send <| mkIleanInfoFinalNotification m st.allInfoTrees
+      ctx.chanOut.send <| mkIleanInfoFinalNotification doc.meta st.allInfoTrees
       return .pure ()
+    publishDiagnostics := do
+      ctx.chanOut.send <| mkPublishDiagnosticsNotification doc.meta <|
+        (← doc.diagnosticsRef.get).map (·.toDiagnostic)
     go node st cont := do
       if (← IO.checkCanceled) then
         return .pure ()
-      let diagnostics := st.diagnostics ++
-        node.element.diagnostics.interactiveDiags.map (·.toDiagnostic)
-      if st.hasBlocked && !node.element.diagnostics.msgLog.isEmpty then
-        ctx.chanOut.send <| mkPublishDiagnosticsNotification m diagnostics
+      let idiags ←
+        if let some cached := node.element.diagnostics.id?.bind prevDiagnosticsCache.find? then
+          pure cached
+        else
+          let idiags ← node.element.diagnostics.msgLog.toList.toArray.mapM
+            (Widget.msgToInteractiveDiagnostic doc.meta.text · ctx.clientHasWidgets)
+          if let some id := node.element.diagnostics.id? then
+            doc.diagnosticsCacheRef.modify (·.insert id idiags)
+          pure idiags
+      if !node.element.diagnostics.msgLog.isEmpty then
+        doc.diagnosticsRef.modify (· ++ idiags)
+        if st.hasBlocked then
+          publishDiagnostics
       -- we assume that only the last node in the tree sets `isFatal`
-      let mut st := { st with diagnostics, isFatal := node.element.isFatal }
+      let mut st := { st with isFatal := 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 m newInfoTrees
+          ctx.chanOut.send <| mkIleanInfoUpdateNotification doc.meta newInfoTrees
           newInfoTrees := #[]
         st := { st with newInfoTrees, allInfoTrees := st.allInfoTrees.push itree }
       goSeq st cont node.children.toList
@@ -161,9 +171,9 @@ section Elab
       | t::ts => do
         let mut st := st
         unless (← IO.hasFinished t.task) do
-          ctx.chanOut.send <| mkFileProgressAtPosNotification m t.range.start
+          ctx.chanOut.send <| mkFileProgressAtPosNotification doc.meta t.range.start
           if !st.hasBlocked then
-            ctx.chanOut.send <| mkPublishDiagnosticsNotification m st.diagnostics
+            publishDiagnostics
             st := { st with hasBlocked := true }
         BaseIO.bindTask t.task fun node =>
           go node st (goSeq · cont ts)
@@ -196,11 +206,11 @@ section Initialization
       if let some path := System.Uri.fileUriToPath? meta.uri then
         mainModuleName ← moduleNameOfFileName path none
     catch _ => pure ()
-    let maxDocVersion ← IO.mkRef 0
-    let chanOut ← mkLspOutputChannel maxDocVersion
+    let maxDocVersionRef ← IO.mkRef 0
+    let nextDiagsIdRef ← IO.mkRef 0
+    let chanOut ← mkLspOutputChannel maxDocVersionRef
     let processor ← Language.Lean.mkIncrementalProcessor {
-      opts, mainModuleName
-      clientHasWidgets
+      opts, mainModuleName, nextDiagsIdRef
       fileSetupHandler? := some fun imports => do
         let result ← setupFile meta imports fun stderrLine => do
           let progressDiagnostic := {
@@ -221,13 +231,17 @@ section Initialization
       initParams
       processor
       clientHasWidgets
-      maxDocVersion
+      maxDocVersionRef
       opts
     }
-    let reporter ← reportSnapshots ctx meta (Language.ToSnapshotTree.toSnapshotTree initSnap)
-    let doc : EditableDocument := { meta, initSnap, reporter }
-    return (ctx,
-    { doc
+    let doc : EditableDocumentCore := {
+      meta, initSnap
+      diagnosticsRef := (← IO.mkRef ∅)
+      diagnosticsCacheRef := (← IO.mkRef ∅)
+    }
+    let reporter ← reportSnapshots ctx doc ∅
+    return (ctx, {
+      doc := { doc with reporter }
       pendingRequests    := RBMap.empty
       rpcSessions        := RBMap.empty
       importCachingTask? := none
@@ -264,13 +278,18 @@ section Updates
     modify fun st => { st with pendingRequests := map st.pendingRequests }
 
   /-- Given the new document, updates editable doc state. -/
-  def updateDocument (meta : DocumentMeta) : WorkerM Unit := do
+  def updateDocument (oldDoc : EditableDocument) (meta : DocumentMeta) : WorkerM Unit := do
     let ctx ← read
     let initSnap ← ctx.processor meta.mkInputContext
-    let reporter ← reportSnapshots ctx meta (Language.ToSnapshotTree.toSnapshotTree initSnap)
-    modify fun st => { st with doc := { meta, initSnap, reporter } }
+    let doc : EditableDocumentCore := {
+      meta, initSnap
+      diagnosticsRef := (← IO.mkRef ∅)
+      diagnosticsCacheRef := (← IO.mkRef ∅)
+    }
+    let reporter ← reportSnapshots ctx doc (← oldDoc.diagnosticsCacheRef.get)
+    modify fun st => { st with doc := { doc with reporter } }
     -- we assume versions are monotonous
-    ctx.maxDocVersion.set meta.version
+    ctx.maxDocVersionRef.set meta.version
 end Updates
 
 /- Notifications are handled in the main thread. They may change global worker state
@@ -283,7 +302,7 @@ section NotificationHandling
     let newVersion := docId.version?.getD 0
     if ¬ changes.isEmpty then
       let newDocText := foldDocumentChanges changes oldDoc.meta.text
-      updateDocument ⟨docId.uri, newVersion, newDocText, oldDoc.meta.dependencyBuildMode⟩
+      updateDocument oldDoc ⟨docId.uri, newVersion, newDocText, oldDoc.meta.dependencyBuildMode⟩
 
   def handleCancelRequest (p : CancelParams) : WorkerM Unit := do
     updatePendingRequests (fun pendingRequests => pendingRequests.erase p.id)
@@ -340,6 +359,17 @@ 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
+    let diags ← st.doc.diagnosticsRef.get
+    return diags.filter fun diag =>
+      params.lineRange?.all fun ⟨s, e⟩ =>
+        -- does [s,e) intersect [diag.fullRange.start.line,diag.fullRange.end.line)?
+        s ≤ diag.fullRange.start.line ∧ diag.fullRange.start.line < e ∨
+        diag.fullRange.start.line ≤ s ∧ s < diag.fullRange.end.line
+
   def handleImportCompletionRequest (id : RequestID) (params : CompletionParams)
       : WorkerM (Task (Except Error AvailableImportsCache)) := do
     let ctx ← read
@@ -369,21 +399,40 @@ section MessageHandling
     let ctx ← read
     let st ← get
 
-    if method == "$/lean/rpc/connect" then
-      try
+    -- special cases
+    try
+      match method with
+      -- needs access to `rpcSessions`
+      | "$/lean/rpc/connect" =>
         let ps ← parseParams RpcConnectParams params
         let resp ← handleRpcConnect ps
         ctx.chanOut.send <| .response id (toJson resp)
-      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.doc.initSnap.stx params then
-        let importCachingTask ← handleImportCompletionRequest id params
-        set <| { st with importCachingTask? := some importCachingTask }
         return
+      | "$/lean/rpc/call" =>
+        let params ← parseParams Lsp.RpcCallParams params
+        -- needs access to `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
 
     -- we assume that any other request requires at least the the search path
     -- TODO: move into language-specific request handling
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index 00317dbb396a..a07f1745c2fa 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -47,14 +47,11 @@ private partial def mkCmdSnaps (initSnap : Language.Lean.InitialSnapshot) :
     AsyncList ElabTaskError Snapshot := Id.run do
   let some headerParsed := initSnap.success? | return .nil
   .delayed <| headerParsed.processed.task.bind fun headerProcessed => Id.run do
-    -- NOTE: this throws away interactive diagnostics of header errors but these are not interactive
-    -- anyway
     let some headerSuccess := headerProcessed.success? | return .pure <| .ok .nil
     return .pure <| .ok <| .cons {
       stx := initSnap.stx
       mpState := headerParsed.parserState
       cmdState := headerSuccess.cmdState
-      interactiveDiags := headerProcessed.diagnostics.interactiveDiags
     } <| .delayed <| headerSuccess.next.task.bind go
 where go cmdParsed :=
   cmdParsed.data.sig.task.bind fun sig =>
@@ -63,21 +60,38 @@ where go cmdParsed :=
         stx := cmdParsed.data.stx
         mpState := cmdParsed.data.parserState
         cmdState := finished.cmdState
-        interactiveDiags :=
-          cmdParsed.data.diagnostics.interactiveDiags ++ sig.diagnostics.interactiveDiags
       } (match cmdParsed.next? with
         | some next => .delayed <| next.task.bind go
         | none => .nil)
 
-/-- 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
+/-- A map from Diagnostics ID to resulting interactive objects. -/
+abbrev DiagnosticsCache := RBMap Nat (Array Widget.InteractiveDiagnostic) compare
+
+/--
+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
-  /-- State snapshots after header and each command. -/
+  /-- Initial processing snapshot. -/
   -- TODO: generalize to other languages by moving request handlers into `Language`
   initSnap : Language.Lean.InitialSnapshot
+  /-- Old representation for backward compatibility. -/
   cmdSnaps : AsyncList ElabTaskError Snapshot := mkCmdSnaps initSnap
+  /--
+  Interactive versions of diagnostics reported so far. Filled by `reportSnapshots` and read by
+  `handleGetInteractiveDiagnosticsRequest`.
+  -/
+  diagnosticsRef : IO.Ref (Array Widget.InteractiveDiagnostic)
+  /--
+  Cache of interactive diagnostics reported so far, for reuse in future document versions. May be
+  a subset of `diagnosticsRef` if there any `Snapshot.Diagnostics.id?` are empty.
+  -/
+  diagnosticsCacheRef : IO.Ref DiagnosticsCache
+
+/-- `EditableDocumentCore` with reporter task. -/
+structure EditableDocument extends EditableDocumentCore where
   /--
     Task reporting processing status back to client. We store it here for implementing
     `waitForDiagnostics`. -/
diff --git a/src/Lean/Server/FileWorker/WidgetRequests.lean b/src/Lean/Server/FileWorker/WidgetRequests.lean
index d7eb9376bab0..8e0de0b6c574 100644
--- a/src/Lean/Server/FileWorker/WidgetRequests.lean
+++ b/src/Lean/Server/FileWorker/WidgetRequests.lean
@@ -93,28 +93,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.filter fun diag =>
-        params.lineRange?.all fun ⟨s, e⟩ =>
-          -- does [s,e) intersect [diag.fullRange.start.line,diag.fullRange.end.line)?
-          s ≤ diag.fullRange.start.line ∧ diag.fullRange.start.line < e ∨
-          diag.fullRange.start.line ≤ s ∧ s < diag.fullRange.end.line
-    pure <| diags?.getD #[]
-
-builtin_initialize
-  registerBuiltinRpcProcedure
-    `Lean.Widget.getInteractiveDiagnostics
-    GetInteractiveDiagnosticsParams
-    (Array InteractiveDiagnostic)
-    getInteractiveDiagnostics
-
 structure GetGoToLocationParams where
   kind : GoToKind
   info : WithRpcRef InfoWithCtx
diff --git a/src/Lean/Server/README.md b/src/Lean/Server/README.md
index d40fed380754..7f787507afd1 100644
--- a/src/Lean/Server/README.md
+++ b/src/Lean/Server/README.md
@@ -49,24 +49,35 @@ In Lean 4, the situation is different as `.olean` artifacts are required to exis
 
 The actual processing of the opened file is left to the generic `Lean.Language` interface.
 It currently has a single implementation `Lean.Language.Lean` for the main Lean 4 language but other languages will be loadable via a `#lang` file annotation in the future.
-The interface is shared with the cmdline driver that is used for building Lean files but the core concept of *snapshots* produced by a `Language` is mostly of interest to the worker: snapshots are how processing data is saved between versions of a file such as to make processing on file edits incremental.
+The interface is shared with the cmdline driver that is used for building Lean files but the core concept of *snapshots* produced by a `Language` 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 `Language`; 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 `Language.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 `Language` given the new file content and asynchronously wait on all these nodes and report the processing status stored therein to the client (`Lean.Server.FileWorker.reportSnapshots`).
+After loading the file and after each edit, the server will obtain this tree from the `Language` given the new file content and asynchronously wait on all these nodes and report the processing status (diagnostics and progress information) stored therein to the client (`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 unique ID in `Snapshot.Diagnostics` that the reporting task uses to detect such snapshot reuse and not recreate interactive diagnostic objects in this case.
+
 ## Code style
 
 Comments should exist to denote specifics of our implementation but, for
diff --git a/src/Lean/Server/Rpc/RequestHandling.lean b/src/Lean/Server/Rpc/RequestHandling.lean
index 44d67cba55c4..3ecd3598c931 100644
--- a/src/Lean/Server/Rpc/RequestHandling.lean
+++ b/src/Lean/Server/Rpc/RequestHandling.lean
@@ -31,6 +31,10 @@ private unsafe def evalRpcProcedureUnsafe (env : Environment) (opts : Options) (
 opaque evalRpcProcedure (env : Environment) (opts : Options) (procName : Name) :
     Except String RpcProcedure
 
+def runRpcProcedure (p : RpcProcedure) :
+    (sessionId : UInt64) → Json → RequestM (RequestTask Json) :=
+  p.wrapper
+
 open RequestM in
 def handleRpcCall (p : Lsp.RpcCallParams) : RequestM (RequestTask Json) := do
   -- The imports are finished at this point, because the handleRequest function
diff --git a/src/Lean/Server/Snapshots.lean b/src/Lean/Server/Snapshots.lean
index d33fda462cc0..983343a30819 100644
--- a/src/Lean/Server/Snapshots.lean
+++ b/src/Lean/Server/Snapshots.lean
@@ -29,10 +29,6 @@ structure Snapshot where
   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 : Array Widget.InteractiveDiagnostic
 
 namespace Snapshot
 

From 8bd6612a94d071199941b04fee44f16ff5db273e Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 1 Feb 2024 11:39:26 +0100
Subject: [PATCH 46/60] clean up Lean.Server changes

---
 src/Lean/Server/FileWorker.lean          | 41 +++++++++++++-----------
 src/Lean/Server/Rpc/RequestHandling.lean |  4 ---
 2 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index c6ff0377a16e..b7b412dc885c 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -107,6 +107,7 @@ section Elab
       document edit in order to reduce flickering"
   }
 
+  open Language in
   /--
     Reports status of a snapshot tree incrementally to the user: progress,
     diagnostics, .ilean reference information.
@@ -125,7 +126,7 @@ section Elab
     BaseIO.bindTask t fun _ =>
       start
   where
-    start := go (Language.toSnapshotTree doc.initSnap) { : ReportSnapshotsState } fun st => do
+    start := go (toSnapshotTree doc.initSnap) { : ReportSnapshotsState } fun st => do
       -- callback at the end of reporting
       if st.isFatal then
         ctx.chanOut.send <| mkFileProgressAtPosNotification doc.meta 0 .fatalError
@@ -140,24 +141,26 @@ section Elab
     publishDiagnostics := do
       ctx.chanOut.send <| mkPublishDiagnosticsNotification doc.meta <|
         (← doc.diagnosticsRef.get).map (·.toDiagnostic)
-    go node st cont := do
+    go (node : SnapshotTree) (st : ReportSnapshotsState)
+        (cont : ReportSnapshotsState → BaseIO (Task Unit)) : BaseIO (Task Unit) := do
       if (← IO.checkCanceled) then
         return .pure ()
-      let idiags ←
-        if let some cached := node.element.diagnostics.id?.bind prevDiagnosticsCache.find? then
-          pure cached
-        else
-          let idiags ← node.element.diagnostics.msgLog.toList.toArray.mapM
-            (Widget.msgToInteractiveDiagnostic doc.meta.text · ctx.clientHasWidgets)
-          if let some id := node.element.diagnostics.id? then
-            doc.diagnosticsCacheRef.modify (·.insert id idiags)
-          pure idiags
+
       if !node.element.diagnostics.msgLog.isEmpty then
+        let idiags ←
+          if let some cached := node.element.diagnostics.id?.bind prevDiagnosticsCache.find? then
+            pure cached
+          else
+            let idiags ← node.element.diagnostics.msgLog.toList.toArray.mapM
+              (Widget.msgToInteractiveDiagnostic doc.meta.text · ctx.clientHasWidgets)
+            if let some id := node.element.diagnostics.id? then
+              doc.diagnosticsCacheRef.modify (·.insert id idiags)
+            pure idiags
         doc.diagnosticsRef.modify (· ++ idiags)
         if st.hasBlocked then
           publishDiagnostics
-      -- we assume that only the last node in the tree sets `isFatal`
-      let mut st := { st with isFatal := node.element.isFatal }
+
+      let mut st := { st with isFatal := st.isFatal || node.element.isFatal }
 
       if let some itree := node.element.infoTree? then
         let mut newInfoTrees := st.newInfoTrees.push itree
@@ -165,8 +168,10 @@ section Elab
           ctx.chanOut.send (← mkIleanInfoUpdateNotification doc.meta newInfoTrees)
           newInfoTrees := #[]
         st := { st with newInfoTrees, allInfoTrees := st.allInfoTrees.push itree }
+
       goSeq st cont node.children.toList
-    goSeq st cont
+    goSeq (st : ReportSnapshotsState) (cont : ReportSnapshotsState → BaseIO (Task Unit)) :
+        List (SnapshotTask SnapshotTree) → BaseIO (Task Unit)
       | [] => cont st
       | t::ts => do
         let mut st := st
@@ -215,6 +220,7 @@ section Initialization
         let result ← 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
@@ -253,7 +259,6 @@ section Initialization
         elaboration tasks here. -/
     mkLspOutputChannel maxDocVersion : IO (IO.Channel JsonRpc.Message) := do
       let chanOut ← IO.Channel.new
-      -- most recent document version seen in notifications
       let _ ← chanOut.forAsync (prio := .dedicated) fun msg => do
         -- discard outdated notifications; note that in contrast to responses, notifications can
         -- always be silently discarded
@@ -288,7 +293,7 @@ section Updates
     }
     let reporter ← reportSnapshots ctx doc (← oldDoc.diagnosticsCacheRef.get)
     modify fun st => { st with doc := { doc with reporter } }
-    -- we assume versions are monotonous
+    -- we assume version updates are monotonous and that we are on the main thread
     ctx.maxDocVersionRef.set meta.version
 end Updates
 
@@ -402,7 +407,7 @@ section MessageHandling
     -- special cases
     try
       match method with
-      -- needs access to `rpcSessions`
+      -- needs access to `WorkerState.rpcSessions`
       | "$/lean/rpc/connect" =>
         let ps ← parseParams RpcConnectParams params
         let resp ← handleRpcConnect ps
@@ -410,7 +415,7 @@ section MessageHandling
         return
       | "$/lean/rpc/call" =>
         let params ← parseParams Lsp.RpcCallParams params
-        -- needs access to `diagnosticsRef`
+        -- 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
diff --git a/src/Lean/Server/Rpc/RequestHandling.lean b/src/Lean/Server/Rpc/RequestHandling.lean
index 3ecd3598c931..44d67cba55c4 100644
--- a/src/Lean/Server/Rpc/RequestHandling.lean
+++ b/src/Lean/Server/Rpc/RequestHandling.lean
@@ -31,10 +31,6 @@ private unsafe def evalRpcProcedureUnsafe (env : Environment) (opts : Options) (
 opaque evalRpcProcedure (env : Environment) (opts : Options) (procName : Name) :
     Except String RpcProcedure
 
-def runRpcProcedure (p : RpcProcedure) :
-    (sessionId : UInt64) → Json → RequestM (RequestTask Json) :=
-  p.wrapper
-
 open RequestM in
 def handleRpcCall (p : Lsp.RpcCallParams) : RequestM (RequestTask Json) := do
   -- The imports are finished at this point, because the handleRequest function

From a0a65b6b93cd65aaffa4ade5599c2caa0593e7aa Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Tue, 13 Feb 2024 14:58:35 +0100
Subject: [PATCH 47/60] refactor: remove unique ID indirection for interactive
 diagnostics caching

---
 src/Lean/Elab/Frontend.lean           |  1 -
 src/Lean/Language/Basic.lean          | 16 +++++-----
 src/Lean/Server/FileWorker.lean       | 45 ++++++++++++++++-----------
 src/Lean/Server/FileWorker/Utils.lean |  5 ---
 src/Lean/Server/README.md             |  3 +-
 5 files changed, 37 insertions(+), 33 deletions(-)

diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index ecd7abe90782..4f0ce722136a 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -140,7 +140,6 @@ def runFrontend
   let ctx := { inputCtx with
     mainModuleName, opts, trustLevel
     fileSetupHandler? := none
-    nextDiagsIdRef := (← IO.mkRef 0)
   }
   let snap ← lang.process none ctx
   let snaps := Language.toSnapshotTree snap
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index 7a51263c574d..0b8c034e171c 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -46,17 +46,20 @@ structure Snapshot.Diagnostics where
   /-- Non-interactive message log. -/
   msgLog : MessageLog
   /--
-  Unique ID used by the file worker for caching diagnostics per message log. If `none`, no caching
-  is done, which should only be used for messages not containing any interactive elements.
+  Dynamic mutable slot usable by the language server for caching interactive diagnostics.  If
+  `none`, no caching is done, which should only be used for messages not containing any interactive
+  elements.
+
+  See also section "Communication" in Lean/Server/README.md.
   -/
-  id? : Option Nat
+  cacheRef? : Option (IO.Ref (Option Dynamic))
 deriving Inhabited
 
 /-- The empty set of diagnostics. -/
 def Snapshot.Diagnostics.empty : Snapshot.Diagnostics where
   msgLog := .empty
   -- nothing to cache
-  id? := none
+  cacheRef? := none
 
 /--
   The base class of all snapshots: all the generic information the language server needs about a
@@ -192,8 +195,6 @@ structure ModuleProcessingContext where
   opts : Options
   /-- Kernel trust level. -/
   trustLevel : UInt32 := 0
-  /-- Next ID to be used for `Snapshot.Diagnostics.id?`. -/
-  nextDiagsIdRef : IO.Ref Nat
   /--
     Callback available in server mode for building imports and retrieving per-library options using
     `lake setup-file`. -/
@@ -211,8 +212,7 @@ diagnostics.
 -/
 def Snapshot.Diagnostics.ofMessageLog (msgLog : Lean.MessageLog) :
     ProcessingM Snapshot.Diagnostics := do
-  let id ← (← read).nextDiagsIdRef.modifyGet fun id => (id, id + 1)
-  return { msgLog, id? := some id }
+  return { msgLog, cacheRef? := some (← IO.mkRef none) }
 
 end Language
 open Language
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index b7b412dc885c..0f50a8842fff 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -107,11 +107,22 @@ section Elab
       document edit in order to reduce flickering"
   }
 
+  /--
+  Type of cache stored in `Snapshot.Diagnostics.cacheRef?`.
+
+  See also section "Communication" in Lean/Server/README.md.
+  -/
+  structure CachedInteractiveDiagnostics 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
@@ -119,8 +130,8 @@ section Elab
     * 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)
-      (prevDiagnosticsCache : DiagnosticsCache) : BaseIO (Task Unit) := do
+  private partial def reportSnapshots (ctx : WorkerContext) (doc : EditableDocumentCore) :
+      BaseIO (Task Unit) := do
     let t ← BaseIO.asTask do
       IO.sleep (server.reportDelayMs.get ctx.opts).toUInt32
     BaseIO.bindTask t fun _ =>
@@ -147,16 +158,17 @@ section Elab
         return .pure ()
 
       if !node.element.diagnostics.msgLog.isEmpty then
-        let idiags ←
-          if let some cached := node.element.diagnostics.id?.bind prevDiagnosticsCache.find? then
-            pure cached
+        let diags ←
+          if let some cached ← node.element.diagnostics.cacheRef?.bindM fun cacheRef => do
+              return (← cacheRef.get).bind (·.get? CachedInteractiveDiagnostics) then
+            pure cached.diags
           else
-            let idiags ← node.element.diagnostics.msgLog.toList.toArray.mapM
+            let diags ← node.element.diagnostics.msgLog.toList.toArray.mapM
               (Widget.msgToInteractiveDiagnostic doc.meta.text · ctx.clientHasWidgets)
-            if let some id := node.element.diagnostics.id? then
-              doc.diagnosticsCacheRef.modify (·.insert id idiags)
-            pure idiags
-        doc.diagnosticsRef.modify (· ++ idiags)
+            if let some cacheRef := node.element.diagnostics.cacheRef? then
+              cacheRef.set <| some <| .mk { diags : CachedInteractiveDiagnostics }
+            pure diags
+        doc.diagnosticsRef.modify (· ++ diags)
         if st.hasBlocked then
           publishDiagnostics
 
@@ -212,10 +224,9 @@ section Initialization
         mainModuleName ← moduleNameOfFileName path none
     catch _ => pure ()
     let maxDocVersionRef ← IO.mkRef 0
-    let nextDiagsIdRef ← IO.mkRef 0
     let chanOut ← mkLspOutputChannel maxDocVersionRef
     let processor ← Language.Lean.mkIncrementalProcessor {
-      opts, mainModuleName, nextDiagsIdRef
+      opts, mainModuleName
       fileSetupHandler? := some fun imports => do
         let result ← setupFile meta imports fun stderrLine => do
           let progressDiagnostic := {
@@ -243,9 +254,8 @@ section Initialization
     let doc : EditableDocumentCore := {
       meta, initSnap
       diagnosticsRef := (← IO.mkRef ∅)
-      diagnosticsCacheRef := (← IO.mkRef ∅)
     }
-    let reporter ← reportSnapshots ctx doc ∅
+    let reporter ← reportSnapshots ctx doc
     return (ctx, {
       doc := { doc with reporter }
       pendingRequests    := RBMap.empty
@@ -283,15 +293,14 @@ section Updates
     modify fun st => { st with pendingRequests := map st.pendingRequests }
 
   /-- Given the new document, updates editable doc state. -/
-  def updateDocument (oldDoc : EditableDocument) (meta : DocumentMeta) : WorkerM Unit := do
+  def updateDocument (meta : DocumentMeta) : WorkerM Unit := do
     let ctx ← read
     let initSnap ← ctx.processor meta.mkInputContext
     let doc : EditableDocumentCore := {
       meta, initSnap
       diagnosticsRef := (← IO.mkRef ∅)
-      diagnosticsCacheRef := (← IO.mkRef ∅)
     }
-    let reporter ← reportSnapshots ctx doc (← oldDoc.diagnosticsCacheRef.get)
+    let reporter ← reportSnapshots ctx doc
     modify fun st => { st with doc := { doc with reporter } }
     -- we assume version updates are monotonous and that we are on the main thread
     ctx.maxDocVersionRef.set meta.version
@@ -307,7 +316,7 @@ section NotificationHandling
     let newVersion := docId.version?.getD 0
     if ¬ changes.isEmpty then
       let newDocText := foldDocumentChanges changes oldDoc.meta.text
-      updateDocument oldDoc ⟨docId.uri, newVersion, newDocText, oldDoc.meta.dependencyBuildMode⟩
+      updateDocument ⟨docId.uri, newVersion, newDocText, oldDoc.meta.dependencyBuildMode⟩
 
   def handleCancelRequest (p : CancelParams) : WorkerM Unit := do
     updatePendingRequests (fun pendingRequests => pendingRequests.erase p.id)
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index a07f1745c2fa..cfe91933cb9a 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -84,11 +84,6 @@ structure EditableDocumentCore where
   `handleGetInteractiveDiagnosticsRequest`.
   -/
   diagnosticsRef : IO.Ref (Array Widget.InteractiveDiagnostic)
-  /--
-  Cache of interactive diagnostics reported so far, for reuse in future document versions. May be
-  a subset of `diagnosticsRef` if there any `Snapshot.Diagnostics.id?` are empty.
-  -/
-  diagnosticsCacheRef : IO.Ref DiagnosticsCache
 
 /-- `EditableDocumentCore` with reporter task. -/
 structure EditableDocument extends EditableDocumentCore where
diff --git a/src/Lean/Server/README.md b/src/Lean/Server/README.md
index 7f787507afd1..97004261b632 100644
--- a/src/Lean/Server/README.md
+++ b/src/Lean/Server/README.md
@@ -76,7 +76,8 @@ 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 unique ID in `Snapshot.Diagnostics` that the reporting task uses to detect such snapshot reuse and not recreate interactive diagnostic objects in this case.
+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
 

From 58d5b58de6b04d2b9535dffb0459aab3cae3b342 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Tue, 20 Feb 2024 12:53:52 +0100
Subject: [PATCH 48/60] Remove `Language` interface for now, to be refined in
 `hash-lang`

---
 src/Lean/Elab/Frontend.lean           |  5 ++---
 src/Lean/Language/Basic.lean          | 26 +++-----------------------
 src/Lean/Language/Lean.lean           | 10 +++-------
 src/Lean/Server/FileWorker.lean       |  2 +-
 src/Lean/Server/FileWorker/Utils.lean |  1 -
 src/Lean/Server/README.md             |  9 ++++-----
 6 files changed, 13 insertions(+), 40 deletions(-)

diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index ad88979edaaa..258480f8b3ac 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -110,7 +110,6 @@ def runFrontend
     : IO (Environment × Bool) := do
   let inputCtx := Parser.mkInputContext input fileName
   -- TODO: replace with `#lang` processing
-  let lang := Language.Lean
   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
@@ -142,7 +141,7 @@ def runFrontend
     mainModuleName, opts, trustLevel
     fileSetupHandler? := none
   }
-  let snap ← lang.process none ctx
+  let snap ← Language.Lean.process none ctx
   let snaps := Language.toSnapshotTree snap
   snaps.runAndReport opts
   if let some ileanFileName := ileanFileName? then
@@ -154,7 +153,7 @@ def runFrontend
   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 := lang.getFinalEnv? snap |>.getD (← mkEmptyEnvironment)
+  let env := Language.Lean.getFinalEnv? snap |>.getD (← mkEmptyEnvironment)
   pure (env, !hasErrors)
 
 
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index 0b8c034e171c..0804c2d8c72c 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -215,34 +215,14 @@ def Snapshot.Diagnostics.ofMessageLog (msgLog : Lean.MessageLog) :
   return { msgLog, cacheRef? := some (← IO.mkRef none) }
 
 end Language
-open Language
-
-/-- Definition of a language processor that can be driven by the cmdline or language server. -/
-structure Language where
-  /--
-    Type of snapshot returned by `process`. It can be converted to a graph of homogeneous snapshot
-    types via `ToSnapshotTree`.  -/
-  InitialSnapshot : Type
-  /-- Instance for transforming the initial snapshot into a snapshot tree for reporting. -/
-  [instToSnapshotTree : ToSnapshotTree InitialSnapshot]
-  /--
-    Processes input into snapshots, potentially reusing information from a previous run.
-    Constructing the initial snapshot is assumed to be cheap enough that it can be done
-    synchronously, which simplifies use of this function. -/
-  process (old? : Option InitialSnapshot) : ProcessingM InitialSnapshot
-  -- TODO: is this the right interface for other languages as well?
-  /-- Gets final environment, if any, that is to be used for persisting, code generation, etc. -/
-  getFinalEnv? : InitialSnapshot → Option Environment
-
-instance (lang : Language) : ToSnapshotTree lang.InitialSnapshot := lang.instToSnapshotTree
 
 /--
   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 (lang : Language) (ctx : ModuleProcessingContext) :
-    BaseIO (Parser.InputContext → BaseIO lang.InitialSnapshot) := do
+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 ← lang.process (← oldRef.get) { ctx, ictx with }
+    let snap ← process (← oldRef.get) { ctx, ictx with }
     oldRef.set (some snap)
     return snap
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index b18457e4849f..0149460c8457 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -245,7 +245,7 @@ General notes:
   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 processLean (old? : Option InitialSnapshot) : ProcessingM InitialSnapshot := do
+partial def process (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
@@ -497,7 +497,8 @@ where
         }
       }
 
-private partial def getFinalEnv? (snap : InitialSnapshot) : Option Environment := do
+/-- Waits for and returns final environment, if importing was successful. -/
+partial def getFinalEnv? (snap : InitialSnapshot) : Option Environment := do
   let snap ← snap.success?
   let snap ← snap.processed.get.success?
   goCmd snap.next.get
@@ -508,8 +509,3 @@ where goCmd snap :=
     snap.data.sig.get.finished.get.cmdState.env
 
 end Lean
-
-/-- The Lean language processor. -/
-abbrev Lean : Language where
-  process := Lean.processLean
-  getFinalEnv? := Lean.getFinalEnv?
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 3d000caf5540..66f5e30c9998 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -227,7 +227,7 @@ section Initialization
     catch _ => pure ()
     let maxDocVersionRef ← IO.mkRef 0
     let chanOut ← mkLspOutputChannel maxDocVersionRef
-    let processor ← Language.Lean.mkIncrementalProcessor {
+    let processor ← Language.mkIncrementalProcessor Language.Lean.process {
       opts, mainModuleName
       fileSetupHandler? := some fun imports => do
         let result ← setupFile meta imports fun stderrLine => do
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index 9902540634e5..84b4f018e6b3 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -76,7 +76,6 @@ structure EditableDocumentCore where
   /-- The document. -/
   meta       : DocumentMeta
   /-- Initial processing snapshot. -/
-  -- TODO: generalize to other languages by moving request handlers into `Language`
   initSnap : Language.Lean.InitialSnapshot
   /-- Old representation for backward compatibility. -/
   cmdSnaps : AsyncList ElabTaskError Snapshot := mkCmdSnaps initSnap
diff --git a/src/Lean/Server/README.md b/src/Lean/Server/README.md
index 97004261b632..34ed71e50ce5 100644
--- a/src/Lean/Server/README.md
+++ b/src/Lean/Server/README.md
@@ -47,15 +47,14 @@ In Lean 4, the situation is different as `.olean` artifacts are required to exis
 
 ### Worker architecture
 
-The actual processing of the opened file is left to the generic `Lean.Language` interface.
-It currently has a single implementation `Lean.Language.Lean` for the main Lean 4 language but other languages will be loadable via a `#lang` file annotation in the future.
-The interface is shared with the cmdline driver that is used for building Lean files but the core concept of *snapshots* produced by a `Language` 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 `Language`; 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.
+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 `Language.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 `Language` given the new file content and asynchronously wait on all these nodes and report the processing status (diagnostics and progress information) stored therein to the client (`Lean.Server.FileWorker.reportSnapshots`).
+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 these nodes and report the processing status (diagnostics and progress information) stored therein to the client (`Lean.Server.FileWorker.reportSnapshots`).
 
 (TODO: the remainder is currently hard-coded for the `Lean` language)
 

From 4665f41cbb7813844484668df1fde3bd96ecf867 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Wed, 21 Feb 2024 17:35:36 +0100
Subject: [PATCH 49/60] Move `setup-file` logic back from language processor
 into file worker

---
 src/Lean/Elab/Frontend.lean     |  8 ++--
 src/Lean/Language/Basic.lean    | 28 ++++++++++---
 src/Lean/Language/Lean.lean     | 66 ++++++-----------------------
 src/Lean/Server/FileWorker.lean | 73 ++++++++++++++++++++++++---------
 4 files changed, 93 insertions(+), 82 deletions(-)

diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index 258480f8b3ac..16a438551a23 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -137,11 +137,9 @@ def runFrontend
 
     return (s.commandState.env, !s.commandState.messages.hasErrors)
 
-  let ctx := { inputCtx with
-    mainModuleName, opts, trustLevel
-    fileSetupHandler? := none
-  }
-  let snap ← Language.Lean.process none ctx
+  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
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index 0804c2d8c72c..e2701d728c24 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -195,16 +195,17 @@ structure ModuleProcessingContext where
   opts : Options
   /-- Kernel trust level. -/
   trustLevel : UInt32 := 0
-  /--
-    Callback available in server mode for building imports and retrieving per-library options using
-    `lake setup-file`. -/
-  fileSetupHandler? : Option (Array Import → IO Server.FileWorker.FileSetupResult)
 
 /-- 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 := ReaderT ProcessingContext BaseIO
+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, caching derived interactive
@@ -213,6 +214,23 @@ diagnostics.
 def Snapshot.Diagnostics.ofMessageLog (msgLog : Lean.MessageLog) :
     ProcessingM Snapshot.Diagnostics := do
   return { msgLog, cacheRef? := 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
 
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 0149460c8457..78f032ad9518 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -142,18 +142,16 @@ partial def CommandParsedSnapshot.cancel (snap : CommandParsedSnapshot) : BaseIO
     let _ ← IO.mapTask (sync := true) (·.cancel) next.task
 
 /-- State after successful importing. -/
-structure HeaderProcessedSucessfully where
+structure HeaderProcessedSuccessfully where
   /-- The resulting initial elaboration state. -/
   cmdState : Command.State
-  /-- The search path communicated by `lake setup-file`, if any. -/
-  srcSearchPath : SearchPath
   /-- First command task (there is always at least a terminal command). -/
   next : SnapshotTask CommandParsedSnapshot
 
 /-- State after the module header has been processed including imports. -/
 structure HeaderProcessedSnapshot extends Snapshot where
   /-- State after successful importing. -/
-  success? : Option HeaderProcessedSucessfully
+  success? : Option HeaderProcessedSuccessfully
   isFatal := success?.isNone
 instance : ToSnapshotTree HeaderProcessedSnapshot where
   toSnapshotTree s := ⟨s.toSnapshot, #[] |>
@@ -182,7 +180,7 @@ instance : ToSnapshotTree HeaderParsedSnapshot where
 
 /-- Shortcut accessor to the final header state, if successful. -/
 def HeaderParsedSnapshot.processedSuccessfully (snap : HeaderParsedSnapshot) :
-    SnapshotTask (Option HeaderProcessedSucessfully) :=
+    SnapshotTask (Option HeaderProcessedSuccessfully) :=
   snap.success?.bind (·.processed.map (sync := true) (·.success?)) |>.getD (.pure none)
 
 /-- Initial snapshot of the Lean language processor: a "header parsed" snapshot. -/
@@ -201,7 +199,7 @@ abbrev LeanProcessingM := LeanProcessingT BaseIO
 instance : MonadLift LeanProcessingM (LeanProcessingT IO) where
   monadLift := fun act ctx => act ctx
 
-instance : MonadLift ProcessingM LeanProcessingM where
+instance : MonadLift (ProcessingT m) (LeanProcessingT m) where
   monadLift := fun act ctx => act ctx.toProcessingContext
 
 /--
@@ -211,15 +209,6 @@ compared to it.
 def isBeforeEditPos (pos : String.Pos) : LeanProcessingM Bool := do
   return (← read).firstDiffPos?.any (pos < ·)
 
-private def diagnosticsOfHeaderError (msg : String) : ProcessingM Snapshot.Diagnostics := do
-  let msgLog := MessageLog.empty.add {
-    fileName := "<input>"
-    pos := ⟨0, 0⟩
-    endPos := dbgTraceVal <| (← 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. -/
@@ -229,9 +218,6 @@ private def withHeaderExceptions (ex : Snapshot → α) (act : LeanProcessingT I
   | .error e => return ex { diagnostics := (← diagnosticsOfHeaderError e.toString) }
   | .ok a => return a
 
-/-- 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
-
 /-- Entry point of the Lean language processor. -/
 /-
 General notes:
@@ -245,7 +231,10 @@ General notes:
   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 (old? : Option InitialSnapshot) : ProcessingM InitialSnapshot := do
+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
@@ -320,39 +309,11 @@ where
     SnapshotTask.ofIO ⟨0, ctx.input.endPos⟩ <|
     ReaderT.run (r := ctx) <|  -- re-enter reader in new task
     withHeaderExceptions (α := HeaderProcessedSnapshot) ({ · with success? := none }) do
-      -- use `lake setup-file` if in language server
-      let fileSetupResult ← if let some handler := ctx.fileSetupHandler? then
-        handler (Elab.headerToImports stx)
-      else
-        pure { kind := .success, srcSearchPath := ∅, fileOptions := .empty }
-      match fileSetupResult.kind with
-      | .importsOutOfDate =>
-        return {
-          diagnostics := (← diagnosticsOfHeaderError
-            "Imports are out of date and must be rebuilt; \
-             use the \"Restart File\" command in your editor.")
-          success? := none
-        }
-      | .error msg =>
-        return {
-          diagnostics := (← diagnosticsOfHeaderError msg)
-          success? := none
-        }
-      | _ => pure ()
-
-      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 { diagnostics := .empty, success? := none }
-
+      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) fileSetupResult.fileOptions
-
+      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
@@ -361,7 +322,7 @@ where
         return { diagnostics, success? := none }
 
       let headerEnv := headerEnv.setMainModule ctx.mainModuleName
-      let cmdState := Elab.Command.mkState headerEnv msgLog opts
+      let cmdState := Elab.Command.mkState headerEnv msgLog ctx.opts
       let cmdState := { cmdState with infoState := {
         enabled := true
         trees := #[Elab.InfoTree.context (.commandCtx {
@@ -383,7 +344,6 @@ where
         infoTree? := cmdState.infoState.trees[0]!
         success? := some {
           cmdState
-          srcSearchPath := fileSetupResult.srcSearchPath
           next := (← parseCmd none parserState cmdState)
         }
       }
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 66f5e30c9998..81bd4fd56253 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -207,6 +207,7 @@ structure AvailableImportsCache where
 
 structure WorkerState where
   doc                : EditableDocument
+  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
@@ -215,6 +216,53 @@ 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
+def setupImports (meta : DocumentMeta) (chanOut : Channel JsonRpc.Message)
+    (srcSearchPathPromise : Promise SearchPath) (stx : Syntax) :
+    Language.ProcessingT IO (Except Language.Lean.HeaderProcessedSnapshot Options) := do
+  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.")
+      success? := none
+    }
+  | .error msg =>
+    return .error {
+      diagnostics := (← diagnosticsOfHeaderError msg)
+      success? := none
+    }
+  | _ => pure ()
+
+  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, success? := none }
+
+  srcSearchPathPromise.resolve fileSetupResult.srcSearchPath
+  return .ok fileSetupResult.fileOptions
+
 /- Worker initialization sequence. -/
 section Initialization
   def initializeWorker (meta : DocumentMeta) (o e : FS.Stream) (initParams : InitializeParams) (opts : Options)
@@ -227,22 +275,10 @@ section Initialization
     catch _ => pure ()
     let maxDocVersionRef ← IO.mkRef 0
     let chanOut ← mkLspOutputChannel maxDocVersionRef
-    let processor ← Language.mkIncrementalProcessor Language.Lean.process {
-      opts, mainModuleName
-      fileSetupHandler? := some fun imports => do
-        let result ← 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 #[]
-        return result
-    }
+    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 := {
       chanOut
@@ -260,6 +296,7 @@ section Initialization
     let reporter ← reportSnapshots ctx doc
     return (ctx, {
       doc := { doc with reporter }
+      srcSearchPathTask  := srcSearchPathPromise.result
       pendingRequests    := RBMap.empty
       rpcSessions        := RBMap.empty
       importCachingTask? := none
@@ -452,9 +489,7 @@ section MessageHandling
 
     -- we assume that any other request requires at least the the search path
     -- TODO: move into language-specific request handling
-    let srcSearchPathTask :=
-      st.doc.initSnap.processedSuccessfully.map (·.map (·.srcSearchPath) |>.getD ∅)
-    let t ← IO.bindTask srcSearchPathTask.task fun srcSearchPath => do
+    let t ← IO.bindTask st.srcSearchPathTask fun srcSearchPath => do
       let rc : RequestContext :=
         { rpcSessions := st.rpcSessions
           srcSearchPath

From 88306673bb8da7bf7deee4796a125df0c74ecb78 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 22 Feb 2024 12:46:42 +0100
Subject: [PATCH 50/60] Run `lake setup-file` at most once per worker process,
 like before

In particular, we shouldn't run two instances in parallel. Would be nice
to find a better solution in the future.
---
 src/Lean/Server/FileWorker.lean | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 81bd4fd56253..deaee81b5dc6 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -223,6 +223,16 @@ open Language Lean in
 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, success? := none }
+
   let imports := Elab.headerToImports stx
   let fileSetupResult ← setupFile meta imports fun stderrLine => do
     let progressDiagnostic := {
@@ -250,16 +260,6 @@ def setupImports (meta : DocumentMeta) (chanOut : Channel JsonRpc.Message)
     }
   | _ => pure ()
 
-  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, success? := none }
-
   srcSearchPathPromise.resolve fileSetupResult.srcSearchPath
   return .ok fileSetupResult.fileOptions
 

From 2c99cf5b2cd1b312b03aaf5a0ed172e265b16801 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 29 Feb 2024 18:05:58 +0100
Subject: [PATCH 51/60] Fix cancellation in reportSnapshots

---
 src/Lean/Server/FileWorker.lean       | 64 +++++++++++++++------------
 src/Lean/Server/FileWorker/Utils.lean |  5 ++-
 2 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 175202758b92..8010ee78492f 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -101,6 +101,7 @@ section Elab
     newInfoTrees : Array Elab.InfoTree := #[]
     /-- Whether we should finish with a fatal progress notification. -/
     isFatal := false
+  deriving Inhabited
 
   register_builtin_option server.reportDelayMs : Nat := {
     defValue := 200
@@ -132,32 +133,33 @@ section Elab
     * 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) :
-      BaseIO (Task Unit) := do
+  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.opts).toUInt32
-    BaseIO.bindTask t fun _ =>
-      start
+    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.isFatal 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
-    start := go (toSnapshotTree doc.initSnap) { : ReportSnapshotsState } fun st => do
-      -- callback at the end of reporting
-      if st.isFatal 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 ()
     publishDiagnostics := do
       ctx.chanOut.send <| mkPublishDiagnosticsNotification doc.meta <|
         (← doc.diagnosticsRef.get).map (·.toDiagnostic)
-    go (node : SnapshotTree) (st : ReportSnapshotsState)
-        (cont : ReportSnapshotsState → BaseIO (Task Unit)) : BaseIO (Task Unit) := do
+    go (node : SnapshotTree) (st : ReportSnapshotsState) : BaseIO (Task ReportSnapshotsState) := do
       if (← IO.checkCanceled) then
-        return .pure ()
+        return .pure st
 
       if !node.element.diagnostics.msgLog.isEmpty then
         let diags ←
@@ -183,10 +185,10 @@ section Elab
           newInfoTrees := #[]
         st := { st with newInfoTrees, allInfoTrees := st.allInfoTrees.push itree }
 
-      goSeq st cont node.children.toList
-    goSeq (st : ReportSnapshotsState) (cont : ReportSnapshotsState → BaseIO (Task Unit)) :
-        List (SnapshotTask SnapshotTree) → BaseIO (Task Unit)
-      | [] => cont st
+      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
@@ -194,8 +196,8 @@ section Elab
           if !st.hasBlocked then
             publishDiagnostics
             st := { st with hasBlocked := true }
-        BaseIO.bindTask t.task fun node =>
-          go node st (goSeq · cont ts)
+        BaseIO.bindTask t.task fun t => do
+          BaseIO.bindTask (← go t st) (goSeq · ts)
 end Elab
 
 -- Pending requests are tracked so they can be canceled
@@ -207,6 +209,8 @@ structure AvailableImportsCache where
 
 structure WorkerState where
   doc                : EditableDocument
+  /-- 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
@@ -295,9 +299,11 @@ section Initialization
       meta, initSnap
       diagnosticsRef := (← IO.mkRef ∅)
     }
-    let reporter ← reportSnapshots ctx doc
+    let reporterCancelTk ← CancelToken.new
+    let reporter ← reportSnapshots ctx doc reporterCancelTk
     return (ctx, {
       doc := { doc with reporter }
+      reporterCancelTk
       srcSearchPathTask  := srcSearchPathPromise.result
       pendingRequests    := RBMap.empty
       rpcSessions        := RBMap.empty
@@ -335,14 +341,16 @@ section Updates
 
   /-- Given the new document, updates editable doc state. -/
   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 reporter ← reportSnapshots ctx doc
-    modify fun st => { st with doc := { doc with reporter } }
+    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
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index 84b4f018e6b3..df5e1b6fd8b8 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -38,9 +38,12 @@ def check [MonadExceptOf ElabTaskError m] [MonadLiftT (ST RealWorld) m] [Monad m
   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
 
 -- TEMP: translate from new heterogeneous snapshot tree to old homogeneous async list

From 44e51634c2786d8081a42433584a8420abd1fc0b Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 1 Mar 2024 15:38:31 +0100
Subject: [PATCH 52/60] refactor: lazy initialization of request caches

This avoids an additional worker-language ping-pong of getting the
header environment into the worker as well as the startup delay
---
 src/Lean/Server/Completion.lean               | 53 ++++++++++++-------
 src/Lean/Server/FileWorker.lean               |  2 -
 .../Server/FileWorker/RequestHandling.lean    |  1 -
 src/Lean/Server/Requests.lean                 | 25 ++-------
 4 files changed, 39 insertions(+), 42 deletions(-)

diff --git a/src/Lean/Server/Completion.lean b/src/Lean/Server/Completion.lean
index 6c981c7a2593..b328ac6e5bfd 100644
--- a/src/Lean/Server/Completion.lean
+++ b/src/Lean/Server/Completion.lean
@@ -81,39 +81,53 @@ open Elab
 open Meta
 open FuzzyMatching
 
+abbrev EligibleHeaderDecls := HashMap Name ConstantInfo
+
 /-- Cached header declarations for which `allowCompletion headerEnv decl` is true. -/
-builtin_initialize eligibleHeaderDeclsRef : IO.Ref (HashMap Name ConstantInfo) ← IO.mkRef {}
+builtin_initialize eligibleHeaderDeclsRef : IO.Ref (Option EligibleHeaderDecls) ←
+  IO.mkRef none
 
-/-- 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
+/--
+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
+  if let some eligibleHeaderDecls ← eligibleHeaderDeclsRef.get then
+    return eligibleHeaderDecls
+
+  -- enter critical section to avoid running the computation multiple times
+  -- `take` safety: we call `set` in all branches, without a chance for exceptions in between
+  if let some eligibleHeaderDecls ← unsafe eligibleHeaderDeclsRef.take then
+    -- someone else was quicker!
+    eligibleHeaderDeclsRef.set eligibleHeaderDecls
+    return eligibleHeaderDecls
+  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 headerEnv declName then
+          if allowCompletion env declName then
             eligibleHeaderDecls.insert declName c
           else
             eligibleHeaderDecls
+  eligibleHeaderDeclsRef.set (some eligibleHeaderDecls)
+  return 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 (← getEnv)).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 (env : Environment) (eligibleHeaderDecls : EligibleHeaderDecls)
+    (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 +385,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 env eligibleHeaderDecls 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 env eligibleHeaderDecls resolvedId then
         if let some score := matchAtomic id openedId then
           addUnresolvedCompletionItemForDecl openedId.getString! resolvedId score
     | OpenDecl.simple ns _      =>
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 8010ee78492f..05e929732e74 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -238,8 +238,6 @@ def setupImports (meta : DocumentMeta) (chanOut : Channel JsonRpc.Message)
     return .error { diagnostics := .empty, success? := none }
 
   let imports := Elab.headerToImports stx
-  -- Prepare header-based caches that requests may use
-  --runHeaderCachingHandlers headerEnv
   let fileSetupResult ← setupFile meta imports fun stderrLine => do
     let progressDiagnostic := {
       range      := ⟨⟨0, 0⟩, ⟨0, 0⟩⟩
diff --git a/src/Lean/Server/FileWorker/RequestHandling.lean b/src/Lean/Server/FileWorker/RequestHandling.lean
index cd6348b62304..b1a196f41134 100644
--- a/src/Lean/Server/FileWorker/RequestHandling.lean
+++ b/src/Lean/Server/FileWorker/RequestHandling.lean
@@ -640,7 +640,6 @@ builtin_initialize
     CompletionParams
     CompletionList
     handleCompletion
-    Completion.fillEligibleHeaderDecls
   registerLspRequestHandler
     "completionItem/resolve"
     CompletionItem
diff --git a/src/Lean/Server/Requests.lean b/src/Lean/Server/Requests.lean
index 9e848cc0d109..2cd55959fe2f 100644
--- a/src/Lean/Server/Requests.lean
+++ b/src/Lean/Server/Requests.lean
@@ -181,13 +181,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 {}
@@ -207,9 +202,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
@@ -220,8 +213,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
@@ -252,14 +245,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

From ad6b2251348bc9ffbfe384f30d836fa8845edc7a Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 1 Mar 2024 16:56:49 +0100
Subject: [PATCH 53/60] fix: traces on partial syntax

---
 src/Lean/Elab/Command.lean                       | 3 ++-
 tests/lean/partialSyntaxTraces.lean              | 5 +++++
 tests/lean/partialSyntaxTraces.lean.expected.out | 6 ++++++
 3 files changed, 13 insertions(+), 1 deletion(-)
 create mode 100644 tests/lean/partialSyntaxTraces.lean
 create mode 100644 tests/lean/partialSyntaxTraces.lean.expected.out

diff --git a/src/Lean/Elab/Command.lean b/src/Lean/Elab/Command.lean
index 7f218dc347c2..8e5a595f57b8 100644
--- a/src/Lean/Elab/Command.lean
+++ b/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
diff --git a/tests/lean/partialSyntaxTraces.lean b/tests/lean/partialSyntaxTraces.lean
new file mode 100644
index 000000000000..535493cc5d0f
--- /dev/null
+++ b/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
diff --git a/tests/lean/partialSyntaxTraces.lean.expected.out b/tests/lean/partialSyntaxTraces.lean.expected.out
new file mode 100644
index 000000000000..15ff15a418fd
--- /dev/null
+++ b/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.def "def" (Command.declId `f []) (Command.optDeclSig [] []) (Command.whereStructInst <missing>)))
+[Elab.command] 

From 2745d55762ace02fa26c972638cd7235ea61375d Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Wed, 6 Mar 2024 18:39:30 +0100
Subject: [PATCH 54/60] address some comments

---
 src/Init/Data/Option/Basic.lean       |  1 +
 src/Lean/Data/Lsp/Ipc.lean            |  2 +-
 src/Lean/Elab/Frontend.lean           |  2 +-
 src/Lean/Elab/InfoTree/Main.lean      |  1 +
 src/Lean/Language/Basic.lean          |  5 ++--
 src/Lean/Language/Lean.lean           |  5 ++--
 src/Lean/Message.lean                 |  7 +++++-
 src/Lean/Server/FileWorker.lean       | 35 ++++++++++++++++++---------
 src/Lean/Server/FileWorker/Utils.lean |  5 +---
 src/Lean/Server/README.md             |  2 +-
 10 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/src/Init/Data/Option/Basic.lean b/src/Init/Data/Option/Basic.lean
index c7e2316d5c2b..738d9c50c3a1 100644
--- a/src/Init/Data/Option/Basic.lean
+++ b/src/Init/Data/Option/Basic.lean
@@ -37,6 +37,7 @@ 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)
diff --git a/src/Lean/Data/Lsp/Ipc.lean b/src/Lean/Data/Lsp/Ipc.lean
index e086c2b29f10..b6ffa557a554 100644
--- a/src/Lean/Data/Lsp/Ipc.lean
+++ b/src/Lean/Data/Lsp/Ipc.lean
@@ -89,7 +89,7 @@ Waits for the worker to emit all diagnostic notifications for the current docume
 returns the last notification, if any.
 
 We used to return all notifications but with debouncing in the server, this would not be
-deterministic anymore.
+deterministic anymore as what messages are dropped depends on wall-clock timing.
  -/
 partial def collectDiagnostics (waitForDiagnosticsId : RequestID := 0) (target : DocumentUri) (version : Nat)
 : IpcM (Option (Notification PublishDiagnosticsParams)) := do
diff --git a/src/Lean/Elab/Frontend.lean b/src/Lean/Elab/Frontend.lean
index 16a438551a23..1e5992ae4572 100644
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@@ -151,7 +151,7 @@ def runFrontend
   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.getFinalEnv? snap |>.getD (← mkEmptyEnvironment)
+  let env := Language.Lean.waitForFinalEnv? snap |>.getD (← mkEmptyEnvironment)
   pure (env, !hasErrors)
 
 
diff --git a/src/Lean/Elab/InfoTree/Main.lean b/src/Lean/Elab/InfoTree/Main.lean
index 507a8aba755d..6e85ddf87372 100644
--- a/src/Lean/Elab/InfoTree/Main.lean
+++ b/src/Lean/Elab/InfoTree/Main.lean
@@ -94,6 +94,7 @@ partial def InfoTree.substitute (tree : InfoTree) (assignment : PersistentHashMa
     | none      => hole id
     | some tree => substitute tree assignment
 
+/-- 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
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index e2701d728c24..61a86fea4261 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -208,12 +208,13 @@ instance : MonadLift ProcessingM (ProcessingT IO) where
   monadLift := fun act ctx => act ctx
 
 /--
-Creates snapshot message log from non-interactive message log, caching derived interactive
-diagnostics.
+Creates snapshot message log from non-interactive message log, also allocating a mutable cell
+that can be used by the server to cache interactive diagnostics derived from the log.
 -/
 def Snapshot.Diagnostics.ofMessageLog (msgLog : Lean.MessageLog) :
     ProcessingM Snapshot.Diagnostics := do
   return { msgLog, cacheRef? := 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 {
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 78f032ad9518..369425d766f7 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -8,6 +8,7 @@ recompilation
 Authors: Sebastian Ullrich
 -/
 
+prelude
 import Lean.Language.Basic
 import Lean.Parser.Module
 import Lean.Elab.Command
@@ -322,7 +323,7 @@ where
         return { diagnostics, success? := none }
 
       let headerEnv := headerEnv.setMainModule ctx.mainModuleName
-      let cmdState := Elab.Command.mkState headerEnv msgLog ctx.opts
+      let cmdState := Elab.Command.mkState headerEnv msgLog opts
       let cmdState := { cmdState with infoState := {
         enabled := true
         trees := #[Elab.InfoTree.context (.commandCtx {
@@ -458,7 +459,7 @@ where
       }
 
 /-- Waits for and returns final environment, if importing was successful. -/
-partial def getFinalEnv? (snap : InitialSnapshot) : Option Environment := do
+partial def waitForFinalEnv? (snap : InitialSnapshot) : Option Environment := do
   let snap ← snap.success?
   let snap ← snap.processed.get.success?
   goCmd snap.next.get
diff --git a/src/Lean/Message.lean b/src/Lean/Message.lean
index d51a1d9c5d94..1fcc1cb21af5 100644
--- a/src/Lean/Message.lean
+++ b/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
 
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 05e929732e74..b49dbfe11a7f 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -66,12 +66,16 @@ structure WorkerContext where
   Latest document version received by the client, used for filtering out notifications from
   previous versions.
   -/
-  maxDocVersionRef : IO.Ref Nat
+  maxDocVersionRef : IO.Ref Int
   hLog             : FS.Stream
   initParams       : InitializeParams
   processor        : Parser.InputContext → BaseIO Lean.Language.Lean.InitialSnapshot
   clientHasWidgets : Bool
-  opts             : Options
+  /--
+  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 -/
 
@@ -107,7 +111,9 @@ section Elab
     defValue := 200
     group := "server"
     descr := "(server) time in milliseconds to wait before reporting progress and diagnostics on \
-      document edit in order to reduce flickering"
+      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`."
   }
 
   /--
@@ -136,7 +142,7 @@ section Elab
   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.opts).toUInt32
+      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
@@ -224,6 +230,10 @@ abbrev WorkerM := ReaderT WorkerContext <| StateRefT WorkerState IO
 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
@@ -291,7 +301,7 @@ section Initialization
       processor
       clientHasWidgets
       maxDocVersionRef
-      opts
+      cmdlineOpts := opts
     }
     let doc : EditableDocumentCore := {
       meta, initSnap
@@ -317,13 +327,14 @@ section Initialization
       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 Nat := do
-          let doc ← match msg with
-            | .notification "textDocument/publishDiagnostics" (some <| .obj params) => some params
-            | .notification "$/lean/fileProgress" (some <| .obj params) =>
-              params.find compare "textDocument" |>.bind (·.getObj?.toOption)
-            | _ => none
-          doc.find compare "version" |>.bind (·.getNat?.toOption)
+        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
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index df5e1b6fd8b8..6067430f73c0 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -68,16 +68,13 @@ where go cmdParsed :=
         | some next => .delayed <| next.task.bind go
         | none => .nil)
 
-/-- A map from Diagnostics ID to resulting interactive objects. -/
-abbrev DiagnosticsCache := RBMap Nat (Array Widget.InteractiveDiagnostic) compare
-
 /--
 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
+  meta     : DocumentMeta
   /-- Initial processing snapshot. -/
   initSnap : Language.Lean.InitialSnapshot
   /-- Old representation for backward compatibility. -/
diff --git a/src/Lean/Server/README.md b/src/Lean/Server/README.md
index 34ed71e50ce5..ba552540a761 100644
--- a/src/Lean/Server/README.md
+++ b/src/Lean/Server/README.md
@@ -52,7 +52,7 @@ The interface is shared with the cmdline driver that is used for building Lean f
 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 `Language.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`.
+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 these nodes and report the processing status (diagnostics and progress information) stored therein to the client (`Lean.Server.FileWorker.reportSnapshots`).
 

From 491b125a4c36bf8548b5a44ee038bd11eddefa4b Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 7 Mar 2024 14:12:12 +0100
Subject: [PATCH 55/60] remove ElabTaskError.aborted

---
 .../Server/FileWorker/RequestHandling.lean    |  5 ++---
 src/Lean/Server/FileWorker/Utils.lean         | 19 ++-----------------
 src/Lean/Server/Requests.lean                 | 14 +++++---------
 3 files changed, 9 insertions(+), 29 deletions(-)

diff --git a/src/Lean/Server/FileWorker/RequestHandling.lean b/src/Lean/Server/FileWorker/RequestHandling.lean
index b1a196f41134..959ae9f5a94e 100644
--- a/src/Lean/Server/FileWorker/RequestHandling.lean
+++ b/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
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index 6067430f73c0..6af9f9baf5db 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -15,16 +15,6 @@ namespace Lean.Server.FileWorker
 open Snapshots
 open IO
 
-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
 
@@ -33,11 +23,6 @@ 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) : BaseIO Unit :=
   tk.ref.set true
 
@@ -48,7 +33,7 @@ end CancelToken
 
 -- TEMP: translate from new heterogeneous snapshot tree to old homogeneous async list
 private partial def mkCmdSnaps (initSnap : Language.Lean.InitialSnapshot) :
-    AsyncList ElabTaskError Snapshot := Id.run do
+    AsyncList IO.Error Snapshot := Id.run do
   let some headerParsed := initSnap.success? | return .nil
   .delayed <| headerParsed.processed.task.bind fun headerProcessed => Id.run do
     let some headerSuccess := headerProcessed.success? | return .pure <| .ok .nil
@@ -78,7 +63,7 @@ structure EditableDocumentCore where
   /-- Initial processing snapshot. -/
   initSnap : Language.Lean.InitialSnapshot
   /-- Old representation for backward compatibility. -/
-  cmdSnaps : AsyncList ElabTaskError Snapshot := mkCmdSnaps initSnap
+  cmdSnaps : AsyncList IO.Error Snapshot := mkCmdSnaps initSnap
   /--
   Interactive versions of diagnostics reported so far. Filled by `reportSnapshots` and read by
   `handleGetInteractiveDiagnosticsRequest`.
diff --git a/src/Lean/Server/Requests.lean b/src/Lean/Server/Requests.lean
index 2cd55959fe2f..906f7f1cc358 100644
--- a/src/Lean/Server/Requests.lean
+++ b/src/Lean/Server/Requests.lean
@@ -105,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
 
@@ -122,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. -/

From c64fe35950758a629926254d679dd73d975e9ee0 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 7 Mar 2024 14:12:23 +0100
Subject: [PATCH 56/60] address more comments

---
 src/Lean/Language/Basic.lean              | 40 ++++-------------------
 src/Lean/Server/FileWorker.lean           | 22 ++++++-------
 src/Lean/Server/FileWorker/SetupFile.lean | 12 ++++++-
 3 files changed, 29 insertions(+), 45 deletions(-)

diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index 61a86fea4261..5e0ce5d4c618 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -13,53 +13,27 @@ import Lean.Parser.Types
 
 set_option linter.missingDocs true
 
--- early declarations for `fileSetupHandler?` below; as the field is used only by the worker, we
--- leave them in that namespace
-namespace Lean.Server.FileWorker
-
-/-- 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
-
-end Lean.Server.FileWorker
-
 namespace Lean.Language
 
-/-- `MessageLog` with caching of interactive diagnostics. -/
+/-- `MessageLog` with interactive diagnostics. -/
 structure Snapshot.Diagnostics where
   /-- Non-interactive message log. -/
   msgLog : MessageLog
   /--
-  Dynamic mutable slot usable by the language server for caching interactive diagnostics.  If
-  `none`, no caching is done, which should only be used for messages not containing any interactive
-  elements.
+  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.
   -/
-  cacheRef? : Option (IO.Ref (Option Dynamic))
+  interactiveDiagsRef? : Option (IO.Ref (Option Dynamic))
 deriving Inhabited
 
 /-- The empty set of diagnostics. -/
 def Snapshot.Diagnostics.empty : Snapshot.Diagnostics where
   msgLog := .empty
   -- nothing to cache
-  cacheRef? := none
+  interactiveDiagsRef? := none
 
 /--
   The base class of all snapshots: all the generic information the language server needs about a
@@ -213,7 +187,7 @@ that can be used by the server to cache interactive diagnostics derived from the
 -/
 def Snapshot.Diagnostics.ofMessageLog (msgLog : Lean.MessageLog) :
     ProcessingM Snapshot.Diagnostics := do
-  return { msgLog, cacheRef? := some (← IO.mkRef none) }
+  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
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index b49dbfe11a7f..db90f0f1906b 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -103,8 +103,8 @@ section Elab
     allInfoTrees : Array Elab.InfoTree := #[]
     /-- New info trees encountered since we last sent a .ilean update notification. -/
     newInfoTrees : Array Elab.InfoTree := #[]
-    /-- Whether we should finish with a fatal progress notification. -/
-    isFatal := false
+    /-- Whether we encountered any snapshot with `Snapshot.isFatal`. -/
+    hasFatal := false
   deriving Inhabited
 
   register_builtin_option server.reportDelayMs : Nat := {
@@ -117,11 +117,11 @@ This option can only be set on the command line, not in the lakefile or via `set
   }
 
   /--
-  Type of cache stored in `Snapshot.Diagnostics.cacheRef?`.
+  Type of state stored in `Snapshot.Diagnostics.cacheRef?`.
 
   See also section "Communication" in Lean/Server/README.md.
   -/
-  structure CachedInteractiveDiagnostics where
+  structure MemorizedInteractiveDiagnostics where
     diags : Array Widget.InteractiveDiagnostic
   deriving TypeName
 
@@ -149,7 +149,7 @@ This option can only be set on the command line, not in the lakefile or via `set
           return .pure ()
 
         -- callback at the end of reporting
-        if st.isFatal then
+        if st.hasFatal then
           ctx.chanOut.send <| mkFileProgressAtPosNotification doc.meta 0 .fatalError
         else
           ctx.chanOut.send <| mkFileProgressDoneNotification doc.meta
@@ -169,20 +169,20 @@ This option can only be set on the command line, not in the lakefile or via `set
 
       if !node.element.diagnostics.msgLog.isEmpty then
         let diags ←
-          if let some cached ← node.element.diagnostics.cacheRef?.bindM fun cacheRef => do
-              return (← cacheRef.get).bind (·.get? CachedInteractiveDiagnostics) then
-            pure cached.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.toList.toArray.mapM
               (Widget.msgToInteractiveDiagnostic doc.meta.text · ctx.clientHasWidgets)
-            if let some cacheRef := node.element.diagnostics.cacheRef? then
-              cacheRef.set <| some <| .mk { diags : CachedInteractiveDiagnostics }
+            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 isFatal := st.isFatal || node.element.isFatal }
+      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
diff --git a/src/Lean/Server/FileWorker/SetupFile.lean b/src/Lean/Server/FileWorker/SetupFile.lean
index 09fe62f96aeb..0e390835ffc2 100644
--- a/src/Lean/Server/FileWorker/SetupFile.lean
+++ b/src/Lean/Server/FileWorker/SetupFile.lean
@@ -9,7 +9,6 @@ import Lean.Server.Utils
 import Lean.Util.FileSetupInfo
 import Lean.Util.LakePath
 import Lean.LoadDynlib
-import Lean.Language.Basic
 
 namespace Lean.Server.FileWorker
 
@@ -53,6 +52,17 @@ partial def runLakeSetupFile
   let exitCode ← lakeProc.wait
   return ⟨spawnArgs, exitCode, stdout, stderr⟩
 
+inductive FileSetupResultKind where
+  | success
+  | noLakefile
+  | importsOutOfDate
+  | error (msg : String)
+
+structure FileSetupResult where
+  kind          : FileSetupResultKind
+  srcSearchPath : SearchPath
+  fileOptions   : Options
+
 def FileSetupResult.ofSuccess (pkgSearchPath : SearchPath) (fileOptions : Options)
     : IO FileSetupResult := do return {
   kind          := FileSetupResultKind.success

From 34ff994b0e13a32d8af88c9817829cb6e58cd2dc Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 7 Mar 2024 14:58:22 +0100
Subject: [PATCH 57/60] address more comments

---
 src/Lean/Language/Basic.lean | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index 5e0ce5d4c618..cc2e232f30eb 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -8,6 +8,7 @@ driver. See the [server readme](../Server/README.md#worker-architecture) for an
 Authors: Sebastian Ullrich
 -/
 
+prelude
 import Lean.Message
 import Lean.Parser.Types
 
@@ -32,7 +33,7 @@ deriving Inhabited
 /-- The empty set of diagnostics. -/
 def Snapshot.Diagnostics.empty : Snapshot.Diagnostics where
   msgLog := .empty
-  -- nothing to cache
+  -- nothing to memorize
   interactiveDiagsRef? := none
 
 /--
@@ -183,7 +184,7 @@ instance : MonadLift ProcessingM (ProcessingT IO) where
 
 /--
 Creates snapshot message log from non-interactive message log, also allocating a mutable cell
-that can be used by the server to cache interactive diagnostics derived from the log.
+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

From 27c90ba40b70b69f833be3b52e1c93f5331b79ae Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Thu, 7 Mar 2024 16:45:14 +0100
Subject: [PATCH 58/60] address even more comments

---
 src/Lean/DeclarationRange.lean        |   4 +-
 src/Lean/Language/Basic.lean          |   2 +-
 src/Lean/Language/Lean.lean           | 111 +++++++++++++-------------
 src/Lean/Server/FileWorker.lean       |  14 +++-
 src/Lean/Server/FileWorker/Utils.lean |  12 +--
 src/Lean/Server/README.md             |   2 +-
 src/Lean/Server/References.lean       |   6 +-
 src/Lean/Server/Utils.lean            |   4 +
 8 files changed, 84 insertions(+), 71 deletions(-)

diff --git a/src/Lean/DeclarationRange.lean b/src/Lean/DeclarationRange.lean
index ebd2e3f7dbfa..551614398bf6 100644
--- a/src/Lean/DeclarationRange.lean
+++ b/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
diff --git a/src/Lean/Language/Basic.lean b/src/Lean/Language/Basic.lean
index cc2e232f30eb..326b50089860 100644
--- a/src/Lean/Language/Basic.lean
+++ b/src/Lean/Language/Basic.lean
@@ -64,7 +64,7 @@ structure SnapshotTask (α : Type) where
   task : Task α
 deriving Nonempty
 
-/-- Creates a snapshot task from a range and a `BaseIO` action. -/
+/-- 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
diff --git a/src/Lean/Language/Lean.lean b/src/Lean/Language/Lean.lean
index 369425d766f7..82fa7349a81c 100644
--- a/src/Lean/Language/Lean.lean
+++ b/src/Lean/Language/Lean.lean
@@ -98,10 +98,10 @@ instance : ToSnapshotTree CommandFinishedSnapshot where
 -- TODO: tactics
 structure CommandSignatureProcessedSnapshot extends Snapshot where
   /-- State after processing is finished. -/
-  finished : SnapshotTask CommandFinishedSnapshot
+  finishedSnap : SnapshotTask CommandFinishedSnapshot
 deriving Nonempty
 instance : ToSnapshotTree CommandSignatureProcessedSnapshot where
-  toSnapshotTree s := ⟨s.toSnapshot, #[s.finished.map (sync := true) toSnapshotTree]⟩
+  toSnapshotTree s := ⟨s.toSnapshot, #[s.finishedSnap.map (sync := true) toSnapshotTree]⟩
 
 /-- State after a command has been parsed. -/
 structure CommandParsedSnapshotData extends Snapshot where
@@ -110,13 +110,14 @@ structure CommandParsedSnapshotData extends Snapshot where
   /-- Resulting parser state. -/
   parserState : Parser.ModuleParserState
   /-- Signature processing task. -/
-  sig : SnapshotTask CommandSignatureProcessedSnapshot
+  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) (next? : Option (SnapshotTask CommandParsedSnapshot))
+  | mk (data : CommandParsedSnapshotData)
+    (nextCmdSnap? : Option (SnapshotTask CommandParsedSnapshot))
 deriving Nonempty
 /-- The snapshot data. -/
 abbrev CommandParsedSnapshot.data : CommandParsedSnapshot → CommandParsedSnapshotData
@@ -129,7 +130,7 @@ abbrev CommandParsedSnapshot.next? : CommandParsedSnapshot →
 partial instance : ToSnapshotTree CommandParsedSnapshot where
   toSnapshotTree := go where
     go s := ⟨s.data.toSnapshot,
-      #[s.data.sig.map (sync := true) toSnapshotTree] |>
+      #[s.data.sigSnap.map (sync := true) toSnapshotTree] |>
         pushOpt (s.next?.map (·.map (sync := true) go))⟩
 
 /-- Cancels all significant computations from this snapshot onwards. -/
@@ -137,33 +138,33 @@ partial def CommandParsedSnapshot.cancel (snap : CommandParsedSnapshot) : BaseIO
   -- This is the only relevant computation right now
   -- TODO: cancel additional elaboration tasks if we add them without switching to implicit
   -- cancellation
-  snap.data.sig.cancel
+  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 HeaderProcessedSuccessfully where
+structure HeaderProcessedState where
   /-- The resulting initial elaboration state. -/
   cmdState : Command.State
   /-- First command task (there is always at least a terminal command). -/
-  next : SnapshotTask CommandParsedSnapshot
+  firstCmdSnap : SnapshotTask CommandParsedSnapshot
 
 /-- State after the module header has been processed including imports. -/
 structure HeaderProcessedSnapshot extends Snapshot where
   /-- State after successful importing. -/
-  success? : Option HeaderProcessedSuccessfully
-  isFatal := success?.isNone
+  result? : Option HeaderProcessedState
+  isFatal := result?.isNone
 instance : ToSnapshotTree HeaderProcessedSnapshot where
   toSnapshotTree s := ⟨s.toSnapshot, #[] |>
-    pushOpt (s.success?.map (·.next.map (sync := true) toSnapshotTree))⟩
+    pushOpt (s.result?.map (·.firstCmdSnap.map (sync := true) toSnapshotTree))⟩
 
 /-- State after successfully parsing the module header. -/
-structure HeaderParsedSucessfully where
+structure HeaderParsedState where
   /-- Resulting parser state. -/
   parserState : Parser.ModuleParserState
   /-- Header processing task. -/
-  processed : SnapshotTask HeaderProcessedSnapshot
+  processedSnap : SnapshotTask HeaderProcessedSnapshot
 
 /-- State after the module header has been parsed. -/
 structure HeaderParsedSnapshot extends Snapshot where
@@ -172,17 +173,17 @@ structure HeaderParsedSnapshot extends Snapshot where
   /-- Resulting syntax tree. -/
   stx : Syntax
   /-- State after successful parsing. -/
-  success? : Option HeaderParsedSucessfully
-  isFatal := success?.isNone
+  result? : Option HeaderParsedState
+  isFatal := result?.isNone
 
 instance : ToSnapshotTree HeaderParsedSnapshot where
   toSnapshotTree s := ⟨s.toSnapshot,
-    #[] |> pushOpt (s.success?.map (·.processed.map (sync := true) toSnapshotTree))⟩
+    #[] |> pushOpt (s.result?.map (·.processedSnap.map (sync := true) toSnapshotTree))⟩
 
 /-- Shortcut accessor to the final header state, if successful. -/
-def HeaderParsedSnapshot.processedSuccessfully (snap : HeaderParsedSnapshot) :
-    SnapshotTask (Option HeaderProcessedSuccessfully) :=
-  snap.success?.bind (·.processed.map (sync := true) (·.success?)) |>.getD (.pure none)
+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
@@ -247,36 +248,36 @@ where
     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.success? then
-        return { old with ictx, success? := some { oldSuccess with
-          processed := (← oldSuccess.processed.bindIO (sync := true) fun oldProcessed => do
-            if let some oldProcSuccess := oldProcessed.success? then
+      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.next.bindIO (sync := true) fun oldCmd =>
-                return .pure { oldProcessed with success? := some { oldProcSuccess with
-                  next := (← parseCmd oldCmd oldSuccess.parserState oldProcSuccess.cmdState ctx) } }
+              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.processedSuccessfully.get? 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.next.get? then
+        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, success? := none }) do
+    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)
-          success? := none
+          result? := none
         }
 
       -- semi-fast path: go to next snapshot if syntax tree is unchanged AND we're still in front
@@ -289,18 +290,18 @@ where
         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.success? then
-          oldSuccess.processed.cancel
-          let _ ← BaseIO.mapTask (t := oldSuccess.processed.task) fun processed => do
-            if let some oldProcSuccess := processed.success? then
-              let _ ← BaseIO.mapTask (·.cancel) oldProcSuccess.next.task
+        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)
-        success? := some {
+        result? := some {
           parserState
-          processed := (← processHeader stx parserState)
+          processedSnap := (← processHeader stx parserState)
         }
       }
 
@@ -309,7 +310,7 @@ where
     let ctx ← read
     SnapshotTask.ofIO ⟨0, ctx.input.endPos⟩ <|
     ReaderT.run (r := ctx) <|  -- re-enter reader in new task
-    withHeaderExceptions (α := HeaderProcessedSnapshot) ({ · with success? := none }) do
+    withHeaderExceptions (α := HeaderProcessedSnapshot) ({ · with result? := none }) do
       let opts ← match (← setupImports stx) with
         | .ok opts => pure opts
         | .error snap => return snap
@@ -320,7 +321,7 @@ where
         ctx.toInputContext ctx.trustLevel
       let diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog)
       if msgLog.hasErrors then
-        return { diagnostics, success? := none }
+        return { diagnostics, result? := none }
 
       let headerEnv := headerEnv.setMainModule ctx.mainModuleName
       let cmdState := Elab.Command.mkState headerEnv msgLog opts
@@ -343,9 +344,9 @@ where
       return {
         diagnostics
         infoTree? := cmdState.infoState.trees[0]!
-        success? := some {
+        result? := some {
           cmdState
-          next := (← parseCmd none parserState cmdState)
+          firstCmdSnap := (← parseCmd none parserState cmdState)
         }
       }
 
@@ -359,18 +360,18 @@ where
       -- 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 (next? := none) {
+      return .pure <| .mk (nextCmdSnap? := none) {
         diagnostics := .empty, stx := .missing, parserState
-        sig := .pure {
+        sigSnap := .pure {
           diagnostics := .empty
-          finished := .pure { diagnostics := .empty, cmdState } } }
+          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)
-          (next? := (← old.data.sig.bindIO (sync := true) fun oldSig =>
-            oldSig.finished.bindIO (sync := true) fun oldFinished =>
+          (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
@@ -400,16 +401,16 @@ where
         -- on first change, make sure to cancel all further old tasks
         old.cancel
 
-      let sig ← processCmdSignature stx cmdState msgLog.hasErrors beginPos ctx
+      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 <$> (sig.bind (·.finished)).bindIO fun finished =>
+        else some <$> (sigSnap.bind (·.finishedSnap)).bindIO fun finished =>
           parseCmd none parserState finished.cmdState ctx
-      return .mk (next? := next?) {
+      return .mk (nextCmdSnap? := next?) {
         diagnostics := (← Snapshot.Diagnostics.ofMessageLog msgLog ctx.toProcessingContext)
         stx
         parserState
-        sig
+        sigSnap
       }
 
   processCmdSignature (stx : Syntax) (cmdState : Command.State) (hasParseError : Bool)
@@ -450,7 +451,7 @@ where
       let cmdState := { cmdState with messages }
       return {
         diagnostics := .empty
-        finished := .pure {
+        finishedSnap := .pure {
           diagnostics :=
             (← Snapshot.Diagnostics.ofMessageLog cmdState.messages ctx.toProcessingContext)
           infoTree? := some cmdState.infoState.trees[0]!
@@ -460,13 +461,13 @@ where
 
 /-- Waits for and returns final environment, if importing was successful. -/
 partial def waitForFinalEnv? (snap : InitialSnapshot) : Option Environment := do
-  let snap ← snap.success?
-  let snap ← snap.processed.get.success?
-  goCmd snap.next.get
+  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.sig.get.finished.get.cmdState.env
+    snap.data.sigSnap.get.finishedSnap.get.cmdState.env
 
 end Lean
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index db90f0f1906b..c09fb727e794 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/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
@@ -192,6 +193,7 @@ This option can only be set on the command line, not in the lakefile or via `set
         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
@@ -245,7 +247,7 @@ def setupImports (meta : DocumentMeta) (chanOut : Channel JsonRpc.Message)
     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, success? := none }
+    return .error { diagnostics := .empty, result? := none }
 
   let imports := Elab.headerToImports stx
   let fileSetupResult ← setupFile meta imports fun stderrLine => do
@@ -265,12 +267,12 @@ def setupImports (meta : DocumentMeta) (chanOut : Channel JsonRpc.Message)
       diagnostics := (← Language.diagnosticsOfHeaderError
         "Imports are out of date and must be rebuilt; \
           use the \"Restart File\" command in your editor.")
-      success? := none
+      result? := none
     }
   | .error msg =>
     return .error {
       diagnostics := (← diagnosticsOfHeaderError msg)
-      success? := none
+      result? := none
     }
   | _ => pure ()
 
@@ -435,7 +437,13 @@ section MessageHandling
   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 =>
       params.lineRange?.all fun ⟨s, e⟩ =>
         -- does [s,e) intersect [diag.fullRange.start.line,diag.fullRange.end.line)?
diff --git a/src/Lean/Server/FileWorker/Utils.lean b/src/Lean/Server/FileWorker/Utils.lean
index 6af9f9baf5db..13c89720dd05 100644
--- a/src/Lean/Server/FileWorker/Utils.lean
+++ b/src/Lean/Server/FileWorker/Utils.lean
@@ -34,17 +34,17 @@ end 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.success? | return .nil
-  .delayed <| headerParsed.processed.task.bind fun headerProcessed => Id.run do
-    let some headerSuccess := headerProcessed.success? | return .pure <| .ok .nil
+  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.next.task.bind go
+    } <| .delayed <| headerSuccess.firstCmdSnap.task.bind go
 where go cmdParsed :=
-  cmdParsed.data.sig.task.bind fun sig =>
-    sig.finished.task.map fun finished =>
+  cmdParsed.data.sigSnap.task.bind fun sig =>
+    sig.finishedSnap.task.map fun finished =>
       .ok <| .cons {
         stx := cmdParsed.data.stx
         mpState := cmdParsed.data.parserState
diff --git a/src/Lean/Server/README.md b/src/Lean/Server/README.md
index ba552540a761..4bac3026e17b 100644
--- a/src/Lean/Server/README.md
+++ b/src/Lean/Server/README.md
@@ -54,7 +54,7 @@ In languages with less strict ordering and less syntax extensibility, there may
 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 these nodes and report the processing status (diagnostics and progress information) stored therein to the client (`Lean.Server.FileWorker.reportSnapshots`).
+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)
 
diff --git a/src/Lean/Server/References.lean b/src/Lean/Server/References.lean
index 22a29d6c3a21..721f55e03e2f 100644
--- a/src/Lean/Server/References.lean
+++ b/src/Lean/Server/References.lean
@@ -42,9 +42,9 @@ def toLspRefInfo (i : RefInfo) : BaseIO Lsp.RefInfo := do
   let refToRefInfoLocation (ref : Reference) : BaseIO RefInfo.Location := do
     let parentDeclName? := ref.ci.parentDecl?
     let .ok parentDeclRanges? ← EIO.toBaseIO <| ref.ci.runCoreM do
-      let some parentDeclName := parentDeclName?
-        | return none
-      findDeclarationRanges? parentDeclName
+        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 {
diff --git a/src/Lean/Server/Utils.lean b/src/Lean/Server/Utils.lean
index 639403c860b2..62ea33930be0 100644
--- a/src/Lean/Server/Utils.lean
+++ b/src/Lean/Server/Utils.lean
@@ -114,6 +114,7 @@ def applyDocumentChange (oldText : FileMap) : (change : Lsp.TextDocumentContentC
 def foldDocumentChanges (changes : Array Lsp.TextDocumentContentChangeEvent) (oldText : FileMap) : FileMap :=
   changes.foldl applyDocumentChange oldText
 
+/-- Constructs a `textDocument/publishDiagnostics` notification. -/
 def mkPublishDiagnosticsNotification (m : DocumentMeta) (diagnostics : Array Lsp.Diagnostic) :
     JsonRpc.Notification Lsp.PublishDiagnosticsParams where
   method := "textDocument/publishDiagnostics"
@@ -123,6 +124,7 @@ def mkPublishDiagnosticsNotification (m : DocumentMeta) (diagnostics : Array Lsp
     diagnostics := diagnostics
   }
 
+/-- Constructs a `$/lean/fileProgress` notification. -/
 def mkFileProgressNotification (m : DocumentMeta) (processing : Array LeanFileProgressProcessingInfo) :
     JsonRpc.Notification Lsp.LeanFileProgressParams where
   method := "$/lean/fileProgress"
@@ -131,11 +133,13 @@ def mkFileProgressNotification (m : DocumentMeta) (processing : Array LeanFilePr
     processing
   }
 
+/-- 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 }]
 
+/-- Constructs a `$/lean/fileProgress` notification marking processing as done. -/
 def mkFileProgressDoneNotification (m : DocumentMeta) : JsonRpc.Notification Lsp.LeanFileProgressParams :=
   mkFileProgressNotification m #[]
 

From f7db58928dbf95695a806fd8666d08e34704a5ba Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 8 Mar 2024 10:38:33 +0100
Subject: [PATCH 59/60] Replace silly critical section code with `modifyGet`

---
 src/Lean/Server/Completion.lean | 33 +++++++++++++--------------------
 1 file changed, 13 insertions(+), 20 deletions(-)

diff --git a/src/Lean/Server/Completion.lean b/src/Lean/Server/Completion.lean
index b328ac6e5bfd..95e4daeda4be 100644
--- a/src/Lean/Server/Completion.lean
+++ b/src/Lean/Server/Completion.lean
@@ -92,26 +92,19 @@ Returns the declarations in the header for which `allowCompletion env decl` is t
 if not already cached.
 -/
 def getEligibleHeaderDecls (env : Environment) : IO EligibleHeaderDecls := do
-  if let some eligibleHeaderDecls ← eligibleHeaderDeclsRef.get then
-    return eligibleHeaderDecls
-
-  -- enter critical section to avoid running the computation multiple times
-  -- `take` safety: we call `set` in all branches, without a chance for exceptions in between
-  if let some eligibleHeaderDecls ← unsafe eligibleHeaderDeclsRef.take then
-    -- someone else was quicker!
-    eligibleHeaderDeclsRef.set eligibleHeaderDecls
-    return eligibleHeaderDecls
-  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
-  eligibleHeaderDeclsRef.set (some eligibleHeaderDecls)
-  return eligibleHeaderDecls
+  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]

From c26033c6acd13c37387aa79e9a5db3a1e37ddcc1 Mon Sep 17 00:00:00 2001
From: Sebastian Ullrich <sebasti@nullri.ch>
Date: Fri, 8 Mar 2024 10:57:45 +0100
Subject: [PATCH 60/60] Address review comments

---
 src/Lean/Server/Completion.lean           | 8 ++++----
 src/Lean/Server/FileWorker.lean           | 5 +----
 src/Lean/Server/FileWorker/SetupFile.lean | 9 +++++++++
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/Lean/Server/Completion.lean b/src/Lean/Server/Completion.lean
index 95e4daeda4be..7e3480ebfc9f 100644
--- a/src/Lean/Server/Completion.lean
+++ b/src/Lean/Server/Completion.lean
@@ -110,14 +110,14 @@ def getEligibleHeaderDecls (env : Environment) : IO EligibleHeaderDecls := do
 private def forEligibleDeclsM [Monad m] [MonadEnv m] [MonadLiftT (ST IO.RealWorld) m]
     [MonadLiftT IO m] (f : Name → ConstantInfo → m PUnit) : m PUnit := do
   let env ← getEnv
-  (← getEligibleHeaderDecls (← getEnv)).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 (env : Environment) (eligibleHeaderDecls : EligibleHeaderDecls)
+private def allowCompletion (eligibleHeaderDecls : EligibleHeaderDecls) (env : Environment)
     (declName : Name) : Bool :=
   eligibleHeaderDecls.contains declName ||
     env.constants.map₂.contains declName && Lean.Meta.allowCompletion env declName
@@ -382,13 +382,13 @@ private def idCompletionCore
   -- Auxiliary function for `alias`
   let addAlias (alias : Name) (declNames : List Name) (score : Float) : M Unit := do
     declNames.forM fun declName => do
-      if allowCompletion env eligibleHeaderDecls 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 env eligibleHeaderDecls resolvedId then
+      if allowCompletion eligibleHeaderDecls env resolvedId then
         if let some score := matchAtomic id openedId then
           addUnresolvedCompletionItemForDecl openedId.getString! resolvedId score
     | OpenDecl.simple ns _      =>
diff --git a/src/Lean/Server/FileWorker.lean b/src/Lean/Server/FileWorker.lean
index 2d5913d10687..685b62c6a102 100644
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@@ -165,16 +165,13 @@ This option can only be set on the command line, not in the lakefile or via `set
       ctx.chanOut.send <| mkPublishDiagnosticsNotification doc.meta <|
         (← doc.diagnosticsRef.get).map (·.toDiagnostic)
     go (node : SnapshotTree) (st : ReportSnapshotsState) : BaseIO (Task ReportSnapshotsState) := do
-      if (← IO.checkCanceled) then
-        return .pure st
-
       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.toList.toArray.mapM
+            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 }
diff --git a/src/Lean/Server/FileWorker/SetupFile.lean b/src/Lean/Server/FileWorker/SetupFile.lean
index 0e390835ffc2..4c972c5711fe 100644
--- a/src/Lean/Server/FileWorker/SetupFile.lean
+++ b/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)