fix: handle overlapping definitions

src/Lean/Server/FileWorker/RequestHandling.lean

@@ -233,14 +233,14 @@ partial def handleDocumentHighlight (p : DocumentHighlightParams)
   let highlightRefs? (snaps : Array Snapshot) (pos : Lsp.Position) : Option (Array DocumentHighlight) := Id.run do
     let trees := snaps.map (·.infoTree)
     let refs := findModuleRefs text trees
-    let some ident ← refs.findAt? p.position
-      | none
-    let some info ← refs.find? ident
-      | none
     let mut ranges := #[]
-    if let some definition := info.definition then
-      ranges := ranges.push definition
-    ranges := ranges.append info.usages
+    for ident in ← refs.findAt p.position do
+      if let some info ← refs.find? ident then
+        if let some definition := info.definition then
+          ranges := ranges.push definition
+        ranges := ranges.append info.usages
+    if ranges.isEmpty then
+      return none
     some <| ranges.map ({ range := ·, kind? := DocumentHighlightKind.text })
 
   withWaitFindSnap doc (fun s => s.endPos > pos)

src/Lean/Server/References.lean

@@ -65,11 +65,12 @@ def addRef (self : ModuleRefs) (ref : Reference) : ModuleRefs :=
   let refInfo := self.findD ref.ident RefInfo.empty
   self.insert ref.ident (refInfo.addRef ref)
 
-def findAt? (self : ModuleRefs) (pos : Lsp.Position) : Option RefIdent := Id.run do
+def findAt (self : ModuleRefs) (pos : Lsp.Position) : Array RefIdent := Id.run do
+  let mut result := #[]
   for (ident, info) in self.toList do
     if info.contains pos then
-      return some ident
-  none
+      result := result.push ident
+  result
 
 end Lean.Lsp.ModuleRefs
 
@@ -215,10 +216,10 @@ def allRefs (self : References) : HashMap Name ModuleRefs :=
   let ileanRefs := self.ileans.toList.foldl (init := HashMap.empty) fun m (name, _, refs) => m.insert name refs
   self.workers.toList.foldl (init := ileanRefs) fun m (name, _, refs) => m.insert name refs
 
-def findAt? (self : References) (module : Name) (pos : Lsp.Position) : Option RefIdent := Id.run do
+def findAt (self : References) (module : Name) (pos : Lsp.Position) : Array RefIdent := Id.run do
   if let some refs := self.allRefs.find? module then
-    return refs.findAt? pos
-  none
+    return refs.findAt pos
+  #[]
 
 def referringTo (self : References) (ident : RefIdent) (srcSearchPath : SearchPath)
     (includeDefinition : Bool := true) : IO (Array Location) := do

src/Lean/Server/Watchdog.lean

@@ -368,23 +368,27 @@ end ServerM
 
 section RequestHandling
 
-def findDefinition? (p : TextDocumentPositionParams) : ServerM <| Option Location := do
+def findDefinitions (p : TextDocumentPositionParams) : ServerM <| Array Location := do
+  let mut definitions := #[]
   if let some path := p.textDocument.uri.toPath? then
     let srcSearchPath := (← read).srcSearchPath
     if let some module ← searchModuleNameOfFileName path srcSearchPath then
       let references ← (← read).references.get
-      if let some ident := references.findAt? module p.position then
-        return ← references.definitionOf? ident srcSearchPath
-  return none
+      for ident in references.findAt module p.position do
+        if let some definition ← references.definitionOf? ident srcSearchPath then
+          definitions := definitions.push definition
+  return definitions
 
 def handleReference (p : ReferenceParams) : ServerM (Array Location) := do
+  let mut result := #[]
   if let some path := p.textDocument.uri.toPath? then
     let srcSearchPath := (← read).srcSearchPath
     if let some module ← searchModuleNameOfFileName path srcSearchPath then
       let references ← (← read).references.get
-      if let some ident := references.findAt? module p.position then
-        return ← references.referringTo ident srcSearchPath p.context.includeDeclaration
-  return #[]
+      for ident in references.findAt module p.position do
+        let identRefs ← references.referringTo ident srcSearchPath p.context.includeDeclaration
+        result := result.append identRefs
+  return result
 
 -- TODO Better matching https://github.com/leanprover/lean4/issues/960
 def handleWorkspaceSymbol (p : WorkspaceSymbolParams) : ServerM (Array SymbolInformation) := do
@@ -543,8 +547,9 @@ section MessageHandling
     -- go-to-type-definition.
     if method == "textDocument/definition" || method == "textDocument/declaration" then
       let params ← parseParams TextDocumentPositionParams params
-      if let some definition ← findDefinition? params then
-        (← read).hOut.writeLspResponse ⟨id, #[definition]⟩
+      let definitions ← findDefinitions params
+      if !definitions.isEmpty then
+        (← read).hOut.writeLspResponse ⟨id, definitions⟩
         return
     match method with
       | "textDocument/references" => handle ReferenceParams (Array Location) handleReference