From 74718e51d6a9367bb70128f40acd26cc4c9e9070 Mon Sep 17 00:00:00 2001
From: BenYip <yebenmy@gmail.com>
Date: Tue, 21 Jan 2025 23:27:32 +0800
Subject: [PATCH] Support custom locations request (#587)

* Support custom locations request

* specify server name in FindLocations

* Rename LspFindLocations
---
 autoload/lsp/buffer.vim    | 11 +++++++++++
 autoload/lsp/lsp.vim       | 18 ++++++++++++++++++
 autoload/lsp/lspserver.vim | 24 ++++++++++++++++++++++++
 3 files changed, 53 insertions(+)

diff --git a/autoload/lsp/buffer.vim b/autoload/lsp/buffer.vim
index 5a07654..f48c3b1 100644
--- a/autoload/lsp/buffer.vim
+++ b/autoload/lsp/buffer.vim
@@ -191,4 +191,15 @@ export def CurbufGetServerChecked(feature: string = null_string): dict<any>
   return lspserver
 enddef
 
+export def CurbufGetServerByName(name: string): dict<any>
+  var lspservers: list<dict<any>> = CurbufGetServers()
+
+  for lspserver in lspservers
+    if lspserver.name == name
+      return lspserver
+    endif
+  endfor
+  return {}
+enddef
+
 # vim: tabstop=8 shiftwidth=2 softtabstop=2
diff --git a/autoload/lsp/lsp.vim b/autoload/lsp/lsp.vim
index 1ed7b8d..69029f2 100644
--- a/autoload/lsp/lsp.vim
+++ b/autoload/lsp/lsp.vim
@@ -832,6 +832,24 @@ export def ShowReferences(peek: bool)
   lspserver.showReferences(peek)
 enddef
 
+# send custom locations request
+def g:LspFindLocations(server_name: string, peek: bool, method: string, args: dict<any> = {})
+  var lspserver: dict<any> = buf.CurbufGetServerByName(server_name)
+  if lspserver->empty()
+    return
+  endif
+  if !lspserver.running
+    util.ErrMsg($'Language server "{server_name}" is not running')
+    return
+  endif
+  if !lspserver.ready
+    util.ErrMsg($'Language server "{server_name}" is not ready')
+    return
+  endif
+
+  lspserver.findLocations(peek, method, args)
+enddef
+
 # highlight all the places where a symbol is referenced
 def g:LspDocHighlight(bnr: number = bufnr(), cmdmods: string = '')
   var lspserver: dict<any> = buf.CurbufGetServerChecked('documentHighlight')
diff --git a/autoload/lsp/lspserver.vim b/autoload/lsp/lspserver.vim
index 1d4b0ed..11e8d67 100644
--- a/autoload/lsp/lspserver.vim
+++ b/autoload/lsp/lspserver.vim
@@ -1006,6 +1006,29 @@ def ShowReferences(lspserver: dict<any>, peek: bool): void
   symbol.ShowLocations(lspserver, reply.result, peek, 'Symbol References')
 enddef
 
+# send custom locations request
+def FindLocations(lspserver: dict<any>, peek: bool, method: string, args: dict<any>): void
+  var param: dict<any>
+  param = lspserver.getTextDocPosition(true)->extend(args)
+  var reply = lspserver.rpc(method, param)
+
+  # Result: Location[] | null
+  if reply->empty() || reply.result->empty()
+    util.WarnMsg('No references found')
+    return
+  endif
+
+  if lspserver.needOffsetEncoding
+    # Decode the position encoding in all the reference locations
+    reply.result->map((_, loc) => {
+      lspserver.decodeLocation(loc)
+      return loc
+    })
+  endif
+
+  symbol.ShowLocations(lspserver, reply.result, peek, 'Symbol References')
+enddef
+
 # process the 'textDocument/documentHighlight' reply from the LSP server
 # Result: DocumentHighlight[] | null
 def DocHighlightReply(lspserver: dict<any>, docHighlightReply: any,
@@ -1948,6 +1971,7 @@ export def NewLspServer(serverParams: dict<any>): dict<any>
     didSaveFile: function(DidSaveFile, [lspserver]),
     hover: function(ShowHoverInfo, [lspserver]),
     showReferences: function(ShowReferences, [lspserver]),
+    findLocations: function(FindLocations, [lspserver]),
     docHighlight: function(DocHighlight, [lspserver]),
     getDocSymbols: function(GetDocSymbols, [lspserver]),
     textDocFormat: function(TextDocFormat, [lspserver]),
-- 
2.48.1