From 7233bb2ec07506b6a6e57dfe4541f1c4e5647fd2 Mon Sep 17 00:00:00 2001
From: Satoru Kitaguchi <64164948+satorunooshie@users.noreply.github.com>
Date: Sat, 9 Sep 2023 14:25:51 +0900
Subject: [PATCH] Add LspAddTreeReferences (#1484)

---
 autoload/lsp/ui/vim.vim | 20 ++++++++++++++++++--
 doc/vim-lsp.txt         | 10 +++++++++-
 plugin/lsp.vim          |  5 +++--
 3 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/autoload/lsp/ui/vim.vim b/autoload/lsp/ui/vim.vim
index 785db3ec7..f986f53b9 100644
--- a/autoload/lsp/ui/vim.vim
+++ b/autoload/lsp/ui/vim.vim
@@ -35,12 +35,17 @@ function! lsp#ui#vim#definition(in_preview, ...) abort
     call s:list_location('definition', l:ctx)
 endfunction
 
-function! lsp#ui#vim#references() abort
-    let l:ctx = { 'jump_if_one': 0 }
+function! lsp#ui#vim#references(ctx) abort
+    let l:ctx = extend({ 'jump_if_one': 0 }, a:ctx)
     let l:request_params = { 'context': { 'includeDeclaration': v:false } }
     call s:list_location('references', l:ctx, l:request_params)
 endfunction
 
+function! lsp#ui#vim#add_tree_references() abort
+    let l:ctx = { 'add_tree': v:true }
+    call lsp#ui#vim#references(l:ctx)
+endfunction
+
 function! s:list_location(method, ctx, ...) abort
     " typeDefinition => type definition
     let l:operation = substitute(a:method, '\u', ' \l\0', 'g')
@@ -307,10 +312,21 @@ function! s:handle_location(ctx, server, type, data) abort "ctx = {counter, list
                 echo 'Retrieved ' . a:type
                 redraw
             elseif !a:ctx['in_preview']
+                if get(a:ctx, 'add_tree', v:false)
+                    let l:qf = getqflist({'idx' : 0, 'items': []})
+                    let l:pos = l:qf.idx
+                    let l:parent = l:qf.items
+                    let l:level = count(l:parent[l:pos-1].text, g:lsp_tree_incoming_prefix)
+                    let a:ctx['list'] = extend(l:parent, map(a:ctx['list'], 'extend(v:val, {"text": repeat("' . g:lsp_tree_incoming_prefix . '", l:level+1) . v:val.text})'), l:pos)
+                endif
                 call setqflist([])
                 call setqflist(a:ctx['list'])
                 echo 'Retrieved ' . a:type
                 botright copen
+                if get(a:ctx, 'add_tree', v:false)
+                    " move the cursor to the newly added item
+                    execute l:pos + 1
+                endif
             else
                 let l:lines = readfile(l:loc['filename'])
                 if has_key(l:loc,'viewstart') " showing a locationLink
diff --git a/doc/vim-lsp.txt b/doc/vim-lsp.txt
index e8bf2659b..2e720a540 100644
--- a/doc/vim-lsp.txt
+++ b/doc/vim-lsp.txt
@@ -114,6 +114,7 @@ CONTENTS                                                  *vim-lsp-contents*
       lsp#document_hover_preview_winid()  |lsp#document_hover_preview_winid()|
     Commands                              |vim-lsp-commands|
       LspAddTreeCallHierarchyIncoming       |:LspAddTreeCallHierarchyIncoming|
+      LspAddTreeReferences                  |:LspAddTreeReferences|
       LspCallHierarchyIncoming              |:LspCallHierarchyIncoming|
       LspCallHierarchyOutgoing              |:LspCallHierarchyOutgoing|
       LspCodeAction                         |:LspCodeAction|
@@ -877,7 +878,9 @@ g:lsp_tree_incoming_prefix                          *g:lsp_tree_incoming_prefix*
     Type: |String|
     Default: `"<= "`
 
-    Specifies the prefix of items added by |LspAddTreeCallHierarchyIncoming|.
+    Specifies the prefix of items added by following commands.
+	* |LspAddTreeCallHierarchyIncoming|
+	* |LspAddTreeReferences|
 
     Example: >
 	let g:lsp_tree_incoming_prefix = "← "
@@ -1680,6 +1683,11 @@ LspAddTreeCallHierarchyIncoming           *:LspAddTreeCallHierarchyIncoming*
 Just like |LspCallHierarchyIncoming| , but instead of making a new list the
 result is appended to the current list.
 
+LspAddTreeReferences                                 *:LspAddTreeReferences*
+
+Just like |LspReferences| , but instead of making a new list the result is
+appended to the current list.
+
 LspCallHierarchyIncoming                         *:LspCallHierarchyIncoming*
 
 Find incoming call hierarchy for the symbol under cursor.
diff --git a/plugin/lsp.vim b/plugin/lsp.vim
index 0485b2b62..11a2b4428 100644
--- a/plugin/lsp.vim
+++ b/plugin/lsp.vim
@@ -128,7 +128,8 @@ command! -nargs=* LspNextWarning call lsp#internal#diagnostics#movement#_next_wa
 command! -nargs=* LspPreviousWarning call lsp#internal#diagnostics#movement#_previous_warning(<f-args>)
 command! -nargs=* LspNextDiagnostic call lsp#internal#diagnostics#movement#_next_diagnostics(<f-args>)
 command! -nargs=* LspPreviousDiagnostic call lsp#internal#diagnostics#movement#_previous_diagnostics(<f-args>)
-command! LspReferences call lsp#ui#vim#references()
+command! LspReferences call lsp#ui#vim#references({})
+command! LspAddTreeReferences call lsp#ui#vim#add_tree_references()
 command! LspRename call lsp#ui#vim#rename()
 command! LspTypeDefinition call lsp#ui#vim#type_definition(0, <q-mods>)
 command! LspTypeHierarchy call lsp#internal#type_hierarchy#show()
@@ -189,7 +190,7 @@ nnoremap <silent> <plug>(lsp-next-diagnostic) :<c-u>call lsp#internal#diagnostic
 nnoremap <silent> <plug>(lsp-next-diagnostic-nowrap) :<c-u>call lsp#internal#diagnostics#movement#_next_diagnostics("-wrap=0")<cr>
 nnoremap <silent> <plug>(lsp-previous-diagnostic) :<c-u>call lsp#internal#diagnostics#movement#_previous_diagnostics()<cr>
 nnoremap <silent> <plug>(lsp-previous-diagnostic-nowrap) :<c-u>call lsp#internal#diagnostics#movement#_previous_diagnostics("-wrap=0")<cr>
-nnoremap <silent> <plug>(lsp-references) :<c-u>call lsp#ui#vim#references()<cr>
+nnoremap <silent> <plug>(lsp-references) :<c-u>call lsp#ui#vim#references({})<cr>
 nnoremap <silent> <plug>(lsp-rename) :<c-u>call lsp#ui#vim#rename()<cr>
 nnoremap <silent> <plug>(lsp-type-definition) :<c-u>call lsp#ui#vim#type_definition(0)<cr>
 nnoremap <silent> <plug>(lsp-type-hierarchy) :<c-u>call lsp#internal#type_hierarchy#show()<cr>