feat(lsp): completion side effects

Author: MariaSolOs

Diff for: runtime/lua/vim/lsp/_completion.lua

@@ -4,6 +4,9 @@ local lsp = vim.lsp
 local protocol = lsp.protocol
 local ms = protocol.Methods
 
+local rtt_ms = 50
+local ns_to_ms = 0.000001
+
 --- @alias vim.lsp.CompletionResult lsp.CompletionList | lsp.CompletionItem[]
 
 -- TODO(mariasolos): Remove this declaration once we figure out a better way to handle
@@ -14,11 +17,96 @@ local ms = protocol.Methods
 --- @field insertTextMode lsp.InsertTextMode?
 --- @field data any
 
+--- @class vim.lsp.completion.BufHandle
+--- @field clients table<integer, vim.lsp.Client>
+--- @field triggers table<string, vim.lsp.Client[]>
+
+--- @type table<integer, vim.lsp.completion.BufHandle>
+local buf_handles = {}
+
+--- @class vim.lsp.completion.Context
+--- @field last_request_time integer?
+--- @field pending_requests function[]
+--- @field cursor { [1]: integer, [2]: integer }?
+--- @field isIncomplete boolean
+--- @field suppress_completeDone boolean
+local Context = {
+  pending_requests = {},
+  isIncomplete = false,
+  suppress_completeDone = false,
+}
+
+function Context:cancel_pending()
+  for _, cancel in ipairs(self.pending_requests) do
+    cancel()
+  end
+
+  self.pending_requests = {}
+end
+
+function Context:reset()
+  -- Note that the cursor isn't reset here, it needs to survive a `CompleteDone` event.
+  self.last_request_time = nil
+  self:cancel_pending()
+end
+
+--- @type uv.uv_timer_t?
+local completion_timer = nil
+
+--- @param timer uv.uv_timer_t?
+--- @return uv.uv_timer_t
+local function new_timer(timer)
+  timer = assert(vim.uv.new_timer())
+  return timer
+end
+
+--- @param timer uv.uv_timer_t?
+local function reset_timer(timer)
+  if timer then
+    timer:stop()
+    timer:close()
+  end
+
+  timer = nil
+end
+
+--- @param window integer
+--- @param warmup integer
+--- @return fun(sample: number): number
+local function exp_avg(window, warmup)
+  local count = 0
+  local sum = 0
+  local value = 0
+
+  return function(sample)
+    if count < warmup then
+      count = count + 1
+      sum = sum + sample
+      value = sum / count
+    else
+      local factor = 2.0 / (window + 1)
+      value = value * (1 - factor) + sample * factor
+    end
+    return value
+  end
+end
+local compute_new_average = exp_avg(10, 10)
+
+--- @return number
+local function next_debounce()
+  if not Context.last_request_time then
+    return rtt_ms
+  end
+
+  local ms_since_request = (vim.uv.hrtime() - Context.last_request_time) * ns_to_ms
+  return math.max((ms_since_request - rtt_ms) * -1, 0)
+end
+
 ---@param input string unparsed snippet
 ---@return string parsed snippet
 local function parse_snippet(input)
   local ok, parsed = pcall(function()
-    return vim.lsp._snippet_grammar.parse(input)
+    return lsp._snippet_grammar.parse(input)
   end)
   return ok and tostring(parsed) or input
 end
@@ -159,7 +247,7 @@ local function adjust_start_col(lnum, line, items, encoding)
     end
   end
   if min_start_char then
-    return vim.lsp.util._str_byteindex_enc(line, min_start_char, encoding)
+    return lsp.util._str_byteindex_enc(line, min_start_char, encoding)
   else
     return nil
   end
