From 643546b82b4bc0c29ca869f81af868a019723d83 Mon Sep 17 00:00:00 2001
From: Chinmay Dalal <dalal.chinmay.0101@gmail.com>
Date: Sun, 11 Jun 2023 15:23:37 +0530
Subject: [PATCH] feat(lsp): add handlers for inlay hints (#23736)

initial support; public API left for a follow-up PR
---
 runtime/doc/lsp.txt                           |   2 +
 runtime/doc/news.txt                          |   2 +
 runtime/lua/vim/lsp.lua                       |  17 +-
 runtime/lua/vim/lsp/_inlay_hint.lua           | 217 ++++++++++++++++++
 runtime/lua/vim/lsp/handlers.lua              |  26 +++
 runtime/lua/vim/lsp/protocol.lua              |   9 +
 runtime/lua/vim/lsp/types.lua                 |  25 ++
 scripts/gen_vimdoc.py                         |   1 +
 .../functional/plugin/lsp/inlay_hint_spec.lua | 132 +++++++++++
 test/functional/plugin/lsp_spec.lua           |   4 +-
 10 files changed, 429 insertions(+), 6 deletions(-)
 create mode 100644 runtime/lua/vim/lsp/_inlay_hint.lua
 create mode 100644 test/functional/plugin/lsp/inlay_hint_spec.lua

diff --git a/runtime/doc/lsp.txt b/runtime/doc/lsp.txt
index 7248d03196..27ef38743b 100644
--- a/runtime/doc/lsp.txt
+++ b/runtime/doc/lsp.txt
@@ -197,6 +197,7 @@ specification. These LSP requests/notifications are defined by default:
     textDocument/formatting
     textDocument/hover
     textDocument/implementation*
+    textDocument/inlayHint
     textDocument/publishDiagnostics
     textDocument/rangeFormatting
     textDocument/references
@@ -208,6 +209,7 @@ specification. These LSP requests/notifications are defined by default:
     window/showDocument
     window/showMessageRequest
     workspace/applyEdit
+    workspace/inlayHint/refresh
     workspace/symbol
 
 * NOTE: These are sometimes not implemented by servers.
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 20d20d9a90..ac9e60637a 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -84,6 +84,8 @@ The following new APIs or features were added.
 
 • |nvim_set_keymap()| and |nvim_del_keymap()| now support abbreviations.
 
+• Added |lsp-handler| for inlay hints: `textDocument/inlayHint` and
+  `workspace/inlayHint/refresh`
 ==============================================================================
 CHANGED FEATURES                                                 *news-changed*
 
diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index 1d9a91801a..6f9a6c460b 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -17,6 +17,7 @@ local if_nil = vim.F.if_nil
 
 local lsp = {
   protocol = protocol,
+  _inlay_hint = require('vim.lsp._inlay_hint'),
 
   handlers = default_handlers,
 
@@ -60,6 +61,8 @@ lsp._request_name_to_capability = {
   ['textDocument/documentHighlight'] = { 'documentHighlightProvider' },
   ['textDocument/semanticTokens/full'] = { 'semanticTokensProvider' },
   ['textDocument/semanticTokens/full/delta'] = { 'semanticTokensProvider' },
+  ['textDocument/inlayHint'] = { 'inlayHintProvider' },
+  ['inlayHint/resolve'] = { 'inlayHintProvider', 'resolveProvider' },
 }
 
 -- TODO improve handling of scratch buffers with LSP attached.
@@ -1498,16 +1501,20 @@ function lsp.start_client(config)
     end
     -- Ensure pending didChange notifications are sent so that the server doesn't operate on a stale state
     changetracking.flush(client, bufnr)
+    local version = util.buf_versions[bufnr]
     bufnr = resolve_bufnr(bufnr)
     if log.debug() then
       log.debug(log_prefix, 'client.request', client_id, method, params, handler, bufnr)
     end
     local success, request_id = rpc.request(method, params, function(err, result)
-      handler(
-        err,
-        result,
-        { method = method, client_id = client_id, bufnr = bufnr, params = params }
-      )
+      local context = {
+        method = method,
+        client_id = client_id,
+        bufnr = bufnr,
+        params = params,
+        version = version,
+      }
+      handler(err, result, context)
     end, function(request_id)
       local request = client.requests[request_id]
       request.type = 'complete'
diff --git a/runtime/lua/vim/lsp/_inlay_hint.lua b/runtime/lua/vim/lsp/_inlay_hint.lua
new file mode 100644
index 0000000000..aa6ec9aca8
--- /dev/null
+++ b/runtime/lua/vim/lsp/_inlay_hint.lua
@@ -0,0 +1,217 @@
+local util = require('vim.lsp.util')
+local log = require('vim.lsp.log')
+local api = vim.api
+local M = {}
+
+---@class lsp._inlay_hint.bufstate
+---@field version integer
+---@field client_hint table<integer, table<integer, lsp.InlayHint[]>> client_id -> (lnum -> hints)
+
+---@type table<integer, lsp._inlay_hint.bufstate>
+local hint_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 namespace = api.nvim_create_namespace('vim_lsp_inlayhint')
+
+M.__explicit_buffers = {}
+
+--- |lsp-handler| for the method `textDocument/inlayHint`
+--- Store hints for a specific buffer and client
+--- Resolves unresolved hints
+---@private
+function M.on_inlayhint(err, result, ctx, _)
+  if err then
+    local _ = log.error() and log.error('inlayhint', err)
+    return
+  end
+  local bufnr = ctx.bufnr
+  if util.buf_versions[bufnr] ~= ctx.version then
+    return
+  end
+  local client_id = ctx.client_id
+  if not result then
+    return
+  end
+  local bufstate = hint_cache_by_buf[bufnr]
+  if not bufstate then
+    bufstate = {
+      client_hint = vim.defaulttable(),
+      version = ctx.version,
+    }
+    hint_cache_by_buf[bufnr] = bufstate
+    api.nvim_buf_attach(bufnr, false, {
+      on_detach = function(_, b)
+        api.nvim_buf_clear_namespace(b, namespace, 0, -1)
+        hint_cache_by_buf[b] = nil
+      end,
+      on_reload = function(_, b)
+        api.nvim_buf_clear_namespace(b, namespace, 0, -1)
+        hint_cache_by_buf[b] = nil
+      end,
+    })
+  end
+  local hints_by_client = bufstate.client_hint
+  local client = vim.lsp.get_client_by_id(client_id)
+
+  local new_hints_by_lnum = vim.defaulttable()
+  local num_unprocessed = #result
+  if num_unprocessed == 0 then
+    hints_by_client[client_id] = {}
+    bufstate.version = ctx.version
+    api.nvim__buf_redraw_range(bufnr, 0, -1)
+    return
+  end
+
+  local lines = api.nvim_buf_get_lines(bufnr, 0, -1, false)
+  ---@private
+  local function pos_to_byte(position)
+    local col = position.character
+    if col > 0 then
+      local line = lines[position.line + 1] or ''
+      local ok, convert_result
+      ok, convert_result = pcall(util._str_byteindex_enc, line, col, client.offset_encoding)
+      if ok then
+        return convert_result
+      end
+      return math.min(#line, col)
+    end
+    return col
+  end
+
+  for _, hint in ipairs(result) do
+    local lnum = hint.position.line
+    hint.position.character = pos_to_byte(hint.position)
+    table.insert(new_hints_by_lnum[lnum], hint)
+  end
+
+  hints_by_client[client_id] = new_hints_by_lnum
+  bufstate.version = ctx.version
+  api.nvim__buf_redraw_range(bufnr, 0, -1)
+end
+
+---@private
+local function resolve_bufnr(bufnr)
+  return bufnr == 0 and api.nvim_get_current_buf() or bufnr
+end
+
+--- Refresh inlay hints for a buffer
+---
+--- It is recommended to trigger this using an autocmd or via keymap.
+---@param opts (nil|table) Optional arguments
+---  - bufnr (integer, default: 0): Buffer whose hints to refresh
+---  - only_visible (boolean, default: false): Whether to only refresh hints for the visible regions of the buffer
+---
+--- Example:
+--- <pre>vim
+---   autocmd BufEnter,InsertLeave,BufWritePost <buffer> lua vim.lsp._inlay_hint.refresh()
+--- </pre>
+---
+---@private
+function M.refresh(opts)
+  opts = opts or {}
+  local bufnr = opts.bufnr or 0
+  local only_visible = opts.only_visible or false
+  bufnr = resolve_bufnr(bufnr)
+  M.__explicit_buffers[bufnr] = true
+  local buffer_windows = {}
+  for _, winid in ipairs(api.nvim_list_wins()) do
+    if api.nvim_win_get_buf(winid) == bufnr then
+      table.insert(buffer_windows, winid)
+    end
+  end
+  for _, window in ipairs(buffer_windows) do
+    local first = vim.fn.line('w0', window)
+    local last = vim.fn.line('w$', window)
+    local params = {
+      textDocument = util.make_text_document_params(bufnr),
+      range = {
+        start = { line = first - 1, character = 0 },
+        ['end'] = { line = last, character = 0 },
+      },
+    }
+    vim.lsp.buf_request(bufnr, 'textDocument/inlayHint', params)
+  end
+  if not only_visible then
+    local params = {
+      textDocument = util.make_text_document_params(bufnr),
+      range = {
+        start = { line = 0, character = 0 },
+        ['end'] = { line = api.nvim_buf_line_count(bufnr), character = 0 },
+      },
+    }
+    vim.lsp.buf_request(bufnr, 'textDocument/inlayHint', params)
+  end
+end
+
+--- Clear inlay hints
+---
+---@param client_id integer|nil filter by client_id. All clients if nil
+---@param bufnr integer|nil filter by buffer. All buffers if nil
+---@private
+function M.clear(client_id, bufnr)
+  local buffers = bufnr and { resolve_bufnr(bufnr) } or vim.tbl_keys(hint_cache_by_buf)
+  for _, iter_bufnr in ipairs(buffers) do
+    M.__explicit_buffers[iter_bufnr] = false
+    local bufstate = hint_cache_by_buf[iter_bufnr]
+    local client_lens = (bufstate or {}).client_hint or {}
+    local client_ids = client_id and { client_id } or vim.tbl_keys(client_lens)
+    for _, iter_client_id in ipairs(client_ids) do
+      if bufstate then
+        bufstate.client_hint[iter_client_id] = {}
+      end
+    end
+    api.nvim_buf_clear_namespace(iter_bufnr, namespace, 0, -1)
+  end
+  vim.cmd('redraw!')
+end
+
+api.nvim_set_decoration_provider(namespace, {
+  on_win = function(_, _, bufnr, topline, botline)
+    local bufstate = hint_cache_by_buf[bufnr]
+    if not bufstate then
+      return
+    end
+
+    if bufstate.version ~= util.buf_versions[bufnr] then
+      return
+    end
+    local hints_by_client = bufstate.client_hint
+    api.nvim_buf_clear_namespace(bufnr, namespace, 0, -1)
+
+    for lnum = topline, botline do
+      for _, hints_by_lnum in pairs(hints_by_client) do
+        local line_hints = hints_by_lnum[lnum] or {}
+        for _, hint in pairs(line_hints) do
+          local text = ''
+          if type(hint.label) == 'string' then
+            text = hint.label
+          else
+            for _, part in ipairs(hint.label) do
+              text = text .. part.value
+            end
+          end
+          if hint.paddingLeft then
+            text = ' ' .. text
+          end
+          if hint.paddingRight then
+            text = text .. ' '
+          end
+          api.nvim_buf_set_extmark(bufnr, namespace, lnum, hint.position.character, {
+            virt_text_pos = 'inline',
+            ephemeral = false,
+            virt_text = {
+              { text, 'LspInlayHint' },
+            },
+            hl_mode = 'combine',
+          })
+        end
+      end
+    end
+  end,
+})
+
+return M
diff --git a/runtime/lua/vim/lsp/handlers.lua b/runtime/lua/vim/lsp/handlers.lua
index 19338ae8f0..44a9a58aca 100644
--- a/runtime/lua/vim/lsp/handlers.lua
+++ b/runtime/lua/vim/lsp/handlers.lua
@@ -219,6 +219,10 @@ M['textDocument/codeLens'] = function(...)
   return require('vim.lsp.codelens').on_codelens(...)
 end
 
+M['textDocument/inlayHint'] = function(...)
+  return require('vim.lsp._inlay_hint').on_inlayhint(...)
+end
+
 --see: https://microsoft.github.io/language-server-protocol/specifications/specification-current/#textDocument_references
 M['textDocument/references'] = function(_, result, ctx, config)
   if not result or vim.tbl_isempty(result) then
@@ -612,6 +616,28 @@ M['window/showDocument'] = function(_, result, ctx, _)
   return { success = success or false }
 end
 
+---@see https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#workspace_inlayHint_refresh
+M['workspace/inlayHint/refresh'] = function(err, _, ctx)
+  local inlay_hint = require('vim.lsp._inlay_hint')
+  if not inlay_hint.__explicit_buffers[ctx.bufnr] then
+    return vim.NIL
+  end
+  if err then
+    return vim.NIL
+  end
+
+  for _, bufnr in ipairs(vim.lsp.get_buffers_by_client_id(ctx.client_id)) do
+    for _, winid in ipairs(api.nvim_list_wins()) do
+      if api.nvim_win_get_buf(winid) == bufnr then
+        inlay_hint.refresh({ bufnr = bufnr })
+        break
+      end
+    end
+  end
+
+  return vim.NIL
+end
+
 -- Add boilerplate error validation and logging for all of these.
 for k, fn in pairs(M) do
   M[k] = function(err, result, ctx, config)
diff --git a/runtime/lua/vim/lsp/protocol.lua b/runtime/lua/vim/lsp/protocol.lua
index 172d43e483..b3a7903420 100644
--- a/runtime/lua/vim/lsp/protocol.lua
+++ b/runtime/lua/vim/lsp/protocol.lua
@@ -641,6 +641,12 @@ function protocol.make_client_capabilities()
       },
     },
     textDocument = {
+      inlayHint = {
+        dynamicRegistration = false,
+        resolveSupport = {
+          properties = {},
+        },
+      },
       semanticTokens = {
         dynamicRegistration = false,
         tokenTypes = {
@@ -853,6 +859,9 @@ function protocol.make_client_capabilities()
         dynamicRegistration = true,
         relativePatternSupport = true,
       },
+      inlayHint = {
+        refreshSupport = true,
+      },
     },
     experimental = nil,
     window = {
diff --git a/runtime/lua/vim/lsp/types.lua b/runtime/lua/vim/lsp/types.lua
index ef85a0d10f..108aeeb922 100644
--- a/runtime/lua/vim/lsp/types.lua
+++ b/runtime/lua/vim/lsp/types.lua
@@ -69,3 +69,28 @@
 --- @field method string
 
 --- @alias lsp.UnregistrationParams {unregisterations: lsp.Unregistration[]}
+
+---@class lsp.Location
+---@field uri string
+---@field range lsp.Range
+
+---@class lsp.MarkupContent
+---@field kind string
+---@field value string
+
+---@class lsp.InlayHintLabelPart
+---@field value string
+---@field tooltip? string | lsp.MarkupContent
+---@field location? lsp.Location
+
+---@class lsp.TextEdit
+---@field range lsp.Range
+---@field newText string
+
+---@class lsp.InlayHint
+---@field position lsp.Position
+---@field label string | lsp.InlayHintLabelPart[]
+---@field kind? integer
+---@field textEdits? lsp.TextEdit[]
+---@field paddingLeft? boolean
+---@field paddingRight? boolean
diff --git a/scripts/gen_vimdoc.py b/scripts/gen_vimdoc.py
index b40f8526ea..2d4c8b2234 100755
--- a/scripts/gen_vimdoc.py
+++ b/scripts/gen_vimdoc.py
@@ -223,6 +223,7 @@ CONFIG = {
             'log.lua',
             'rpc.lua',
             'protocol.lua',
+            'inlay_hint.lua'
         ],
         'files': [
             'runtime/lua/vim/lsp',
diff --git a/test/functional/plugin/lsp/inlay_hint_spec.lua b/test/functional/plugin/lsp/inlay_hint_spec.lua
new file mode 100644
index 0000000000..b134095c4f
--- /dev/null
+++ b/test/functional/plugin/lsp/inlay_hint_spec.lua
@@ -0,0 +1,132 @@
+local helpers = require('test.functional.helpers')(after_each)
+local lsp_helpers = require('test.functional.plugin.lsp.helpers')
+local Screen = require('test.functional.ui.screen')
+
+local dedent = helpers.dedent
+local exec_lua = helpers.exec_lua
+local insert = helpers.insert
+
+local clear_notrace = lsp_helpers.clear_notrace
+local create_server_definition = lsp_helpers.create_server_definition
+
+before_each(function()
+  clear_notrace()
+end)
+
+after_each(function()
+  exec_lua("vim.api.nvim_exec_autocmds('VimLeavePre', { modeline = false })")
+end)
+
+describe('inlay hints', function()
+  local screen
+  before_each(function()
+    screen = Screen.new(50, 9)
+    screen:attach()
+  end)
+
+  describe('general', function()
+    local text = dedent([[
+    auto add(int a, int b) { return a + b; }
+
+    int main() {
+        int x = 1;
+        int y = 2;
+        return add(x,y);
+    }
+    }]])
+
+
+    local response = [==[
+    [
+    {"kind":1,"paddingLeft":false,"label":"-> int","position":{"character":22,"line":0},"paddingRight":false},
+    {"kind":2,"paddingLeft":false,"label":"a:","position":{"character":15,"line":5},"paddingRight":true},
+    {"kind":2,"paddingLeft":false,"label":"b:","position":{"character":17,"line":5},"paddingRight":true}
+    ]
+    ]==]
+
+
+    before_each(function()
+      exec_lua(create_server_definition)
+      exec_lua([[
+        local response = ...
+        server = _create_server({
+          capabilities = {
+            inlayHintProvider = true,
+          },
+          handlers = {
+            ['textDocument/inlayHint'] = function()
+              return vim.json.decode(response)
+            end,
+          }
+        })
+      ]], response)
+    end)
+
+    it(
+      'inlay hints are applied when vim.lsp._inlay_hint.refresh() is called',
+      function()
+        exec_lua([[
+        bufnr = vim.api.nvim_get_current_buf()
+        vim.api.nvim_win_set_buf(0, bufnr)
+        client_id = vim.lsp.start({ name = 'dummy', cmd = server.cmd })
+      ]])
+
+        insert(text)
+        exec_lua([[vim.lsp._inlay_hint.refresh({bufnr = bufnr})]])
+        screen:expect({
+          grid = [[
+  auto add(int a, int b)-> int { return a + b; }    |
+                                                    |
+  int main() {                                      |
+      int x = 1;                                    |
+      int y = 2;                                    |
+      return add(a: x,b: y);                        |
+  }                                                 |
+  ^}                                                 |
+                                                    |
+]]
+        })
+      end)
+
+    it(
+      'inlay hints are cleared when vim.lsp._inlay_hint.clear() is called',
+      function()
+        exec_lua([[
+        bufnr = vim.api.nvim_get_current_buf()
+        vim.api.nvim_win_set_buf(0, bufnr)
+        client_id = vim.lsp.start({ name = 'dummy', cmd = server.cmd })
+      ]])
+
+        insert(text)
+        exec_lua([[vim.lsp._inlay_hint.refresh({bufnr = bufnr})]])
+        screen:expect({
+          grid = [[
+  auto add(int a, int b)-> int { return a + b; }    |
+                                                    |
+  int main() {                                      |
+      int x = 1;                                    |
+      int y = 2;                                    |
+      return add(a: x,b: y);                        |
+  }                                                 |
+  ^}                                                 |
+                                                    |
+]]
+        })
+        exec_lua([[vim.lsp._inlay_hint.clear()]])
+        screen:expect({
+          grid = [[
+  auto add(int a, int b) { return a + b; }          |
+                                                    |
+  int main() {                                      |
+      int x = 1;                                    |
+      int y = 2;                                    |
+      return add(x,y);                              |
+  }                                                 |
+  ^}                                                 |
+                                                    |
+]],
+          unchanged = true
+        })
+      end)
+  end)
+end)
diff --git a/test/functional/plugin/lsp_spec.lua b/test/functional/plugin/lsp_spec.lua
index 85138417ff..2f4a703c74 100644
--- a/test/functional/plugin/lsp_spec.lua
+++ b/test/functional/plugin/lsp_spec.lua
@@ -3230,9 +3230,10 @@ describe('LSP', function()
           eq(0, signal, "exit signal")
         end;
         on_handler = function(err, result, ctx)
-          -- Don't compare & assert params, they're not relevant for the testcase
+          -- Don't compare & assert params and version, they're not relevant for the testcase
           -- This allows us to be lazy and avoid declaring them
           ctx.params = nil
+          ctx.version = nil
 
           eq(table.remove(test.expected_handlers), {err, result, ctx}, "expected handler")
           if ctx.method == 'start' then
@@ -3314,6 +3315,7 @@ describe('LSP', function()
         end,
         on_handler = function(err, result, ctx)
           ctx.params = nil -- don't compare in assert
+          ctx.version = nil
           eq(table.remove(expected_handlers), { err, result, ctx })
           if ctx.method == 'start' then
             exec_lua([[