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.

Author: mhuisi

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

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}"
+    let toLocation (a : Array Json) : Except String RefInfo.Location := do
+      if h : a.size ≠ 4 ∧ a.size ≠ 5 then
+        .error s!"Expected list of length 4 or 5, not {a.size}"
       else
-        return {
-          start := {
-            line := ← fromJson? a[i]
-            character := ← fromJson? a[i+1]
+        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]
           }
-          «end» := {
-            line := ← fromJson? a[i+2]
-            character := ← fromJson? a[i+3]
+        else
+          return {
+            startPosLine
+            startPosCharacter
+            endPosLine
+            endPosCharacter
+            parentDecl := ""
           }
-        }
-    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⟩
-      else
-        return ⟨range, none⟩
-
     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
 
 /--

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`.

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

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

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