feat(lsp): Add codelens support

6 files changed, 343 insertions, 0 deletions

diff --git a/runtime/doc/lsp.txt b/runtime/doc/lsp.txt
index 531374620a..075912d4e5 100644
--- a/runtime/doc/lsp.txt
+++ b/runtime/doc/lsp.txt
@@ -1559,6 +1559,50 @@ show_line_diagnostics({opts}, {bufnr}, {line_nr}, {client_id})
 
 
 ==============================================================================
+Lua module: vim.lsp.codelens                                    *lsp-codelens*
+
+display({lenses}, {bufnr}, {client_id})           *vim.lsp.codelens.display()*
+                Display the lenses using virtual text
+
+                Parameters: ~
+                    {lenses}     table of lenses to display ( `CodeLens[] |
+                                 null` )
+                    {bufnr}      number
+                    {client_id}  number
+
+get({bufnr})                                          *vim.lsp.codelens.get()*
+                Return all lenses for the given buffer
+
+                Return: ~
+                    table ( `CodeLens[]` )
+
+                                              *vim.lsp.codelens.on_codelens()*
+on_codelens({err}, {_}, {result}, {client_id}, {bufnr})
+                |lsp-handler| for the method `textDocument/codeLens`
+
+refresh()                                         *vim.lsp.codelens.refresh()*
+                Refresh the codelens for the current buffer
+
+                It is recommended to trigger this using an autocmd or via
+                keymap.
+>
+    autocmd BufEnter,CursorHold,InsertLeave <buffer> lua vim.lsp.codelens.refresh()
+<
+
+run()                                                 *vim.lsp.codelens.run()*
+                Run the code lens in the current line
+
+save({lenses}, {bufnr}, {client_id})                 *vim.lsp.codelens.save()*
+                Store lenses for a specific buffer and client
+
+                Parameters: ~
+                    {lenses}     table of lenses to store ( `CodeLens[] |
+                                 null` )
+                    {bufnr}      number
+                    {client_id}  number
+
+
+==============================================================================
 Lua module: vim.lsp.handlers                                    *lsp-handlers*
 
                                          *vim.lsp.handlers.progress_handler()*
diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index 5a606188dd..75faf9bcc7 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -20,6 +20,7 @@ local lsp = {
 
   buf = require'vim.lsp.buf';
   diagnostic = require'vim.lsp.diagnostic';
+  codelens = require'vim.lsp.codelens';
   util = util;
 
   -- Allow raw RPC access.
diff --git a/runtime/lua/vim/lsp/codelens.lua b/runtime/lua/vim/lsp/codelens.lua
new file mode 100644
index 0000000000..fbd37e3830
--- /dev/null
+++ b/runtime/lua/vim/lsp/codelens.lua
@@ -0,0 +1,231 @@
+local util = require('vim.lsp.util')
+local api = vim.api
+local M = {}
+
+--- bufnr → true|nil
+--- to throttle refreshes to at most one at a time
+local active_refreshes = {}
+
+--- bufnr -> client_id -> lenses
+local lens_cache_by_buf = setmetatable({}, {
+  __index = function(t, b)
+    local key = b > 0 and b or api.nvim_get_current_buf()
+    return rawget(t, key)
+  end
+})
+
+local namespaces = setmetatable({}, {
+  __index = function(t, key)
+    local value = api.nvim_create_namespace('vim_lsp_codelens:' .. key)
+    rawset(t, key, value)
+    return value
+  end;
+})
+
+--@private
+M.__namespaces = namespaces
+
+
+--@private
+local function execute_lens(lens, bufnr, client_id)
+  local line = lens.range.start.line
+  api.nvim_buf_clear_namespace(bufnr, namespaces[client_id], line, line + 1)
+
+  -- Need to use the client that returned the lens → must not use buf_request
+  local client = vim.lsp.get_client_by_id(client_id)
+  assert(client, 'Client is required to execute lens, client_id=' .. client_id)
+  client.request('workspace/executeCommand', lens.command, function(...)
+    local result = vim.lsp.handlers['workspace/executeCommand'](...)
+    M.refresh()
+    return result
+  end, bufnr)
+end
+
+
+--- Return all lenses for the given buffer
+---
+---@return table (`CodeLens[]`)
+function M.get(bufnr)
+  local lenses_by_client = lens_cache_by_buf[bufnr]
+  if not lenses_by_client then return {} end
+  local lenses = {}
+  for _, client_lenses in pairs(lenses_by_client) do
+    vim.list_extend(lenses, client_lenses)
+  end
+  return lenses
+end
+
+
+--- Run the code lens in the current line
+---
+function M.run()
+  local line = api.nvim_win_get_cursor(0)[1]
+  local bufnr = api.nvim_get_current_buf()
+  local options = {}
+  local lenses_by_client = lens_cache_by_buf[bufnr] or {}
+  for client, lenses in pairs(lenses_by_client) do
+    for _, lens in pairs(lenses) do
+      if lens.range.start.line == (line - 1) then
+        table.insert(options, {client=client, lens=lens})
+      end
+    end
+  end
+  if #options == 0 then
+    vim.notify('No executable codelens found at current line')
+  elseif #options == 1 then
+    local option = options[1]
+    execute_lens(option.lens, bufnr, option.client)
+  else
+    local options_strings = {"Code lenses:"}
+    for i, option in ipairs(options) do
+       table.insert(options_strings, string.format('%d. %s', i, option.lens.command.title))
+    end
+    local choice = vim.fn.inputlist(options_strings)
+    if choice < 1 or choice > #options then
+      return
+    end
+    local option = options[choice]
+    execute_lens(option.lens, bufnr, option.client)
+  end
+end
+
+
+--- Display the lenses using virtual text
+---
+---@param lenses table of lenses to display (`CodeLens[] | null`)
+---@param bufnr number
+---@param client_id number
+function M.display(lenses, bufnr, client_id)
+  if not lenses or not next(lenses) then
+    return
+  end
+  local lenses_by_lnum = {}
+  for _, lens in pairs(lenses) do
+    local line_lenses = lenses_by_lnum[lens.range.start.line]
+    if not line_lenses then
+      line_lenses = {}
+      lenses_by_lnum[lens.range.start.line] = line_lenses
+    end
+    table.insert(line_lenses, lens)
+  end
+  local ns = namespaces[client_id]
+  local num_lines = api.nvim_buf_line_count(bufnr)
+  for i = 0, num_lines do
+    local line_lenses = lenses_by_lnum[i]
+    api.nvim_buf_clear_namespace(bufnr, ns, i, i + 1)
+    local chunks = {}
+    for _, lens in pairs(line_lenses or {}) do
+      local text = lens.command and lens.command.title or 'Unresolved lens ...'
+      table.insert(chunks, {text, 'LspCodeLens' })
+    end
+    if #chunks > 0 then
+      api.nvim_buf_set_virtual_text(bufnr, ns, i, chunks, {})
+    end
+  end
+end
+
+
+--- Store lenses for a specific buffer and client
+---
+---@param lenses table of lenses to store (`CodeLens[] | null`)
+---@param bufnr number
+---@param client_id number
+function M.save(lenses, bufnr, client_id)
+  local lenses_by_client = lens_cache_by_buf[bufnr]
+  if not lenses_by_client then
+    lenses_by_client = {}
+    lens_cache_by_buf[bufnr] = lenses_by_client
+    local ns = namespaces[client_id]
+    api.nvim_buf_attach(bufnr, false, {
+      on_detach = function(b) lens_cache_by_buf[b] = nil end,
+      on_lines = function(_, b, _, first_lnum, last_lnum)
+        api.nvim_buf_clear_namespace(b, ns, first_lnum, last_lnum)
+      end
+    })
+  end
+  lenses_by_client[client_id] = lenses
+end
+
+
+--@private
+local function resolve_lenses(lenses, bufnr, client_id, callback)
+  lenses = lenses or {}
+  local num_lens = vim.tbl_count(lenses)
+  if num_lens == 0 then
+    callback()
+    return
+  end
+
+  --@private
+  local function countdown()
+    num_lens = num_lens - 1
+    if num_lens == 0 then
+      callback()
+    end
+  end
+  local ns = namespaces[client_id]
+  local client = vim.lsp.get_client_by_id(client_id)
+  for _, lens in pairs(lenses or {}) do
+    if lens.command then
+      countdown()
+    else
+      client.request('codeLens/resolve', lens, function(_, _, result)
+        if result and result.command then
+          lens.command = result.command
+          -- Eager display to have some sort of incremental feedback
+          -- Once all lenses got resolved there will be a full redraw for all lenses
+          -- So that multiple lens per line are properly displayed
+          api.nvim_buf_set_virtual_text(
+            bufnr,
+            ns,
+            lens.range.start.line,
+            {{ lens.command.title, 'LspCodeLens' },},
+            {}
+          )
+        end
+        countdown()
+      end, bufnr)
+    end
+  end
+end
+
+
+--- |lsp-handler| for the method `textDocument/codeLens`
+---
+function M.on_codelens(err, _, result, client_id, bufnr)
+  assert(not err, vim.inspect(err))
+
+  M.save(result, bufnr, client_id)
+
+  -- Eager display for any resolved (and unresolved) lenses and refresh them
+  -- once resolved.
+  M.display(result, bufnr, client_id)
+  resolve_lenses(result, bufnr, client_id, function()
+    M.display(result, bufnr, client_id)
+    active_refreshes[bufnr] = nil
+  end)
+end
+
+
+--- Refresh the codelens for the current buffer
+---
+--- It is recommended to trigger this using an autocmd or via keymap.
+---
+--- <pre>
+---   autocmd BufEnter,CursorHold,InsertLeave <buffer> lua vim.lsp.codelens.refresh()
+--- </pre>
+---
+function M.refresh()
+  local params = {
+    textDocument = util.make_text_document_params()
+  }
+  local bufnr = api.nvim_get_current_buf()
+  if active_refreshes[bufnr] then
+    return
+  end
+  active_refreshes[bufnr] = true
+  vim.lsp.buf_request(0, 'textDocument/codeLens', params)
+end
+
+
+return M
diff --git a/runtime/lua/vim/lsp/handlers.lua b/runtime/lua/vim/lsp/handlers.lua
index 6ae54ea253..5d38150fc0 100644
--- a/runtime/lua/vim/lsp/handlers.lua
+++ b/runtime/lua/vim/lsp/handlers.lua
@@ -187,6 +187,10 @@ M['textDocument/publishDiagnostics'] = function(...)
   return require('vim.lsp.diagnostic').on_publish_diagnostics(...)
 end
 
+M['textDocument/codeLens'] = function(...)
+  return require('vim.lsp.codelens').on_codelens(...)
+end
+
 --@see https://microsoft.github.io/language-server-protocol/specifications/specification-current/#textDocument_references
 M['textDocument/references'] = function(_, _, result)
   if not result then return end
diff --git a/scripts/gen_vimdoc.py b/scripts/gen_vimdoc.py
index d46306d41a..18a2839702 100755
--- a/scripts/gen_vimdoc.py
+++ b/scripts/gen_vimdoc.py
@@ -154,6 +154,7 @@ CONFIG = {
             'lsp.lua',
             'buf.lua',
             'diagnostic.lua',
+            'codelens.lua',
             'handlers.lua',
             'util.lua',
             'log.lua',
diff --git a/test/functional/plugin/lsp/codelens_spec.lua b/test/functional/plugin/lsp/codelens_spec.lua
new file mode 100644
index 0000000000..e09d93f7cc
--- /dev/null
+++ b/test/functional/plugin/lsp/codelens_spec.lua
@@ -0,0 +1,62 @@
+local helpers = require('test.functional.helpers')(after_each)
+
+local exec_lua = helpers.exec_lua
+local eq = helpers.eq
+
+describe('vim.lsp.codelens', function()
+  before_each(function()
+    helpers.clear()
+    exec_lua('require("vim.lsp")')
+  end)
+  after_each(helpers.clear)
+
+  it('on_codelens_stores_and_displays_lenses', function()
+    local fake_uri = "file://fake/uri"
+    local bufnr = exec_lua([[
+      fake_uri = ...
+      local bufnr = vim.uri_to_bufnr(fake_uri)
+      local lines = {'So', 'many', 'lines'}
+      vim.fn.bufload(bufnr)
+      vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, lines)
+      return bufnr
+    ]], fake_uri)
+
+    exec_lua([[
+      local bufnr = ...
+      local lenses = {
+        {
+          range = {
+            start = { line = 0, character = 0, },
+            ['end'] = { line = 0, character = 0 }
+          },
+          command = { title = 'Lens1', command = 'Dummy' }
+        },
+      }
+      vim.lsp.codelens.on_codelens(nil, 'textDocument/codeLens', lenses, 1, bufnr)
+    ]], bufnr)
+
+    local stored_lenses = exec_lua('return vim.lsp.codelens.get(...)', bufnr)
+    local expected = {
+      {
+        range = {
+          start = { line = 0, character = 0 },
+          ['end'] = { line = 0, character = 0 }
+        },
+        command = {
+          title = 'Lens1',
+          command = 'Dummy',
+        },
+      },
+    }
+    eq(expected, stored_lenses)
+
+    local virtual_text_chunks = exec_lua([[
+      local bufnr = ...
+      local ns = vim.lsp.codelens.__namespaces[1]
+      local extmarks = vim.api.nvim_buf_get_extmarks(bufnr, ns, 0, -1, {})
+      return vim.api.nvim_buf_get_extmark_by_id(bufnr, ns, extmarks[1][1], { details = true })[3].virt_text
+    ]], bufnr)
+
+    eq({[1] = {'Lens1', 'LspCodeLens'}}, virtual_text_chunks)
+  end)
+end)