@@ -246,7 +334,7 @@ function M.omnifunc(findstart, base)
     local params = util.make_position_params(win, client.offset_encoding)
     client.request(ms.textDocument_completion, params, function(err, result)
       if err then
-        vim.lsp.log.warn(err.message)
+        lsp.log.warn(err.message)
       end
       if result and vim.fn.mode() == 'i' then
         local matches
@@ -273,4 +361,257 @@ function M.omnifunc(findstart, base)
   return -2
 end
 
+--- Initializes the completion commands for the given client.
+---
+--- @param client vim.lsp.Client
+local function init_commands(client)
+  local trigger_completion_cmd = 'editor.action.triggerSuggest'
+
+  -- Check if the command is in the global registry or in the client's commands.
+  if not lsp.commands[trigger_completion_cmd] and not client.commands[trigger_completion_cmd] then
+    client.commands[trigger_completion_cmd] = function()
+      pcall(M.trigger_completion)
+    end
+  end
+end
+
+--- @param clients table<integer, vim.lsp.Client>
+--- @param bufnr integer
+--- @param win integer
+--- @param callback fun(responses: table<integer, { err: lsp.ResponseError, result: vim.lsp.CompletionResult }>)
+--- @return function # Cancellation function
+local function request(clients, bufnr, win, callback)
+  local responses = {} --- @type table<integer, { err: lsp.ResponseError, result: any }>
+  local request_ids = {} --- @type table<integer, integer>
+  local remaining_requests = vim.tbl_count(clients)
+
+  for client_id, client in pairs(clients) do
+    local params = lsp.util.make_position_params(win, client.offset_encoding)
+    local ok, request_id = client.request(ms.textDocument_completion, params, function(err, result)
+      responses[client_id] = { err = err, result = result }
+      remaining_requests = remaining_requests - 1
+      if remaining_requests == 0 then
+        callback(responses)
+      end
+    end, bufnr)
+
+    if ok then
+      request_ids[client_id] = request_id
+    end
+  end
+
+  return function()
+    for client_id, request_id in pairs(request_ids) do
+      local client = lsp.get_client_by_id(client_id)
+      if client then
+        client.cancel_request(request_id)
+      end
+    end
+  end
+end
+
+--- @param handle vim.lsp.completion.BufHandle
+local function insert_char_pre_cb(handle)
+  if tonumber(vim.fn.pumvisible()) == 1 then
+    if Context.isIncomplete then
+      reset_timer(completion_timer)
+
+      -- Calling vim.fn.complete while pumvisible will trigger `CompleteDone` for the active completion window,
+      -- so we suppress it to avoid resetting the completion context.
+      Context.suppress_completeDone = true
+
+      local debounce_ms = next_debounce()
+      if debounce_ms == 0 then
+        vim.schedule(M.trigger_completion)
+      else
+        completion_timer = new_timer(completion_timer)
+        completion_timer:start(debounce_ms, 0, vim.schedule_wrap(M.trigger_completion))
+      end
+    end
+
+    return
+  end
+
+  local char = api.nvim_get_vvar('char')
+  if not completion_timer and handle.triggers[char] then
+    completion_timer = assert(vim.uv.new_timer())
+    completion_timer:start(25, 0, function()
+      reset_timer(completion_timer)
+      vim.schedule(M.trigger_completion)
+    end)
+  end
+end
+
+local function text_changed_p_cb()
+  Context.cursor = api.nvim_win_get_cursor(0)
+end
+
+local function text_changed_i_cb()
+  if not Context.cursor or completion_timer then
+    return
+  end
+
+  local cursor = api.nvim_win_get_cursor(0)
+  if cursor[1] == Context.cursor[1] and cursor[2] <= Context.cursor[2] then
+    completion_timer = new_timer(completion_timer)
+    completion_timer:start(150, 0, vim.schedule_wrap(M.trigger_completion))
+  elseif cursor[1] ~= Context.cursor[1] then
+    Context.cursor = nil
+  end
+end
+
+local function insert_leave_cb()
+  reset_timer(completion_timer)
+  Context.cursor = nil
+  Context:reset()
+end
+
+local function complete_done_cb()
+  if Context.suppress_completeDone then
+    Context.suppress_completeDone = false
+    return
+  end
+end
+
+function M.trigger_completion()
+  reset_timer(completion_timer)
+  Context:cancel_pending()
+
+  local win = api.nvim_get_current_win()
+  local bufnr = api.nvim_get_current_buf()
+  local cursor_row, cursor_col = unpack(api.nvim_win_get_cursor(win)) --- @type integer, integer
+  local line = api.nvim_get_current_line()
+  local line_to_cursor = line:sub(1, cursor_col)
+  local clients = (buf_handles[bufnr] or {}).clients or {}
+  local word_boundary = vim.fn.match(line_to_cursor, '\\k*$')
+  local start_time = vim.uv.hrtime()
+  Context.last_request_time = start_time
+
+  local cancel_request = request(clients, bufnr, win, function(responses)
+    local end_time = vim.uv.hrtime()
+    rtt_ms = compute_new_average((end_time - start_time) * ns_to_ms)
+
+    Context.pending_requests = {}
+    Context.isIncomplete = false
+
+    local row_changed = api.nvim_win_get_cursor(win)[1] ~= cursor_row
+    local mode = api.nvim_get_mode().mode
+    if row_changed or not (mode == 'i' or mode == 'ic') then
+      return
+    end
+
+    local matches = {}
+    local server_start_boundary --- @type integer?
+    for client_id, response in pairs(responses) do
+      if response.err then
+        lsp.log.warn(response.err.message)
+      end
+
+      local result = response.result
+      if result then
+        Context.isIncomplete = Context.isIncomplete or result.isIncomplete
+        local client = lsp.get_client_by_id(client_id)
+        local encoding = client and client.offset_encoding or 'utf-16'
+        local client_matches
+        client_matches, server_start_boundary =
+          M._convert_results(line, cursor_row - 1, cursor_col, word_boundary, nil, result, encoding)
+        vim.list_extend(matches, client_matches)
+      end
+    end
+    local start_col = (server_start_boundary or word_boundary) + 1
+    vim.fn.complete(start_col, matches)
+  end)
+
+  table.insert(Context.pending_requests, cancel_request)
+end
+
+--- @param client_id integer Client ID
+--- @param bufnr integer Buffer handle, or 0 for the current buffer
+function M.attach(client_id, bufnr)
+  bufnr = (bufnr == 0 and api.nvim_get_current_buf()) or bufnr
+
+  if not buf_handles[bufnr] then
+    buf_handles[bufnr] = { clients = {}, triggers = {} }
+
+    -- Attach to buffer events.
+    api.nvim_buf_attach(bufnr, false, {
+      on_detach = function(_, buf)
+        buf_handles[buf] = nil
+      end,
+      on_reload = function(_, buf)
+        M.attach(client_id, buf)
+      end,
+    })
+
+    -- Set up autocommands.
+    local group =
+      api.nvim_create_augroup(string.format('vim/lsp/completion-%d', bufnr), { clear = true })
+    api.nvim_create_autocmd('InsertCharPre', {
+      group = group,
+      buffer = bufnr,
+      callback = function()
+        insert_char_pre_cb(buf_handles[bufnr])
+      end,
+    })
+    api.nvim_create_autocmd('TextChangedP', {
+      group = group,
+      buffer = bufnr,
+      callback = function()
+        text_changed_p_cb()
+      end,
+    })
+    api.nvim_create_autocmd('TextChangedI', {
+      group = group,
+      buffer = bufnr,
+      callback = function()
+        text_changed_i_cb()
+      end,
+    })
+    api.nvim_create_autocmd('InsertLeave', {
+      group = group,
+      buffer = bufnr,
+      callback = function()
+        insert_leave_cb()
+      end,
+    })
+    api.nvim_create_autocmd('CompleteDone', {
+      group = group,
+      buffer = bufnr,
+      callback = complete_done_cb,
+    })
+  end
+
+  if not buf_handles[bufnr].clients[client_id] then
+    local client = lsp.get_client_by_id(client_id)
+    assert(client, 'invalid client ID')
+
+    -- Add the new client to the buffer's clients.
+    init_commands(client)
+    buf_handles[bufnr].clients[client_id] = client
+
+    -- Add the new client to the clients that should be triggered by its trigger characters.
+    --- @type string[]
+    local triggers = vim.tbl_get(
+      client.server_capabilities,
+      'completionProvider',
+      'triggerCharacters'
+    ) or {}
+    for _, char in ipairs(triggers) do
+      local clients_for_trigger = buf_handles[bufnr].triggers[char]
+      if not clients_for_trigger then
+        clients_for_trigger = {}
+        buf_handles[bufnr].triggers[char] = clients_for_trigger
+      end
+      local client_exists = vim.iter(clients_for_trigger):any(function(c)
+        return c.id == client_id
+      end)
+      if not client_exists then
+        table.insert(clients_for_trigger, client)
+      end
+    end
+  end
+end
+
+-- TODO(mariasolos): Add M.detach if we decide to use the attach approach.
+
 return M