feat: reduce server memory consumption (#11162)

This PR reduces the memory consumption of the language server (the watchdog process in particular). In Mathlib, it reduces memory consumption by about 1GB. It also fixes two bugs in the call hierarchy: - When an open file had import errors (e.g. from a transitive build failure), the call hierarchy would not display any usages in that file. Now we use the reference information from the .ilean instead. - When a command would not set a parent declaration (e.g. `#check`), the result was filtered from the call hierarchy. Now we display it as `[anonymous]` instead.
2025-12-01 11:53:23 +01:00 · 2025-12-01 11:53:23 +01:00 · af5b47295f
commit af5b47295f
parent 3282ac6f96
14 changed files with 385 additions and 135 deletions
--- a/src/Lean/Data/Lsp/Extra.lean
+++ b/src/Lean/Data/Lsp/Extra.lean
@ -63,8 +63,8 @@ an ILean finalization notification for the worker and the document version desig
 Used for test stability in tests that use the .ileans.
 -/
 structure WaitForILeansParams where
-  uri     : DocumentUri
-  version : Nat
+  uri?     : Option DocumentUri := none
+  version? : Option Nat := none
  deriving FromJson, ToJson

 structure WaitForILeans where
--- a/src/Lean/Data/Lsp/Internal.lean
+++ b/src/Lean/Data/Lsp/Internal.lean
@ -10,6 +10,7 @@ prelude
 public import Lean.Expr
 public import Lean.Data.Lsp.Basic
 public import Lean.Data.JsonRpc
+public import Lean.Data.DeclarationRange

 public section

@ -98,27 +99,129 @@ instance : ToJson RefIdent where

 end RefIdent

-/-- Information about the declaration surrounding a reference. -/
-structure RefInfo.ParentDecl where
-  /-- Name of the declaration surrounding a reference. -/
-  name           : String
-  /-- Range of the declaration surrounding a reference. -/
-  range          : Lsp.Range
-  /-- Selection range of the declaration surrounding a reference. -/
-  selectionRange : Lsp.Range
-  deriving ToJson
+/-- Position information for a declaration. Inlined to reduce memory consumption. -/
+structure DeclInfo where
+  /-- Start line of range. -/
+  rangeStartPosLine : Nat
+  /-- Start character of range. -/
+  rangeStartPosCharacter : Nat
+  /-- End line of range. -/
+  rangeEndPosLine : Nat
+  /-- End character of range. -/
+  rangeEndPosCharacter : Nat
+  /-- Start line of selection range. -/
+  selectionRangeStartPosLine : Nat
+  /-- Start character of selection range. -/
+  selectionRangeStartPosCharacter : Nat
+  /-- End line of selection range. -/
+  selectionRangeEndPosLine : Nat
+  /-- End character of selection range. -/
+  selectionRangeEndPosCharacter : Nat
+
+/-- Converts a set of `DeclarationRanges` to a `DeclInfo`. -/
+def DeclInfo.ofDeclarationRanges (r : DeclarationRanges) : DeclInfo where
+  rangeStartPosLine := r.range.pos.line - 1
+  rangeStartPosCharacter := r.range.charUtf16
+  rangeEndPosLine := r.range.endPos.line - 1
+  rangeEndPosCharacter := r.range.endCharUtf16
+  selectionRangeStartPosLine := r.selectionRange.pos.line - 1
+  selectionRangeStartPosCharacter := r.selectionRange.charUtf16
+  selectionRangeEndPosLine := r.selectionRange.endPos.line - 1
+  selectionRangeEndPosCharacter := r.selectionRange.endCharUtf16
+
+/-- Range of this parent decl. -/
+def DeclInfo.range (i : DeclInfo) : Lsp.Range :=
+  ⟨⟨i.rangeStartPosLine, i.rangeStartPosCharacter⟩, ⟨i.rangeEndPosLine, i.rangeEndPosCharacter⟩⟩
+
+/-- Selection range of this parent decl. -/
+def DeclInfo.selectionRange (i : DeclInfo) : Lsp.Range :=
+  ⟨⟨i.selectionRangeStartPosLine, i.selectionRangeStartPosCharacter⟩,
+    ⟨i.selectionRangeEndPosLine, i.selectionRangeEndPosCharacter⟩⟩
+
+instance : ToJson DeclInfo where
+  toJson i :=
+    Json.arr #[
+      i.rangeStartPosLine,
+      i.rangeStartPosCharacter,
+      i.rangeEndPosLine,
+      i.rangeEndPosCharacter,
+      i.selectionRangeStartPosLine,
+      i.selectionRangeStartPosCharacter,
+      i.selectionRangeEndPosLine,
+      i.selectionRangeEndPosCharacter
+    ]
+
+instance : FromJson DeclInfo where
+  fromJson?
+    | .arr xs => do
+      if xs.size != 8 then
+        throw s!"Expected list of length 8, not length {xs.size}"
+      return {
+        rangeStartPosLine := ← fromJson? xs[0]!
+        rangeStartPosCharacter := ← fromJson? xs[1]!
+        rangeEndPosLine := ← fromJson? xs[2]!
+        rangeEndPosCharacter := ← fromJson? xs[3]!
+        selectionRangeStartPosLine := ← fromJson? xs[4]!
+        selectionRangeStartPosCharacter := ← fromJson? xs[5]!
+        selectionRangeEndPosLine := ← fromJson? xs[6]!
+        selectionRangeEndPosCharacter := ← fromJson? xs[7]!
+      }
+    | _ => throw "Expected list"
+
+/-- Declarations of a file with associated position information. -/
+@[expose] def Decls := Std.TreeMap String DeclInfo
+  deriving EmptyCollection, ForIn Id
+
+instance : ToJson Decls where
+  toJson m := Json.mkObj <| m.toList.map fun (declName, info) => (declName, toJson info)
+
+instance : FromJson Decls where
+  fromJson? j := do
+    let node ← j.getObj?
+    node.foldlM (init := ∅) fun m k v =>
+      return m.insert k (← fromJson? v)

 /--
 Denotes the range of a reference, as well as the parent declaration of the reference.
 If the reference is itself a declaration, then it contains no parent declaration.
