From ca5de9306c00d07cce1daef1f0038c937098bc66 Mon Sep 17 00:00:00 2001
From: Chinmay Dalal <dalal.chinmay.0101@gmail.com>
Date: Tue, 20 Jun 2023 11:36:54 +0530
Subject: [PATCH] feat(lsp): inlay hints #23984

Add automatic refresh and a public interface on top of #23736

 * add on_reload, on_detach handlers in `enable()` buf_attach, and
  LspDetach autocommand in case of manual detach
* unify `__buffers` and `hint_cache_by_buf`
* use callback bufnr in `on_lines` callback, bufstate: remove __index override
* move user-facing functions into vim.lsp.buf, unify enable/disable/toggle

Closes #18086
---
 runtime/doc/lsp.txt                           |   7 +
 runtime/doc/news.txt                          |   5 +-
 runtime/lua/vim/lsp.lua                       |   1 -
 runtime/lua/vim/lsp/_inlay_hint.lua           | 164 +++++++++++++-----
 runtime/lua/vim/lsp/buf.lua                   |  15 ++
 runtime/lua/vim/lsp/handlers.lua              |   3 -
 scripts/gen_vimdoc.py                         |   1 -
 .../functional/plugin/lsp/inlay_hint_spec.lua |  10 +-
 8 files changed, 146 insertions(+), 60 deletions(-)

diff --git a/runtime/doc/lsp.txt b/runtime/doc/lsp.txt
index ca4570e392..73691c0656 100644
--- a/runtime/doc/lsp.txt
+++ b/runtime/doc/lsp.txt
@@ -1296,6 +1296,13 @@ incoming_calls()                                *vim.lsp.buf.incoming_calls()*
     window. If the symbol can resolve to multiple items, the user can pick one
     in the |inputlist()|.
 
