feat: improve workspace symbol performance (#8091)

This PR improves the performance of the workspace symbol request. In my testing on my machine, the time to respond to the workspace symbol request containing just `c` in Mathlib has been reduced to ~1200ms from ~11000ms. We also serve the nearest-matching 1000 symbols instead of just the first 100 now and use the length of the symbol as a tie-breaker for when the fuzzy matching score is equal. Some further improvements might be gained in the future when #8087 is fixed and we can switch back to `qsort`.
2025-05-13 10:29:49 +02:00 · 2025-05-13 10:29:49 +02:00 · efcf94298a
commit efcf94298a
parent f75e36dcdb
4 changed files with 142 additions and 83 deletions
--- a/src/Lean/Server/Completion/CompletionCollectors.lean
+++ b/src/Lean/Server/Completion/CompletionCollectors.lean
@ -121,24 +121,6 @@ end Infrastructure

 section Utils

-  private partial def containsSuccessiveCharacters (a b : String) : Bool :=
-    go ⟨0⟩ ⟨0⟩
-  where
-    go (aPos bPos : String.Pos) : Bool :=
-      if ha : a.atEnd aPos then
-        true
-      else if hb : b.atEnd bPos then
-        false
-      else
-        let ac := a.get' aPos ha
-        let bc := b.get' bPos hb
-        let bPos := b.next' bPos hb
-        if ac == bc then
-          let aPos := a.next' aPos ha
-          go aPos bPos
-        else
-          go aPos bPos
-
  private def normPrivateName? (declName : Name) : MetaM (Option Name) := do
    match privateToUserName? declName with
    | none => return declName
@ -169,7 +151,7 @@ section Utils
    else if let (.str p₁ s₁, .str p₂ s₂) := (id, declName) then
      if p₁ == p₂ then
        -- If the namespaces agree, fuzzy-match on the trailing part
-        if containsSuccessiveCharacters s₁ s₂ then
+        if s₁.charactersIn s₂ then
          return some <| .mkSimple s₂
        else
          return none
@ -177,7 +159,7 @@ section Utils
        -- If `id` is namespace-less, also fuzzy-match declaration names in arbitrary namespaces
        -- (but don't match the namespace itself).
        -- Penalize score by component length of added namespace.
-        if containsSuccessiveCharacters s₁ s₂ then
+        if s₁.charactersIn s₂ then
          return some declName
        else
          return none
@ -202,7 +184,7 @@ section IdCompletionUtils
      false
    else
      match id, declName with
-      | .str .anonymous s₁, .str .anonymous s₂ => containsSuccessiveCharacters s₁ s₂
+      | .str .anonymous s₁, .str .anonymous s₂ => s₁.charactersIn s₂
      | _, _ => false

  /--
@ -236,7 +218,7 @@ section IdCompletionUtils
    else
      match ns, nsFragment with
      | .str p₁ s₁, .str p₂ s₂ =>
-        if p₁ == p₂ then containsSuccessiveCharacters s₂ s₁ else false
+        if p₁ == p₂ then s₂.charactersIn s₁ else false
      | _, _ => false

  def completeNamespaces (ctx : ContextInfo) (id : Name) (danglingDot : Bool) : M Unit := do
@ -543,7 +525,7 @@ def fieldIdCompletion
    let fieldNames := getStructureFieldsFlattened (← getEnv) structName (includeSubobjectFields := false)
    for fieldName in fieldNames do
      let .str _ fieldName := fieldName | continue
-      if ! containsSuccessiveCharacters idStr fieldName then
+      if ! idStr.charactersIn fieldName then
        continue
      let item := { label := fieldName, detail? := "field", documentation? := none, kind? := CompletionItemKind.field }
      addItem item
@ -571,7 +553,7 @@ def optionCompletion
    let opts ← getOptions
    let mut items := #[]
    for ⟨name, decl⟩ in decls do
-      if containsSuccessiveCharacters partialName name.toString then
+      if partialName.charactersIn name.toString then
        let textEdit :=
          if !caps.textDocument?.any (·.completion?.any (·.completionItem?.any (·.insertReplaceSupport?.any (·)))) then
            none -- InsertReplaceEdit not supported by client
--- a/src/Lean/Server/Completion/CompletionUtils.lean
+++ b/src/Lean/Server/Completion/CompletionUtils.lean
@ -7,6 +7,24 @@ prelude
 import Init.Prelude
 import Lean.Meta.WHNF

+partial def String.charactersIn (a b : String) : Bool :=
+  go ⟨0⟩ ⟨0⟩
+where
+  go (aPos bPos : String.Pos) : Bool :=
+    if ha : a.atEnd aPos then
+      true
+    else if hb : b.atEnd bPos then
+      false
+    else
+      let ac := a.get' aPos ha
+      let bc := b.get' bPos hb
+      let bPos := b.next' bPos hb
+      if ac == bc then
+        let aPos := a.next' aPos ha
+        go aPos bPos
+      else
+        go aPos bPos
+
 namespace Lean.Server.Completion
 open Elab

--- a/src/Lean/Server/References.lean
+++ b/src/Lean/Server/References.lean
@ -376,13 +376,36 @@ def findModuleRefs (text : FileMap) (trees : Array InfoTree) (localVars : Bool :

 /-! # Collecting and maintaining reference info from different sources -/

+/-- Reference information from a loaded ILean file. -/
+structure LoadedILean where
+  /-- URI of the module of this ILean. -/
+  moduleUri : DocumentUri
+  /-- Path to the ILean file. -/
+  ileanPath : System.FilePath
+  /-- Reference information from this ILean. -/
+  refs      : Lsp.ModuleRefs
+
 /-- Paths and module references for every module name. Loaded from `.ilean` files. -/
-abbrev ILeanMap := Std.TreeMap Name (System.FilePath × Lsp.ModuleRefs) Name.quickCmp
+abbrev ILeanMap := Std.TreeMap Name LoadedILean Name.quickCmp
+
+/--
+Transient reference information from a file worker.
+We track this information so that we have up-to-date reference information before a file has been
+built.
+-/
+structure TransientWorkerILean where
+  /-- URI of the module that the file worker is associated with. -/
+  moduleUri : DocumentUri
+  /-- Document version for which these references have been collected. -/
+  version   : Nat
+  /-- References provided by the worker. -/
+  refs      : Lsp.ModuleRefs
+
 /--
 Document versions and module references for every module name. Loaded from the current state
 in a file worker.
 -/
-abbrev WorkerRefMap := Std.TreeMap Name (Nat × Lsp.ModuleRefs) Name.quickCmp
+abbrev WorkerRefMap := Std.TreeMap Name TransientWorkerILean Name.quickCmp

 /-- References from ilean files and current ilean information from file workers. -/
 structure References where
@ -397,12 +420,22 @@ namespace References
 def empty : References := { ileans := ∅, workers := ∅ }

 /-- Adds the contents of an ilean file `ilean` at `path` to `self`. -/
-def addIlean (self : References) (path : System.FilePath) (ilean : Ilean) : References :=
-  { self with ileans := self.ileans.insert ilean.module (path, ilean.references) }
+def addIlean
+    (self      : References)
+    (moduleUri : DocumentUri)
+    (path      : System.FilePath)
+    (ilean     : Ilean)
+    : References := { self with
+  ileans := self.ileans.insert ilean.module {
+    moduleUri
+    ileanPath := path
+    refs := ilean.references
+  }
+}

 /-- Removes the ilean file data at `path` from `self`. -/
 def removeIlean (self : References) (path : System.FilePath) : References :=
-  let namesToRemove := self.ileans.filter (fun _ (p, _) => p == path)
+  let namesToRemove := self.ileans.filter (fun _ { ileanPath, .. } => ileanPath == path)
  namesToRemove.foldl (init := self) fun self name _ =>
    { self with ileans := self.ileans.erase name }

@ -411,26 +444,40 @@ Updates the worker references in `self` with the `refs` of the worker managing t
 Replaces the current references with `refs` if `version` is newer than the current version managed
 in `refs` and otherwise merges the reference data if `version` is equal to the current version.
 -/
-def updateWorkerRefs (self : References) (name : Name) (version : Nat) (refs : Lsp.ModuleRefs) : References := Id.run do
-  if let some (currVersion, _) := self.workers[name]? then
+def updateWorkerRefs
+    (self      : References)
+    (name      : Name)
+    (moduleUri : DocumentUri)
+    (version   : Nat)
+    (refs      : Lsp.ModuleRefs)
+    : References := Id.run do
+  if let some { version := currVersion, .. } := self.workers[name]? then
    if version > currVersion then
-      return { self with workers := self.workers.insert name (version, refs) }
+      return { self with workers := self.workers.insert name { moduleUri, version, refs} }
    if version == currVersion then
-      let current := self.workers.getD name (version, Std.TreeMap.empty)
-      let merged := refs.foldl (init := current.snd) fun m ident info =>
+      let current := self.workers.getD name { moduleUri, version, refs := Std.TreeMap.empty }
+      let merged := refs.foldl (init := current.refs) fun m ident info =>
        m.getD ident Lsp.RefInfo.empty |>.merge info |> m.insert ident
-      return { self with workers := self.workers.insert name (version, merged) }
+      return { self with
+        workers := self.workers.insert name { moduleUri, version, refs := merged }
+      }
  return self

 /--
 Replaces the worker references in `self` with the `refs` of the worker managing the module `name`
 if `version` is newer than the current version managed in `refs`.
 -/
-def finalizeWorkerRefs (self : References) (name : Name) (version : Nat) (refs : Lsp.ModuleRefs) : References := Id.run do
-  if let some (currVersion, _) := self.workers[name]? then
+def finalizeWorkerRefs
+    (self      : References)
+    (name      : Name)
+    (moduleUri : DocumentUri)
+    (version   : Nat)
+    (refs      : Lsp.ModuleRefs)
+    : References := Id.run do
+  if let some { version := currVersion, .. } := self.workers[name]? then
    if version < currVersion then
      return self
-  return { self with workers := self.workers.insert name (version, refs) }
+  return { self with workers := self.workers.insert name { moduleUri, version, refs } }

 /-- Erases all worker references in `self` for the worker managing `name`. -/
 def removeWorkerRefs (self : References) (name : Name) : References :=
@ -440,19 +487,23 @@ def removeWorkerRefs (self : References) (name : Name) : References :=
 All references for a module.
 The current references in a file worker take precedence over those in .ilean files.
 -/
-abbrev AllRefsMap := Std.TreeMap Name Lsp.ModuleRefs Name.quickCmp
+abbrev AllRefsMap := Std.TreeMap Name (DocumentUri × Lsp.ModuleRefs) 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 (_, refs) => m.insert name refs
-  self.workers.foldl (init := ileanRefs) fun m name (_, refs) => m.insert name refs
+  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)

 /--
 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 Lsp.ModuleRefs :=
-  self.workers[mod]?.map (·.2) <|> self.ileans[mod]?.map (·.2)
+def getModuleRefs? (self : References) (mod : Name) : Option (DocumentUri × Lsp.ModuleRefs) := do
+  if let some worker := self.workers[mod]? then
+    return (worker.moduleUri, worker.refs)
+  if let some ilean := self.ileans[mod]? then
+    return (ilean.moduleUri, ilean.refs)
+  none

 /--
 Yields all references in `self` for `ident`, as well as the `DocumentUri` that each
@ -468,25 +519,23 @@ def allRefsFor
      let identModuleName := identModule.toName
      match self.getModuleRefs? identModuleName with
      | none => #[]
-      | some refs => #[(identModuleName, refs)]
+      | some (moduleUri, refs) => #[(identModuleName, moduleUri, refs)]
  let mut result := #[]
-  for (module, refs) in refsToCheck do
+  for (module, moduleUri, refs) in refsToCheck do
    let some info := refs.get? ident
      | continue
-    let some uri ← documentUriFromModule? module
-      | continue
-    result := result.push (uri, module, info)
+    result := result.push (moduleUri, module, info)
  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

 /-- Location and parent declaration of a reference. -/
@ -525,34 +574,33 @@ def definitionOf?
  return none

 /-- A match in `References.definitionsMatching`. -/
-structure MatchedDefinition (α β : Type) where
+structure MatchedDefinition (α : Type) where
  /-- Result of `filterMapMod`. -/
-  mod   : α
+  mod    : Name
+  /-- URI for `mod`. -/
+  modUri : DocumentUri
  /-- Result of `filterMapIdent`. -/
-  ident : β
+  ident  : α
  /-- Definition range of matched identifier. -/
-  range : Range
+  range  : Range

 /-- Yields all definitions matching the given `filter`. -/
 def definitionsMatching
    (self           : References)
-    (filterMapMod   : Name → IO (Option α))
-    (filterMapIdent : Name → IO (Option β))
+    (filterMapIdent : Name → IO (Option α))
    (cancelTk?      : Option CancelToken := none)
-    : IO (Array (MatchedDefinition α β)) := do
+    : IO (Array (MatchedDefinition α)) := do
  let mut result := #[]
-  for (module, 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
-    let some a ← filterMapMod module
-      | continue
    for (ident, info) in refs do
      let (RefIdent.const _ nameString, some ⟨definitionRange, _⟩) := (ident, info.definition?)
        | continue
-      let some b ← filterMapIdent nameString.toName
+      let some v ← filterMapIdent nameString.toName
        | continue
-      result := result.push ⟨a, b, definitionRange⟩
+      result := result.push ⟨module, moduleUri, v, definitionRange⟩
  return result

 end References
--- a/src/Lean/Server/Watchdog.lean
+++ b/src/Lean/Server/Watchdog.lean
@ -20,6 +20,8 @@ import Lean.Server.Utils
 import Lean.Server.Requests
 import Lean.Server.References
 import Lean.Server.ServerTask
+import Lean.Server.Completion.CompletionUtils
+import Init.Data.List.Sort

 /-!
 For general server architecture, see `README.md`. This module implements the watchdog process.
@ -454,9 +456,10 @@ section ServerM
    st.hLog.flush

  def handleIleanInfoUpdate (fw : FileWorker) (params : LeanIleanInfoParams) : ServerM Unit := do
-    let some module ← getFileWorkerMod? fw.doc.uri
+    let uri := fw.doc.uri
+    let some module ← getFileWorkerMod? uri
      | return
-    modifyReferences (·.updateWorkerRefs module params.version params.references)
+    modifyReferences (·.updateWorkerRefs module uri params.version params.references)

  def handleIleanInfoFinal (fw : FileWorker) (params : LeanIleanInfoParams) : ServerM Unit := do
    let s ← read
@ -465,7 +468,7 @@ section ServerM
      | return
    s.referenceData.atomically do
      let pendingWaitForILeanRequests ← modifyGet fun rd =>
-        let rd := rd.modifyReferences (·.finalizeWorkerRefs module params.version params.references)
+        let rd := rd.modifyReferences (·.finalizeWorkerRefs module uri params.version params.references)
        let (pending, rest) := rd.pendingWaitForILeanRequests.partition (·.uri == uri)
        let rd := { rd with pendingWaitForILeanRequests := rest }
        let rd := rd.modifyFinalizedWorkerILeanVersions (·.insert uri params.version)
@ -581,12 +584,11 @@ section ServerM
    let task ← ServerTask.IO.asTask do
      let mut queryResults : Array LeanQueriedModule := #[]
      for query in params.queries do
-        let filterMapMod mod := pure <| some mod
        let filterMapIdent decl := pure <| matchAgainstQuery? query decl
-        let symbols ← refs.definitionsMatching filterMapMod filterMapIdent cancelTk
-        let sorted := symbols.qsort fun { ident := m1, .. } { ident := m2, .. } =>
+        let symbols ← refs.definitionsMatching filterMapIdent cancelTk
+        let sorted := symbols.toList.mergeSort fun { ident := m1, .. } { ident := m2, .. } =>
          m1.fastCompare m2 == .gt
-        let result : LeanQueriedModule := sorted.extract 0 10 |>.map fun m => {
+        let result : LeanQueriedModule := sorted.take 10 |>.toArray.map fun m => {
          module := m.mod
          decl := m.ident.decl
          isExactMatch := m.ident.isExactMatch
@ -975,7 +977,7 @@ 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
@ -1014,22 +1016,27 @@ def handleWorkspaceSymbol (p : WorkspaceSymbolParams) : ReaderT ReferenceRequest
  if p.query.isEmpty then
    return #[]
  let references := (← read).references
-  let filterMapMod mod := documentUriFromModule? mod
  let filterMapIdent ident := do
-    let ident := privateToUserName? ident |>.getD ident
-    if let some score := fuzzyMatchScoreWithThreshold? p.query ident.toString then
-      return some (ident.toString, score)
+    let ident := privateToUserName? ident |>.getD ident |>.toString
+    let some score := FuzzyMatching.fuzzyMatchScoreWithThreshold? p.query ident
+      | return none
+    return some (ident, score)
+  let symbols ← references.definitionsMatching filterMapIdent
+  let ltSymbol symbol1 symbol2 :=
+    let (ident1, score1) := symbol1.ident
+    let (ident2, score2) := symbol2.ident
+    if score1 == score2 then
+      ident1.length < ident2.length
    else
-      return none
-  let symbols ← references.definitionsMatching filterMapMod filterMapIdent
-  return symbols
-    |>.qsort (fun { ident := (_, s1), .. } { ident := (_, s2), .. } => s1 > s2)
-    |>.extract 0 100 -- max amount
+      score1 > score2
+  let symbols := symbols.toList.mergeSort ltSymbol
+    |>.extract 0 1000
    |>.map fun m => {
      name := m.ident.1
      kind := SymbolKind.constant
-      location := { uri := m.mod, range := m.range }
+      location := { uri := m.modUri, range := m.range }
    }
+  return symbols.toArray

 def handlePrepareRename (p : PrepareRenameParams) : ReaderT ReferenceRequestContext IO (Option Range) := do
  -- This just checks that the cursor is over a renameable identifier
@ -1133,10 +1140,12 @@ section NotificationHandling
          continue
        try
          let ilean ← Ilean.load path
+          let some moduleUri ← documentUriFromModule? ilean.module
+            | continue
          if let FileChangeType.Changed := c.type then
-            modifyReferences (·.removeIlean path |>.addIlean path ilean)
+            modifyReferences (·.removeIlean path |>.addIlean moduleUri path ilean)
          else
-            modifyReferences (·.addIlean path ilean)
+            modifyReferences (·.addIlean moduleUri path ilean)
        catch
          -- ilean vanished, ignore error
          | .noFileOrDirectory .. => modifyReferences (·.removeIlean path)
@ -1501,8 +1510,10 @@ def startLoadingReferences (referenceData : Std.Mutex ReferenceData) : IO Unit :
    for path in ← oleanSearchPath.findAllWithExt "ilean" do
      try
        let ilean ← Ilean.load path
+        let some moduleUri ← documentUriFromModule? ilean.module
+          | continue
        referenceData.atomically <| modify fun rd =>
-          rd.modifyReferences (·.addIlean path ilean)
+          rd.modifyReferences (·.addIlean moduleUri path ilean)
      catch _ =>
        -- could be a race with the build system, for example
        -- ilean load errors should not be fatal, but we *should* log them