+The position information is inlined to reduce memory consumption.
 -/
 structure RefInfo.Location where
-  /-- Range of the reference. -/
-  range       : Lsp.Range
-  /-- Parent declaration of the reference. `none` if the reference is itself a declaration. -/
-  parentDecl? : Option RefInfo.ParentDecl
+  mk' ::
+  /-- Start line of the range of this location. -/
+  startPosLine : Nat
+  /-- Start character of the range of this location. -/
+  startPosCharacter : Nat
+  /-- End line of the range of this location. -/
+  endPosLine : Nat
+  /-- End character of the range of this location. -/
+  endPosCharacter : Nat
+  /--
+  Parent declaration of the reference. Empty string if the reference is itself a declaration.
+  We do not use `Option` for memory consumption reasons.
+  -/
+  parentDecl : String
 deriving Inhabited

+/-- Creates a `RefInfo.Location`. -/
+def RefInfo.Location.mk (range : Lsp.Range) (parentDecl? : Option String) : RefInfo.Location where
+  startPosLine := range.start.line
+  startPosCharacter := range.start.character
+  endPosLine := range.end.line
+  endPosCharacter := range.end.character
+  parentDecl := parentDecl?.getD ""
+
+/-- Range of this location. -/
+def RefInfo.Location.range (l : RefInfo.Location) : Lsp.Range :=
+  ⟨⟨l.startPosLine, l.startPosCharacter⟩, ⟨l.endPosLine, l.endPosCharacter⟩⟩
+
+/-- Name of the parent declaration of this location. -/
+def RefInfo.Location.parentDecl? (l : RefInfo.Location) : Option String :=
+  if l.parentDecl.isEmpty then
+    none
+  else
+    some l.parentDecl
+
 /-- Definition site and usage sites of a reference. Obtained from `Lean.Server.RefInfo`. -/
 structure RefInfo where
  /-- Definition site of the reference. May be `none` when we cannot find a definition site. -/
@ -128,16 +231,9 @@ structure RefInfo where

 instance : ToJson RefInfo where
  toJson i :=
-    let rangeToList (r : Lsp.Range) : List Nat :=
-      [r.start.line, r.start.character, r.end.line, r.end.character]
-    let parentDeclToList (d : RefInfo.ParentDecl) : List Json :=
-      let name := d.name |> toJson
-      let range := rangeToList d.range |>.map toJson
-      let selectionRange := rangeToList d.selectionRange |>.map toJson
-      [name] ++ range ++ selectionRange
    let locationToList (l : RefInfo.Location) : List Json :=
