From efcf94298a6ed24a312305fbebad244e0ea957ea Mon Sep 17 00:00:00 2001 From: Marc Huisinga Date: Tue, 13 May 2025 10:29:49 +0200 Subject: [PATCH] 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`. --- .../Completion/CompletionCollectors.lean | 30 +---- .../Server/Completion/CompletionUtils.lean | 18 +++ src/Lean/Server/References.lean | 124 ++++++++++++------ src/Lean/Server/Watchdog.lean | 53 +++++--- 4 files changed, 142 insertions(+), 83 deletions(-) diff --git a/src/Lean/Server/Completion/CompletionCollectors.lean b/src/Lean/Server/Completion/CompletionCollectors.lean index 04d0bfd2f0..806c1b924d 100644 --- 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 diff --git a/src/Lean/Server/Completion/CompletionUtils.lean b/src/Lean/Server/Completion/CompletionUtils.lean index 4f260b1c0a..5eb6e6213a 100644 --- 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 diff --git a/src/Lean/Server/References.lean b/src/Lean/Server/References.lean index f9e52cc430..593e4de49f 100644 --- 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 diff --git a/src/Lean/Server/Watchdog.lean b/src/Lean/Server/Watchdog.lean index 2ffa9ef04e..892849761d 100644 --- 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