+inlay_hint({bufnr}, {enable})                       *vim.lsp.buf.inlay_hint()*
+    Enable/disable/toggle inlay hints for a buffer
+
+    Parameters: ~
+      • {bufnr}   (integer) Buffer handle, or 0 for current
+      • {enable}  (boolean|nil) true/false to enable/disable, nil to toggle
+
 list_workspace_folders()                *vim.lsp.buf.list_workspace_folders()*
     List workspace folders.
 
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index cd3b1bd45e..2a25edc4eb 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -93,8 +93,9 @@ The following new APIs and 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`
+• Implemented LSP inlay hints: |vim.lsp.buf.inlay_hint()|
+  https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_inlayHint
+
 ==============================================================================
 CHANGED FEATURES                                                 *news-changed*
 
diff --git a/runtime/lua/vim/lsp.lua b/runtime/lua/vim/lsp.lua
index 761a8406f2..917aeb6604 100644
--- a/runtime/lua/vim/lsp.lua
+++ b/runtime/lua/vim/lsp.lua
@@ -17,7 +17,6 @@ local if_nil = vim.F.if_nil
 
 local lsp = {
   protocol = protocol,
-  _inlay_hint = require('vim.lsp._inlay_hint'),
 
   handlers = default_handlers,
 
diff --git a/runtime/lua/vim/lsp/_inlay_hint.lua b/runtime/lua/vim/lsp/_inlay_hint.lua
index aa6ec9aca8..70d332a1ac 100644
--- a/runtime/lua/vim/lsp/_inlay_hint.lua
+++ b/runtime/lua/vim/lsp/_inlay_hint.lua
@@ -6,22 +6,30 @@ local M = {}
 ---@class lsp._inlay_hint.bufstate
 ---@field version integer
 ---@field client_hint table<integer, table<integer, lsp.InlayHint[]>> client_id -> (lnum -> hints)
+---@field enabled boolean Whether inlay hints are enabled for the buffer
+---@field timer uv.uv_timer_t? Debounce timer associated with the buffer
 
 ---@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 bufstates = {}
 
 local namespace = api.nvim_create_namespace('vim_lsp_inlayhint')
+local augroup = api.nvim_create_augroup('vim_lsp_inlayhint', {})
 
-M.__explicit_buffers = {}
+--- Reset the request debounce timer of a buffer
+---@private
+local function reset_timer(reset_bufnr)
+  local timer = bufstates[reset_bufnr].timer
+  if timer then
+    bufstates[reset_bufnr].timer = nil
+    if not timer:is_closing() then
+      timer:stop()
+      timer:close()
+    end
+  end
+end
 
 --- |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
@@ -36,21 +44,20 @@ function M.on_inlayhint(err, result, ctx, _)
   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
+  local bufstate = bufstates[bufnr]
+  if not (bufstate.client_hint and bufstate.version) then
+    bufstate.client_hint = vim.defaulttable()
+    bufstate.version = ctx.version
     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
+      on_detach = function(_, cb_bufnr)
+        api.nvim_buf_clear_namespace(cb_bufnr, namespace, 0, -1)
+        bufstates[cb_bufnr].version = nil
+        bufstates[cb_bufnr].client_hint = nil
       end,
-      on_reload = function(_, b)
-        api.nvim_buf_clear_namespace(b, namespace, 0, -1)
-        hint_cache_by_buf[b] = nil
+      on_reload = function(_, cb_bufnr)
+        api.nvim_buf_clear_namespace(cb_bufnr, namespace, 0, -1)
+        bufstates[cb_bufnr].version = nil
+        bufstates[cb_bufnr].client_hint = nil
       end,
     })
   end
@@ -95,28 +102,24 @@ end
 
 ---@private
 local function resolve_bufnr(bufnr)
-  return bufnr == 0 and api.nvim_get_current_buf() or bufnr
+  return bufnr > 0 and bufnr or api.nvim_get_current_buf()
 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 bufnr = resolve_bufnr(opts.bufnr or 0)
+  local bufstate = bufstates[bufnr]
+  if not (bufstate and bufstate.enabled) then
+    return
+  end
   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
@@ -148,30 +151,95 @@ function M.refresh(opts)
 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
+---@param bufnr (integer) Buffer handle, or 0 for current
 ---@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
+local function clear(bufnr)
+  bufnr = resolve_bufnr(bufnr)
+  reset_timer(bufnr)
+  local bufstate = bufstates[bufnr]
+  local client_lens = (bufstate or {}).client_hint or {}
+  local client_ids = vim.tbl_keys(client_lens)
+  for _, iter_client_id in ipairs(client_ids) do
+    if bufstate then
+      bufstate.client_hint[iter_client_id] = {}
     end
-    api.nvim_buf_clear_namespace(iter_bufnr, namespace, 0, -1)
   end
-  vim.cmd('redraw!')
+  api.nvim_buf_clear_namespace(bufnr, namespace, 0, -1)
+  api.nvim__buf_redraw_range(bufnr, 0, -1)
+end
+
+---@private
+local function make_request(request_bufnr)
+  reset_timer(request_bufnr)
+  M.refresh({ bufnr = request_bufnr })
+end
+
+--- Enable inlay hints for a buffer
+---@param bufnr (integer) Buffer handle, or 0 for current
+---@private
+function M.enable(bufnr)
+  bufnr = resolve_bufnr(bufnr)
+  local bufstate = bufstates[bufnr]
+  if not (bufstate and bufstate.enabled) then
+    bufstates[bufnr] = { enabled = true, timer = nil }
+    M.refresh({ bufnr = bufnr })
+    api.nvim_buf_attach(bufnr, true, {
+      on_lines = function(_, cb_bufnr)
+        if not bufstates[cb_bufnr].enabled then
+          return true
+        end
+        reset_timer(cb_bufnr)
+        bufstates[cb_bufnr].timer = vim.defer_fn(function()
+          make_request(cb_bufnr)
+        end, 200)
+      end,
+      on_reload = function(_, cb_bufnr)
+        clear(cb_bufnr)
+        bufstates[cb_bufnr] = nil
+        M.refresh({ bufnr = cb_bufnr })
+      end,
+      on_detach = function(_, cb_bufnr)
+        clear(cb_bufnr)
+        bufstates[cb_bufnr] = nil
+      end,
+    })
+    api.nvim_create_autocmd('LspDetach', {
+      buffer = bufnr,
+      callback = function(opts)
+        clear(opts.buf)
+      end,
+      once = true,
+      group = augroup,
+    })
+  end
+end
+
+--- Disable inlay hints for a buffer
+---@param bufnr (integer) Buffer handle, or 0 for current
+---@private
+function M.disable(bufnr)
+  bufnr = resolve_bufnr(bufnr)
+  clear(bufnr)
+  bufstates[bufnr].enabled = nil
+  bufstates[bufnr].timer = nil
+end
+
+--- Toggle inlay hints for a buffer
+---@param bufnr (integer) Buffer handle, or 0 for current
+---@private
+function M.toggle(bufnr)
+  bufnr = resolve_bufnr(bufnr)
+  local bufstate = bufstates[bufnr]
+  if bufstate and bufstate.enabled then
+    M.disable(bufnr)
+  else
+    M.enable(bufnr)
+  end
 end
 
 api.nvim_set_decoration_provider(namespace, {
   on_win = function(_, _, bufnr, topline, botline)
-    local bufstate = hint_cache_by_buf[bufnr]
+    local bufstate = bufstates[bufnr]
     if not bufstate then
       return
     end
diff --git a/runtime/lua/vim/lsp/buf.lua b/runtime/lua/vim/lsp/buf.lua
index e0034cf86e..c3deffc1f9 100644
--- a/runtime/lua/vim/lsp/buf.lua
+++ b/runtime/lua/vim/lsp/buf.lua
@@ -810,4 +810,19 @@ function M.execute_command(command_params)
   request('workspace/executeCommand', command_params)
 end
 
+--- Enable/disable/toggle inlay hints for a buffer
+---@param bufnr (integer) Buffer handle, or 0 for current
+---@param enable (boolean|nil) true/false to enable/disable, nil to toggle
+function M.inlay_hint(bufnr, enable)
+  vim.validate({ enable = { enable, { 'boolean', 'nil' } }, bufnr = { bufnr, 'number' } })
+  local inlay_hint = require('vim.lsp._inlay_hint')
+  if enable then
+    inlay_hint.enable(bufnr)
+  elseif enable == false then
+    inlay_hint.disable(bufnr)
+  else
+    inlay_hint.toggle(bufnr)
+  end
+end
+
 return M
diff --git a/runtime/lua/vim/lsp/handlers.lua b/runtime/lua/vim/lsp/handlers.lua
index 44a9a58aca..284e3ef2d0 100644
--- a/runtime/lua/vim/lsp/handlers.lua
+++ b/runtime/lua/vim/lsp/handlers.lua
@@ -619,9 +619,6 @@ 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
diff --git a/scripts/gen_vimdoc.py b/scripts/gen_vimdoc.py
index 2d4c8b2234..b40f8526ea 100755
--- a/scripts/gen_vimdoc.py
+++ b/scripts/gen_vimdoc.py
@@ -223,7 +223,6 @@ 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
index b134095c4f..103fd8d968 100644
--- a/test/functional/plugin/lsp/inlay_hint_spec.lua
+++ b/test/functional/plugin/lsp/inlay_hint_spec.lua
@@ -63,7 +63,7 @@ describe('inlay hints', function()
     end)
 
     it(
-      'inlay hints are applied when vim.lsp._inlay_hint.refresh() is called',
+      'inlay hints are applied when vim.lsp.buf.inlay_hint(true) is called',
       function()
         exec_lua([[
         bufnr = vim.api.nvim_get_current_buf()
@@ -72,7 +72,7 @@ describe('inlay hints', function()
       ]])
 
         insert(text)
-        exec_lua([[vim.lsp._inlay_hint.refresh({bufnr = bufnr})]])
+        exec_lua([[vim.lsp.buf.inlay_hint(bufnr, true)]])
         screen:expect({
           grid = [[
   auto add(int a, int b)-> int { return a + b; }    |
@@ -89,7 +89,7 @@ describe('inlay hints', function()
       end)
 
     it(
-      'inlay hints are cleared when vim.lsp._inlay_hint.clear() is called',
+      'inlay hints are cleared when vim.lsp.buf.inlay_hint(false) is called',
       function()
         exec_lua([[
         bufnr = vim.api.nvim_get_current_buf()
@@ -98,7 +98,7 @@ describe('inlay hints', function()
       ]])
 
         insert(text)
-        exec_lua([[vim.lsp._inlay_hint.refresh({bufnr = bufnr})]])
+        exec_lua([[vim.lsp.buf.inlay_hint(bufnr, true)]])
         screen:expect({
           grid = [[
   auto add(int a, int b)-> int { return a + b; }    |
@@ -112,7 +112,7 @@ describe('inlay hints', function()
                                                     |
 ]]
         })
-        exec_lua([[vim.lsp._inlay_hint.clear()]])
+        exec_lua([[vim.lsp.buf.inlay_hint(bufnr, false)]])
         screen:expect({
           grid = [[
   auto add(int a, int b) { return a + b; }          |