-      let range := rangeToList l.range |>.map toJson
-      let parentDecl := l.parentDecl?.map parentDeclToList |>.getD []
+      let range := [l.startPosLine, l.startPosCharacter, l.endPosLine, l.endPosCharacter].map toJson
+      let parentDecl := l.parentDecl?.map ([toJson ·]) |>.getD []
      range ++ parentDecl
    Json.mkObj [
      ("definition", toJson $ i.definition?.map locationToList),
@ -147,35 +243,30 @@ instance : ToJson RefInfo where
 instance : FromJson RefInfo where
  -- This implementation is optimized to prevent redundant intermediate allocations.
  fromJson? j := do
-    let toRange (a : Array Json) (i : Nat) : Except String Lsp.Range :=
-      if h : a.size < i + 4 then
-        throw s!"Expected list of length 4, not {a.size}"
-      else
-        return {
-          start := {
-            line := ← fromJson? a[i]
-            character := ← fromJson? a[i+1]
-          }
-          «end» := {
-            line := ← fromJson? a[i+2]
-            character := ← fromJson? a[i+3]
-          }
-        }
-    let toParentDecl (a : Array Json) (i : Nat) : Except String RefInfo.ParentDecl := do
-      let name ← fromJson? a[i]!
-      let range ← toRange a (i + 1)
-      let selectionRange ← toRange a (i + 5)
-      return ⟨name, range, selectionRange⟩
    let toLocation (a : Array Json) : Except String RefInfo.Location := do
-      if a.size != 4 && a.size != 13 then
-        .error "Expected list of length 4 or 13, not {l.size}"
-      let range ← toRange a 0
-      if a.size == 13 then
-        let parentDecl ← toParentDecl a 4
-        return ⟨range, parentDecl⟩
+      if h : a.size ≠ 4 ∧ a.size ≠ 5 then
+        .error s!"Expected list of length 4 or 5, not {a.size}"
      else
-        return ⟨range, none⟩
-
+        let startPosLine ← fromJson? a[0]
+        let startPosCharacter ← fromJson? a[1]
+        let endPosLine ← fromJson? a[2]
+        let endPosCharacter ← fromJson? a[3]
+        if h' : a.size = 5 then
+          return {
+            startPosLine
+            startPosCharacter
+            endPosLine
+            endPosCharacter
+            parentDecl := ← fromJson? a[4]
+          }
+        else
+          return {
+            startPosLine
+            startPosCharacter
+            endPosLine
+            endPosCharacter
+            parentDecl := ""
+          }
    let definition? ← j.getObjValAs? (Option $ Array Json) "definition"
    let definition? ← match definition? with
      | none => pure none
@ -213,15 +304,28 @@ structure LeanILeanHeaderInfoParams where
  directImports : Array ImportInfo
  deriving FromJson, ToJson

+/--
+Used in the `$/lean/ileanHeaderSetupInfo` watchdog <- worker notifications.
+Contains status information about `lake setup-file`.
+-/
+structure LeanILeanHeaderSetupInfoParams where
+  /-- Version of the file these imports are from. -/
+  version        : Nat
+  /-- Whether setting up the header using `lake setup-file` has failed. -/
+  isSetupFailure : Bool
+  deriving FromJson, ToJson
+
 /--
 Used in the `$/lean/ileanInfoUpdate` and `$/lean/ileanInfoFinal` watchdog <- worker notifications.
 Contains the definitions and references of the file managed by a worker.
 -/
 structure LeanIleanInfoParams where
  /-- Version of the file these references are from. -/
-  version        : Nat
+  version    : Nat
  /-- All references for the file. -/
-  references     : ModuleRefs
+  references : ModuleRefs
+  /-- All decls for the file. -/
+  decls      : Decls
  deriving FromJson, ToJson

 /--
--- a/src/Lean/Data/Lsp/Ipc.lean
+++ b/src/Lean/Data/Lsp/Ipc.lean
@ -141,6 +141,19 @@ partial def waitForILeans (waitForILeansId : RequestID := 0) (target : DocumentU
    | _ =>
      pure ()

+partial def waitForWatchdogILeans (waitForILeansId : RequestID := 0) : IpcM Unit := do
+  writeRequest ⟨waitForILeansId, "$/lean/waitForILeans", WaitForILeansParams.mk none none⟩
+  while true do
+    match (← readMessage) with
+    | .response id _ =>
+      if id == waitForILeansId then
+        return
+    | .responseError id _ msg _ =>
+      if id == waitForILeansId then
+        throw $ userError s!"Waiting for ILeans failed: {msg}"
+    | _ =>
+      pure ()
+
 /--
 Waits for a diagnostic notification with a specific message to be emitted. Discards all received
 messages, so should not be combined with `collectDiagnostics`.
--- a/src/Lean/Elab/Frontend.lean
+++ b/src/Lean/Elab/Frontend.lean
@ -199,10 +199,12 @@ def runFrontend
  if let some ileanFileName := ileanFileName? then
    let trees := snaps.getAll.flatMap (match ·.infoTree? with | some t => #[t] | _ => #[])
    let references := Lean.Server.findModuleRefs inputCtx.fileMap trees (localVars := false)
+    let (moduleRefs, decls) ← references.toLspModuleRefs
    let ilean := {
      module        := mainModuleName
      directImports := Server.collectImports ⟨snap.stx⟩
-      references    := ← references.toLspModuleRefs
+      references    := moduleRefs
+      decls
      : Lean.Server.Ilean
    }
    IO.FS.writeFile ileanFileName $ Json.compress $ toJson ilean
--- a/src/Lean/Server/FileWorker.lean
+++ b/src/Lean/Server/FileWorker.lean
@ -151,14 +151,18 @@ section Elab
  -- Placed here instead of Lean.Server.Utils because of an import loop
  private def mkIleanInfoNotification (method : String) (m : DocumentMeta)
      (trees : Array Elab.InfoTree) : BaseIO (JsonRpc.Notification Lsp.LeanIleanInfoParams) := do
-    let references ← findModuleRefs m.text trees (localVars := true) |>.toLspModuleRefs
-    let param := { version := m.version, references }
+    let (references, decls) ← findModuleRefs m.text trees (localVars := true) |>.toLspModuleRefs
+    let param := { version := m.version, references, decls }
    return { method, param }

-  private def mkInitialIleanInfoUpdateNotification (m : DocumentMeta)
+  private def mkIleanHeaderInfoNotification (m : DocumentMeta)
      (directImports : Array ImportInfo) : JsonRpc.Notification Lsp.LeanILeanHeaderInfoParams :=
    { method := "$/lean/ileanHeaderInfo", param := { version := m.version, directImports } }

+  private def mkIleanHeaderSetupInfoNotification (m : DocumentMeta)
+      (isSetupFailure : Bool) : JsonRpc.Notification Lsp.LeanILeanHeaderSetupInfoParams :=
+    { method := "$/lean/ileanHeaderSetupInfo", param := { version := m.version, isSetupFailure } }
+
  private def mkIleanInfoUpdateNotification : DocumentMeta → Array Elab.InfoTree →
      BaseIO (JsonRpc.Notification Lsp.LeanIleanInfoParams) :=
    mkIleanInfoNotification "$/lean/ileanInfoUpdate"
@ -396,7 +400,7 @@ def setupImports
    return .error { diagnostics := .empty, result? := none, metaSnap := default }

  let header := stx.toModuleHeader
-  chanOut.sync.send <| .ofMsg <| mkInitialIleanInfoUpdateNotification doc <| collectImports stx
+  chanOut.sync.send <| .ofMsg <| mkIleanHeaderInfoNotification doc <| collectImports stx
  let fileSetupResult ← setupFile doc header fun stderrLine => do
    let progressDiagnostic := {
      range      := ⟨⟨0, 0⟩, ⟨1, 0⟩⟩
@ -406,6 +410,9 @@ def setupImports
      message    := stderrLine
    }
    chanOut.sync.send <| .ofMsg <| mkPublishDiagnosticsNotification doc #[progressDiagnostic]
+  let isSetupError := fileSetupResult.kind matches .importsOutOfDate
+    || fileSetupResult.kind matches .error ..
+  chanOut.sync.send <| .ofMsg <| mkIleanHeaderSetupInfoNotification doc isSetupError
  -- clear progress notifications in the end
  chanOut.sync.send <| .ofMsg <| mkPublishDiagnosticsNotification doc #[]
  match fileSetupResult.kind with
--- a/src/Lean/Server/FileWorker/RequestHandling.lean
+++ b/src/Lean/Server/FileWorker/RequestHandling.lean
@ -266,12 +266,12 @@ partial def handleDocumentHighlight (p : DocumentHighlightParams)

  let highlightRefs? (snaps : Array Snapshot) : IO (Option (Array DocumentHighlight)) := do
    let trees := snaps.map (·.infoTree)
-    let refs : Lsp.ModuleRefs ← findModuleRefs text trees |>.toLspModuleRefs
+    let (refs, _) ← findModuleRefs text trees |>.toLspModuleRefs
    let mut ranges := #[]
    for ident in refs.findAt p.position (includeStop := true) do
      if let some info := refs.get? ident then
-        if let some ⟨definitionRange, _⟩ := info.definition? then
-          ranges := ranges.push definitionRange
+        if let some loc := info.definition? then
+          ranges := ranges.push loc.range
        ranges := ranges.append <| info.usages.map (·.range)
    if ranges.isEmpty then
      return none
--- a/src/Lean/Server/References.lean
+++ b/src/Lean/Server/References.lean
@ -81,11 +81,12 @@ def addRef (i : RefInfo) (ref : Reference) : RefInfo :=
  | { usages, .. }, { isBinder := false, .. } =>
    { i with usages := usages.push ref }

-/-- Converts `i` to a JSON-serializable `Lsp.RefInfo`. -/
-def toLspRefInfo (i : RefInfo) : BaseIO Lsp.RefInfo := do
-  let refToRefInfoLocation (ref : Reference) : BaseIO RefInfo.Location := do
+/-- Converts `i` to a JSON-serializable `Lsp.RefInfo` and collects its decls. -/
+def toLspRefInfo (i : RefInfo) : StateT Decls BaseIO Lsp.RefInfo := do
+  let refToRefInfoLocation (ref : Reference) : StateT Decls BaseIO RefInfo.Location := do
    let parentDeclName? := ref.ci.parentDecl?
-    let .ok parentDeclRanges? ← EIO.toBaseIO <| ref.ci.runCoreM do
+    let parentDeclNameString? := parentDeclName?.map (·.toString)
+    let .ok parentDeclInfo? ← EIO.toBaseIO <| ref.ci.runCoreM do
        let some parentDeclName := parentDeclName?
          | return none
        -- Use `local` as it avoids unnecessary blocking, which is especially important when called
@ -93,17 +94,15 @@ def toLspRefInfo (i : RefInfo) : BaseIO Lsp.RefInfo := do
        -- `parentDeclName` will not be available in the current environment and we would block only
        -- to return `none` in the end anyway. At the end of elaboration, we rerun this function on
        -- the full info tree with the main environment, so the access will succeed immediately.
-        return declRangeExt.find? (asyncMode := .local) (← getEnv) parentDeclName
+        let some parentDeclRanges := declRangeExt.find? (asyncMode := .local) (← getEnv) parentDeclName
+          | return none
+        return some <| .ofDeclarationRanges parentDeclRanges
      -- we only use `CoreM` to get access to a `MonadEnv`, but these are currently all `IO`
      | unreachable!
-    return {
-      range := ref.range
-      parentDecl? := do
-        let parentDeclName ← parentDeclName?
-        let parentDeclRange := (← parentDeclRanges?).range.toLspRange
-        let parentDeclSelectionRange := (← parentDeclRanges?).selectionRange.toLspRange
-        return ⟨parentDeclName.toString, parentDeclRange, parentDeclSelectionRange⟩
-    }
+    if let some parentDeclNameString := parentDeclNameString? then
+      if let some parentDeclInfo := parentDeclInfo? then
+        modify (·.insert parentDeclNameString parentDeclInfo)
+    return .mk ref.range (parentDeclName?.map (·.toString))
  let definition? ← i.definition.mapM refToRefInfoLocation
  let usages ← i.usages.mapM refToRefInfoLocation
  return {
@ -123,8 +122,8 @@ def addRef (self : ModuleRefs) (ref : Reference) : ModuleRefs :=
  let refInfo := self.getD ref.ident RefInfo.empty
  self.insert ref.ident (refInfo.addRef ref)

-/-- Converts `refs` to a JSON-serializable `Lsp.ModuleRefs`. -/
-def toLspModuleRefs (refs : ModuleRefs) : BaseIO Lsp.ModuleRefs := do
+/-- Converts `refs` to a JSON-serializable `Lsp.ModuleRefs` and collects all decls. -/
+def toLspModuleRefs (refs : ModuleRefs) : BaseIO (Lsp.ModuleRefs × Decls) := StateT.run (s := ∅) do
  let mut refs' := ∅
  for (k, v) in refs do
    refs' := refs'.insert k (← v.toLspRefInfo)
@ -207,13 +206,15 @@ open Elab
 /-- Content of individual `.ilean` files -/
 structure Ilean where
  /-- Version number of the ilean format. -/
-  version       : Nat := 4
+  version       : Nat := 5
  /-- Name of the module that this ilean data has been collected for. -/
  module        : Name
  /-- Direct imports of the module. -/
  directImports : Array Lsp.ImportInfo
  /-- All references of this module. -/
  references    : Lsp.ModuleRefs
+  /-- All declarations of this module. -/
+  decls         : Lsp.Decls
  deriving FromJson, ToJson

 namespace Ilean
@ -495,6 +496,8 @@ structure LoadedILean where
  directImports : DirectImports
  /-- Reference information from this ILean. -/
  refs          : Lsp.ModuleRefs
+  /-- Declarations in the module of the ILean. -/
+  decls         : Lsp.Decls

 /-- Paths and module references for every module name. Loaded from `.ilean` files. -/
 abbrev ILeanMap := Std.TreeMap Name LoadedILean Name.quickCmp
@ -506,13 +509,24 @@ built.
 -/
 structure TransientWorkerILean where
  /-- URI of the module that the file worker is associated with. -/
-  moduleUri     : DocumentUri
+  moduleUri       : DocumentUri
  /-- Document version for which these references have been collected. -/
-  version       : Nat
+  version         : Nat
  /-- Direct imports of the module that the file worker is associated with. -/
-  directImports : DirectImports
+  directImports   : DirectImports
+  /--
+  Whether `lake setup-file` has failed for this worker. `none` if no setup info has been received
+  for this version yet.
+  -/
+  isSetupFailure? : Option Bool
  /-- References provided by the worker. -/
-  refs          : Lsp.ModuleRefs
+  refs            : Lsp.ModuleRefs
+  /-- Declarations provided by the worker. -/
+  decls           : Lsp.Decls
+
+/-- Determines whether this transient worker ILean includes actual references. -/
+def TransientWorkerILean.hasRefs (i : TransientWorkerILean) : Bool :=
+  i.isSetupFailure?.any (fun isSetupFailure => ! isSetupFailure)

 /--
 Document versions and module references for every module name. Loaded from the current state
@ -548,6 +562,7 @@ def addIlean
      ileanPath := path
      directImports
      refs := ilean.references
+      decls := ilean.decls
    }
  }

@ -570,14 +585,41 @@ def updateWorkerImports
    : IO References := do
  let directImports ← DirectImports.convertImportInfos directImports
  let some workerRefs := self.workers[name]?
-    | return { self with workers := self.workers.insert name { moduleUri, version, directImports, refs := ∅} }
+    | return { self with workers := self.workers.insert name { moduleUri, version, directImports, isSetupFailure? := none, refs := ∅, decls := ∅} }
  match compare version workerRefs.version with
  | .lt => return self
-  | .gt => return { self with workers := self.workers.insert name { moduleUri, version, directImports, refs := ∅} }
+  | .gt => return { self with workers := self.workers.insert name { moduleUri, version, directImports, isSetupFailure? := none, refs := ∅, decls := ∅} }
+  | .eq =>
+    let isSetupFailure? := workerRefs.isSetupFailure?
+    let refs := workerRefs.refs
+    let decls := workerRefs.decls
+    return { self with
+      workers := self.workers.insert name { moduleUri, version, directImports, isSetupFailure?, refs, decls }
+    }
+
+/--
+Replaces the direct imports of a worker for the module `name` in `self` with
+a new set of direct imports.
+-/
+def updateWorkerSetupInfo
+    (self           : References)
+    (name           : Name)
+    (moduleUri      : DocumentUri)
+    (version        : Nat)
+    (isSetupFailure : Bool)
+    : IO References := do
+  let isSetupFailure? := some isSetupFailure
+  let some workerRefs := self.workers[name]?
+    | return { self with workers := self.workers.insert name { moduleUri, version, directImports := ∅, isSetupFailure?, refs := ∅, decls := ∅} }
+  let directImports := workerRefs.directImports
+  match compare version workerRefs.version with
+  | .lt => return self
+  | .gt => return { self with workers := self.workers.insert name { moduleUri, version, directImports, isSetupFailure?, refs := ∅, decls := ∅} }
  | .eq =>
    let refs := workerRefs.refs
+    let decls := workerRefs.decls
    return { self with
-      workers := self.workers.insert name { moduleUri, version, directImports, refs }
+      workers := self.workers.insert name { moduleUri, version, directImports, isSetupFailure?, refs, decls }
    }

 /--
@ -591,18 +633,21 @@ def updateWorkerRefs
    (moduleUri : DocumentUri)
    (version   : Nat)
    (refs      : Lsp.ModuleRefs)
+    (decls     : Lsp.Decls)
    : IO References := do
  let some workerRefs := self.workers[name]?
-    | return { self with workers := self.workers.insert name { moduleUri, version, directImports := ∅, refs } }
+    | return { self with workers := self.workers.insert name { moduleUri, version, directImports := ∅, isSetupFailure? := none, refs, decls } }
  let directImports := workerRefs.directImports
+  let isSetupFailure? := workerRefs.isSetupFailure?
  match compare version workerRefs.version with
  | .lt => return self
-  | .gt => return { self with workers := self.workers.insert name { moduleUri, version, directImports, refs } }
+  | .gt => return { self with workers := self.workers.insert name { moduleUri, version, directImports, isSetupFailure?, refs, decls } }
  | .eq =>
    let mergedRefs := refs.foldl (init := workerRefs.refs) fun m ident info =>
      m.getD ident Lsp.RefInfo.empty |>.merge info |> m.insert ident
+    let mergedDecls := workerRefs.decls.insertMany decls
    return { self with
-      workers := self.workers.insert name { moduleUri, version, directImports, refs := mergedRefs }
+      workers := self.workers.insert name { moduleUri, version, directImports, isSetupFailure?, refs := mergedRefs, decls := mergedDecls }
    }

 /--
@ -615,15 +660,17 @@ def finalizeWorkerRefs
    (moduleUri : DocumentUri)
    (version   : Nat)
    (refs      : Lsp.ModuleRefs)
+    (decls     : Lsp.Decls)
    : IO References := do
  let some workerRefs := self.workers[name]?
-    | return { self with workers := self.workers.insert name { moduleUri, version, directImports := ∅, refs } }
+    | return { self with workers := self.workers.insert name { moduleUri, version, directImports := ∅, isSetupFailure? := none, refs, decls } }
  let directImports := workerRefs.directImports
+  let isSetupFailure? := workerRefs.isSetupFailure?
  match compare version workerRefs.version with
  | .lt => return self
-  | .gt => return { self with workers := self.workers.insert name { moduleUri, version, directImports, refs} }
+  | .gt => return { self with workers := self.workers.insert name { moduleUri, version, directImports, isSetupFailure?, refs, decls } }
  | .eq =>
-    return { self with workers := self.workers.insert name { moduleUri, version, directImports, refs } }
+    return { self with workers := self.workers.insert name { moduleUri, version, directImports, isSetupFailure?, refs, decls } }

 /-- Erases all worker references in `self` for the worker managing `name`. -/
 def removeWorkerRefs (self : References) (name : Name) : References :=
@ -633,12 +680,16 @@ def removeWorkerRefs (self : References) (name : Name) : References :=
 Map from each module to all of its references.
 The current references in a file worker take precedence over those in .ilean files.
 -/
-abbrev AllRefsMap := Std.TreeMap Name (DocumentUri × Lsp.ModuleRefs) Name.quickCmp
+abbrev AllRefsMap := Std.TreeMap Name (DocumentUri × Lsp.ModuleRefs × Lsp.Decls) Name.quickCmp

 /-- Yields a map from all modules to all of their references. -/
 def allRefs (self : References) : AllRefsMap :=
-  let ileanRefs := self.ileans.foldl (init := ∅) fun m name { moduleUri, refs, .. } => m.insert name (moduleUri, refs)
-  self.workers.foldl (init := ileanRefs) fun m name { moduleUri, refs, ..} => m.insert name (moduleUri, refs)
+  let ileanRefs := self.ileans.foldl (init := ∅) fun m name { moduleUri, refs, decls, .. } => m.insert name (moduleUri, refs, decls)
+  self.workers.foldl (init := ileanRefs) fun m name i@{ moduleUri, refs, decls, ..} =>
+    if i.hasRefs then
+      m.insert name (moduleUri, refs, decls)
+    else
+      m

 /--
 Map from each module to all of its direct imports.
@ -659,11 +710,12 @@ def allDirectImports (self : References) : AllDirectImportsMap := Id.run do
 Gets the references for `mod`.
 The current references in a file worker take precedence over those in .ilean files.
 -/
-def getModuleRefs? (self : References) (mod : Name) : Option (DocumentUri × Lsp.ModuleRefs) := do
+def getModuleRefs? (self : References) (mod : Name) : Option (DocumentUri × Lsp.ModuleRefs × Lsp.Decls) := do
  if let some worker := self.workers[mod]? then
-    return (worker.moduleUri, worker.refs)
+    if worker.hasRefs then
+      return (worker.moduleUri, worker.refs, worker.decls)
  if let some ilean := self.ileans[mod]? then
-    return (ilean.moduleUri, ilean.refs)
+    return (ilean.moduleUri, ilean.refs, ilean.decls)
  none

 /--
@ -677,6 +729,15 @@ def getDirectImports? (self : References) (mod : Name) : Option DirectImports :=
    return ilean.directImports
  none

+/-- Gets the set of declarations of `mod`. -/
+def getDecls? (self : References) (mod : Name) : Option Decls := do
+  if let some worker := self.workers[mod]? then
+    if worker.hasRefs then
+      return worker.decls
+  if let some ilean := self.ileans[mod]? then
+    return ilean.decls
+  none
+
 /--
 Yields all references in `self` for `ident`, as well as the `DocumentUri` that each
 reference occurs in.
@ -684,32 +745,50 @@ reference occurs in.
 def allRefsFor
    (self  : References)
    (ident : RefIdent)
-    : Array (DocumentUri × Name × Lsp.RefInfo) := Id.run do
+    : Array (DocumentUri × Name × Lsp.RefInfo × Decls) := Id.run do
  let refsToCheck := match ident with
    | RefIdent.const .. => self.allRefs.toArray
    | RefIdent.fvar identModule .. =>
      let identModuleName := identModule.toName
      match self.getModuleRefs? identModuleName with
      | none => #[]
-      | some (moduleUri, refs) => #[(identModuleName, moduleUri, refs)]
+      | some (moduleUri, refs, decls) => #[(identModuleName, moduleUri, refs, decls)]
  let mut result := #[]
-  for (module, moduleUri, refs) in refsToCheck do
+  for (module, moduleUri, refs, decls) in refsToCheck do
    let some info := refs.get? ident
      | continue
-    result := result.push (moduleUri, module, info)
+    result := result.push (moduleUri, module, info, decls)
  return result

 /-- Yields all references in `module` at `pos`. -/
 def findAt (self : References) (module : Name) (pos : Lsp.Position) (includeStop := false) : Array RefIdent := Id.run do
-  if let some (_, refs) := self.getModuleRefs? module then
+  if let some (_, refs, _) := self.getModuleRefs? module then
    return refs.findAt pos includeStop
  #[]

 /-- Yields the first reference in `module` at `pos`. -/
 def findRange? (self : References) (module : Name) (pos : Lsp.Position) (includeStop := false) : Option Range := do
-  let (_, refs) ← self.getModuleRefs? module
+  let (_, refs, _) ← self.getModuleRefs? module
  refs.findRange? pos includeStop

+/-- Parent declaration of an identifier. -/
+structure ParentDecl where
+  /-- Name of the parent declaration. -/
+  name : String
+  /-- Range of the parent declaration. -/
+  range : Lsp.Range
+  /-- Selection range of the parent declaration. -/
+  selectionRange : Lsp.Range
+
+/-- Yields a `ParentDecl` for the declaration `name`. -/
+def ParentDecl.ofDecls? (ds : Decls) (name : String) : Option ParentDecl := do
+  let d ← ds.get? name
+  return {
+    name
+    range := d.range
+    selectionRange := d.selectionRange
+  }
+
 /-- Location and parent declaration of a reference. -/
 structure DocumentRefInfo where
  /-- Location of the reference. -/
@ -717,7 +796,7 @@ structure DocumentRefInfo where
  /-- Module name of the reference. -/
  module      : Name
  /-- Parent declaration of the reference. -/
-  parentInfo? : Option RefInfo.ParentDecl
+  parentInfo? : Option ParentDecl

 /-- Yields locations and parent declaration for all references referring to `ident`. -/
 def referringTo
@ -726,12 +805,16 @@ def referringTo
    (includeDefinition : Bool := true)
    : Array DocumentRefInfo := Id.run do
  let mut result := #[]
-  for (moduleUri, module, info) in self.allRefsFor ident do
+  for (moduleUri, module, info, decls) in self.allRefsFor ident do
    if includeDefinition then
-      if let some ⟨range, parentDeclInfo?⟩ := info.definition? then
-        result := result.push ⟨⟨moduleUri, range⟩, module, parentDeclInfo?⟩
-    for ⟨range, parentDeclInfo?⟩ in info.usages do
-      result := result.push ⟨⟨moduleUri, range⟩, module, parentDeclInfo?⟩
+      if let some loc := info.definition? then
+        let parentDecl? := do
+          ParentDecl.ofDecls? decls <| ← loc.parentDecl?
+        result := result.push ⟨⟨moduleUri, loc.range⟩, module, parentDecl?⟩
+    for loc in info.usages do
+      let parentDecl? := do
+        ParentDecl.ofDecls? decls <| ← loc.parentDecl?
+      result := result.push ⟨⟨moduleUri, loc.range⟩, module, parentDecl?⟩
  return result

 /-- Yields the definition location of `ident`. -/
@ -739,10 +822,12 @@ def definitionOf?
    (self  : References)
    (ident : RefIdent)
    : Option DocumentRefInfo := Id.run do
-  for (moduleUri, module, info) in self.allRefsFor ident do
-    let some ⟨definitionRange, definitionParentDeclInfo?⟩ := info.definition?
+  for (moduleUri, module, info, decls) in self.allRefsFor ident do
+    let some loc := info.definition?
      | continue
-    return some ⟨⟨moduleUri, definitionRange⟩, module, definitionParentDeclInfo?⟩
+    let definitionParentDecl? := do
+        ParentDecl.ofDecls? decls <| ← loc.parentDecl?
+    return some ⟨⟨moduleUri, loc.range⟩, module, definitionParentDecl?⟩
  return none

 /-- A match in `References.definitionsMatching`. -/
@ -763,16 +848,16 @@ def definitionsMatching
    (cancelTk?      : Option CancelToken := none)
    : BaseIO (Array (MatchedDefinition α)) := do
  let mut result := #[]
-  for (module, moduleUri, refs) in self.allRefs do
+  for (module, moduleUri, refs, _) in self.allRefs do
    if let some cancelTk := cancelTk? then
      if ← cancelTk.isSet then
        return result
    for (ident, info) in refs do
-      let (RefIdent.const _ nameString, some ⟨definitionRange, _⟩) := (ident, info.definition?)
+      let (RefIdent.const _ nameString, some loc) := (ident, info.definition?)
        | continue
      let some v := filterMapIdent nameString.toName
        | continue
-      result := result.push ⟨module, moduleUri, v, definitionRange⟩
+      result := result.push ⟨module, moduleUri, v, loc.range⟩
  return result

 /-- Yields all imports that import the given `requestedMod`. -/
--- a/src/Lean/Server/Watchdog.lean
+++ b/src/Lean/Server/Watchdog.lean
@ -533,10 +533,15 @@ section ServerM
    let uri := fw.doc.uri
    modifyReferencesIO (·.updateWorkerImports module uri params.version params.directImports)

+  def handleILeanHeaderSetupInfo (fw : FileWorker) (params : LeanILeanHeaderSetupInfoParams) : ServerM Unit := do
+    let module := fw.doc.mod
+    let uri := fw.doc.uri
+    modifyReferencesIO (·.updateWorkerSetupInfo module uri params.version params.isSetupFailure)
+
  def handleIleanInfoUpdate (fw : FileWorker) (params : LeanIleanInfoParams) : ServerM Unit := do
    let module := fw.doc.mod
    let uri := fw.doc.uri
-    modifyReferencesIO (·.updateWorkerRefs module uri params.version params.references)
+    modifyReferencesIO (·.updateWorkerRefs module uri params.version params.references params.decls)

  def handleIleanInfoFinal (fw : FileWorker) (params : LeanIleanInfoParams) : ServerM Unit := do
    let s ← read
@ -544,7 +549,7 @@ section ServerM
    let uri := fw.doc.uri
    s.referenceData.atomically do
      let rd ← get
-      let rd ← rd.modifyReferencesM (·.finalizeWorkerRefs module uri params.version params.references)
+      let rd ← rd.modifyReferencesM (·.finalizeWorkerRefs module uri params.version params.references params.decls)
      let (pendingWaitForILeanRequests, rest) := rd.pendingWaitForILeanRequests.partition (·.uri == uri)
      let rd := { rd with pendingWaitForILeanRequests := rest }
      let rd := rd.modifyFinalizedWorkerILeanVersions (·.insert uri params.version)
@ -852,6 +857,9 @@ section ServerM
        | .notification "$/lean/ileanHeaderInfo" =>
          if let .ok params := parseNotificationParams? LeanILeanHeaderInfoParams msg then
            handleILeanHeaderInfo fw params
+        | .notification "$/lean/ileanHeaderSetupInfo" =>
+          if let .ok params := parseNotificationParams? LeanILeanHeaderSetupInfoParams msg then
+            handleILeanHeaderSetupInfo fw params
        | .notification "$/lean/ileanInfoUpdate" =>
          if let .ok params := parseNotificationParams? LeanIleanInfoParams msg then
            handleIleanInfoUpdate fw params
@ -1085,17 +1093,19 @@ def handleCallHierarchyIncomingCalls (p : CallHierarchyIncomingCallsParams)
  let references := (← read).references
  let identRefs := references.referringTo (.const itemData.module.toString itemData.name.toString) false

-  let incomingCalls ← identRefs.filterMapM fun ⟨location, refModule, parentDecl?⟩ => do
+  let incomingCalls ← identRefs.mapM fun ⟨location, refModule, parentDecl?⟩ => do

-    let some ⟨parentDeclNameString, parentDeclRange, parentDeclSelectionRange⟩ := parentDecl?
-      | return none
+    let ⟨parentDeclNameString, parentDeclRange, parentDeclSelectionRange⟩ :=
+      match parentDecl? with
+      | some parentDecl => parentDecl
+      | none => ⟨"[anonymous]", location.range, location.range⟩

    let parentDeclName := parentDeclNameString.toName

    -- Remove private header from name
    let label := Lean.privateToUserName? parentDeclName |>.getD parentDeclName

-    return some {
+    return {
      «from» := {
        name           := label.toString
        kind           := SymbolKind.constant
@ -1131,14 +1141,14 @@ def handleCallHierarchyOutgoingCalls (p : CallHierarchyOutgoingCallsParams)

  let references := (← read).references

-  let some (_, refs) := references.getModuleRefs? itemData.module
+  let some (_, refs, _) := references.getModuleRefs? itemData.module
    | return #[]

  let items ← refs.toArray.filterMapM fun ⟨ident, info⟩ => do
    let outgoingUsages := info.usages.filter fun usage => Id.run do
      let some parentDecl := usage.parentDecl?
        | return false
-      return itemData.name == parentDecl.name.toName
+      return itemData.name == parentDecl.toName

    let outgoingUsages := outgoingUsages.map (·.range)
    if outgoingUsages.isEmpty then
@ -1425,7 +1435,13 @@ section MessageHandling
    | "$/lean/waitForILeans" =>
      let rd ← ctx.referenceData.atomically get
      IO.wait rd.loadingTask.task
-      let ⟨uri, version⟩ ← parseParams WaitForILeansParams params
+      let ⟨some uri, some version⟩ ← parseParams WaitForILeansParams params
+        | writeMessage {
+            id
+            result := ⟨⟩
+            : Response WaitForILeans
+          }
+          return
      if let none ← getFileWorker? uri then
        writeMessage {
          id
--- a/src/stdlib_flags.h
+++ b/src/stdlib_flags.h
@ -1,6 +1,5 @@
 #include "util/options.h"

-// stage0 update needed for grind_annotated command
 namespace lean {
 options get_default_options() {
    options opts;
--- a/stage0/src/stdlib_flags.h
+++ b/stage0/src/stdlib_flags.h
@ -1,6 +1,6 @@
 #include "util/options.h"

-// stage0 update needed for grind_annotated command
+// Please update
 namespace lean {
 options get_default_options() {
    options opts;
--- a/tests/bench/ilean_roundtrip.lean
+++ b/tests/bench/ilean_roundtrip.lean
@ -1,14 +1,8 @@
 import Lean.Data.Lsp

 def genModuleRefs (n : Nat) : IO Lean.Lsp.ModuleRefs := do
-  let someLoc : Lean.Lsp.RefInfo.Location := {
-    range := ⟨⟨333, 444⟩, ⟨444, 555⟩⟩
-    parentDecl? := some {
-      name := "A.Reasonably.Long.ParentDecl.Name.barfoo",
-      range := ⟨⟨1111, 2222⟩, ⟨3333, 4444⟩⟩
-      selectionRange := ⟨⟨5555, 6666⟩, ⟨7777, 8888⟩⟩
-    }
-  }
+  let someLoc : Lean.Lsp.RefInfo.Location :=
+    .mk ⟨⟨333, 444⟩, ⟨444, 555⟩⟩ <| some "A.Reasonably.Long.ParentDecl.Name.barfoo"

  let mut someUsages : Array Lean.Lsp.RefInfo.Location := #[]
  for _ in [0, 200] do
--- a/tests/bench/server_startup.lean
+++ b/tests/bench/server_startup.lean
@ -6,7 +6,7 @@ def main : IO Unit := do
    let hIn ← Ipc.stdin
    hIn.write (←FS.readBinFile "server_startup.log")
    hIn.flush
-    let initResp ← Ipc.readResponseAs 0 InitializeResult
+    let _ ← Ipc.readResponseAs 0 InitializeResult
    Ipc.writeNotification ⟨"initialized", InitializedParams.mk⟩

    Ipc.shutdown 1
--- a/tests/bench/speedcenter.exec.velcom.yaml
+++ b/tests/bench/speedcenter.exec.velcom.yaml
@ -266,9 +266,19 @@
 - attributes:
    description: language server startup
    tags: [other]
+  build_config:
+    cmd: ./compile.sh server_startup.lean
  run_config:
    <<: *time
-    cmd: lean -Dlinter.all=false --run server_startup.lean
+    cmd: ./server_startup.lean.out
+- attributes:
+    description: language server startup with ileans
+    tags: [other]
+  build_config:
+    cmd: ./compile.sh watchdogRss.lean
+  run_config:
+    <<: *time
+    cmd: ./watchdogRss.lean.out
 - attributes:
    description: ilean roundtrip
    tags: [other]
--- a/tests/bench/watchdogRss.lean
+++ b/tests/bench/watchdogRss.lean
@ -0,0 +1,20 @@
+import Lean.Data.Lsp
+open Lean.Lsp
+
+def main (_ : List String) : IO Unit := do
+  Ipc.runWith "lean" (#["--server"]) do
+    let capabilities := {
+      textDocument? := some {
+        completion? := some {
+          completionItem? := some {
+            insertReplaceSupport? := true
+          }
+        }
+      }
+    }
+    Ipc.writeRequest ⟨0, "initialize", { capabilities : InitializeParams }⟩
+    discard <| Ipc.readResponseAs 0 InitializeResult
+    Ipc.writeNotification ⟨"initialized", InitializedParams.mk⟩
+    Ipc.waitForWatchdogILeans 1
+    Ipc.shutdown 2
+    discard <| Ipc.waitForExit