feat(lsp): semantic token range improvements #37451

* cache all tokens from various range requests for a given document
  version
  - all new token highlights are merged with previous highlights to
    maintain order and the "marked" property
  - this allows the tokens to stop flickering once they've loaded once
    per document version
* abandon the processing coroutine if the request_id has changed instead
  of relying only on the document version
  - this will improve efficiency if a new range request is made while a
    previous one was processing its result
* apply new highlights from processing coroutine directly to the current
  result when the version hasn't changed
  - this allows new highlights to be immediately drawable once they've
    processed instead of waiting for the whole response to be processed
    at once
* rpc layer was changed to provide the request ID back in success
  callbacks, which is then provided as a request_id field on the handler
  context to lsp handlers

runtime/doc/lsp.txt

@@ -2436,7 +2436,7 @@ Lua module: vim.lsp.rpc                                              *lsp-rpc*
     Client RPC object
 
     Fields: ~
-      • {request}     (`fun(method: string, params: table?, callback: fun(err?: lsp.ResponseError, result: any), notify_reply_callback?: fun(message_id: integer)):boolean,integer?`)
+      • {request}     (`fun(method: string, params: table?, callback: fun(err?: lsp.ResponseError, result: any, request_id: integer), notify_reply_callback?: fun(message_id: integer)):boolean,integer?`)
                       See |vim.lsp.rpc.request()|
       • {notify}      (`fun(method: string, params: any): boolean`) See
                       |vim.lsp.rpc.notify()|

runtime/lua/vim/lsp/_meta.lua

@@ -7,6 +7,7 @@ error('Cannot require a meta file')
 ---@class lsp.HandlerContext
 ---@field method vim.lsp.protocol.Method
 ---@field client_id integer
+---@field request_id? integer
 ---@field bufnr? integer
 ---@field params? any
 ---@field version? integer

runtime/lua/vim/lsp/client.lua

@@ -731,10 +731,11 @@ function Client:request(method, params, handler, bufnr)
   local request_registered = false
 
   -- NOTE: rpc.request might call an in-process (Lua) server, thus may be synchronous.
-  local success, request_id = self.rpc.request(method, params, function(err, result)
+  local success, request_id = self.rpc.request(method, params, function(err, result, request_id)
     handler(err, result, {
       method = method,
       client_id = self.id,
+      request_id = request_id,
       bufnr = bufnr,
       params = params,
       version = version,
@@ -896,7 +897,7 @@ function Client:stop(force)
   end
 
   -- Sending a signal after a process has exited is acceptable.
-  rpc.request('shutdown', nil, function(err, _)
+  rpc.request('shutdown', nil, function(err, _, _)
     if err == nil then
       rpc.notify('exit')
     else

runtime/lua/vim/lsp/rpc.lua

@@ -298,7 +298,7 @@ end
 ---
 ---@param method vim.lsp.protocol.Method The invoked LSP method
 ---@param params table? Parameters for the invoked LSP method
----@param callback fun(err?: lsp.ResponseError, result: any) Callback to invoke
+---@param callback fun(err?: lsp.ResponseError, result: any, message_id: integer) Callback to invoke
 ---@param notify_reply_callback? fun(message_id: integer) Callback to invoke as soon as a request is no longer pending
 ---@return boolean success `true` if request could be sent, `false` if not
 ---@return integer? message_id if request could be sent, `nil` if not
@@ -467,7 +467,8 @@ function Client:handle_body(body)
         M.client_errors.SERVER_RESULT_CALLBACK_ERROR,
         callback,
         decoded.error,
-        decoded.result ~= vim.NIL and decoded.result or nil
+        decoded.result ~= vim.NIL and decoded.result or nil,
+        result_id
       )
     else
       self:on_error(M.client_errors.NO_RESULT_CALLBACK_FOUND, decoded)
@@ -505,7 +506,7 @@ end
 --- @class vim.lsp.rpc.PublicClient
 ---
 --- See [vim.lsp.rpc.request()]
---- @field request fun(method: vim.lsp.protocol.Method.ClientToServer.Request, params: table?, callback: fun(err?: lsp.ResponseError, result: any), notify_reply_callback?: fun(message_id: integer)):boolean,integer?
+--- @field request fun(method: vim.lsp.protocol.Method.ClientToServer.Request, params: table?, callback: fun(err?: lsp.ResponseError, result: any, request_id: integer), notify_reply_callback?: fun(message_id: integer)):boolean,integer?
 ---
 --- See [vim.lsp.rpc.notify()]
 --- @field notify fun(method: vim.lsp.protocol.Method.ClientToServer.Notification, params: any): boolean

runtime/lua/vim/lsp/semantic_tokens.lua

@@ -77,8 +77,9 @@ end
 ---@param bufnr integer
 ---@param client vim.lsp.Client
 ---@param request STActiveRequest
+---@param ranges STTokenRange[]
 ---@return STTokenRange[]
-local function tokens_to_ranges(data, bufnr, client, request)
+local function tokens_to_ranges(data, bufnr, client, request, ranges)
   local legend = client.server_capabilities.semanticTokensProvider.legend
   local token_types = legend.tokenTypes
   local token_modifiers = legend.tokenModifiers
@@ -86,7 +87,9 @@ local function tokens_to_ranges(data, bufnr, client, request)
   local lines = api.nvim_buf_get_lines(bufnr, 0, -1, false)
   -- For all encodings, \r\n takes up two code points, and \n (or \r) takes up one.
   local eol_offset = vim.bo.fileformat[bufnr] == 'dos' and 2 or 1
-  local ranges = {} ---@type STTokenRange[]
+  local version = request.version
+  local request_id = request.request_id
+  local last_insert_idx = 1
 
   local start = uv.hrtime()
   local ms_to_ns = 1e6
@@ -102,14 +105,18 @@ local function tokens_to_ranges(data, bufnr, client, request)
 
       if elapsed_ns > yield_interval_ns then
         vim.schedule(function()
-          coroutine.resume(co, util.buf_versions[bufnr])
+          -- Ensure the request hasn't become stale since the last time the coroutine ran.
+          -- If it's stale, we don't resume the coroutine so it'll be garbage collected.
+          if
+            version == util.buf_versions[bufnr]
+            and request_id == request.request_id
+            and api.nvim_buf_is_valid(bufnr)
+          then
+            coroutine.resume(co)
+          end
         end)
-        if request.version ~= coroutine.yield() then
-          -- request became stale since the last time the coroutine ran.
-          -- abandon it by yielding without a way to resume
-          coroutine.yield()
-        end
 
+        coroutine.yield()
         start = uv.hrtime()
       end
     end
@@ -141,7 +148,8 @@ local function tokens_to_ranges(data, bufnr, client, request)
 
       local end_col = vim.str_byteindex(buf_line, encoding, end_char, false)
 
-      ranges[#ranges + 1] = {
+      ---@type STTokenRange
+      local range = {
         line = line,
         end_line = end_line,
         start_col = start_col,
@@ -150,6 +158,47 @@ local function tokens_to_ranges(data, bufnr, client, request)
         modifiers = modifiers,
         marked = false,
       }
+
+      if last_insert_idx <= #ranges then
+        local needs_insert = true
+        local idx = vim.list.bisect(ranges, { line = range.line }, {
+          lo = last_insert_idx,
+          key = function(highlight)
+            return highlight.line
+          end,
+        })
+        while idx <= #ranges do
+          local token = ranges[idx]
+
+          if
+            token.line > range.line
+            or (token.line == range.line and token.start_col > range.start_col)
+          then
+            break
+          end
+
+          if
+            range.line == token.line
+            and range.start_col == token.start_col
+            and range.end_line == token.end_line
+            and range.end_col == token.end_col
+            and range.type == token.type
+          then
+            needs_insert = false
+            break
+          end
+
+          idx = idx + 1
+        end
+
+        last_insert_idx = idx
+        if needs_insert then
+          table.insert(ranges, last_insert_idx, range)
+        end
+      else
+        last_insert_idx = #ranges + 1
+        ranges[last_insert_idx] = range
+      end
     end
   end
 
@@ -207,7 +256,8 @@ function STHighlighter:cancel_active_request(client_id)
   if state.active_request.request_id then
     local client = assert(vim.lsp.get_client_by_id(client_id))
     client:cancel_request(state.active_request.request_id)
-    state.active_request = {}
+    state.active_request.request_id = nil
+    state.active_request.version = nil
   end
 end
 
@@ -274,8 +324,8 @@ function STHighlighter:send_request()
         -- cancel stale in-flight request
         if active_request.request_id then
           client:cancel_request(active_request.request_id)
-          active_request = {}
+          active_request.request_id = nil
-          state.active_request = active_request
+          active_request.version = nil
         end
 
         ---@type lsp.SemanticTokensParams|lsp.SemanticTokensRangeParams|lsp.SemanticTokensDeltaParams
@@ -301,11 +351,18 @@ function STHighlighter:send_request()
           end
 
           if err or not response then
-            highlighter.client_state[client.id].active_request = {}
+            active_request.request_id = nil
+            active_request.version = nil
             return
           end
 
-          coroutine.wrap(STHighlighter.process_response)(highlighter, response, client, version)
+          coroutine.wrap(STHighlighter.process_response)(
+            highlighter,
+            response,
+            client,
+            ctx.request_id,
+            version
+          )
         end, self.bufnr)
 
         if success then
@@ -359,16 +416,17 @@ end
 ---@async
 ---@param response lsp.SemanticTokens|lsp.SemanticTokensDelta
 ---@param client vim.lsp.Client
+---@param request_id integer
 ---@param version integer
 ---@private
-function STHighlighter:process_response(response, client, version)
+function STHighlighter:process_response(response, client, request_id, version)
   local state = self.client_state[client.id]
   if not state then
     return
   end
 
   -- ignore stale responses
-  if state.active_request.version and version ~= state.active_request.version then
+  if state.active_request.request_id and request_id ~= state.active_request.request_id then
     return
   end
 
@@ -400,25 +458,32 @@ function STHighlighter:process_response(response, client, version)
     tokens = response.data
   end
 
+  local current_result = state.current_result
+  local version_changed = version ~= current_result.version
+  local highlights = {} --- @type STTokenRange[]
+  if current_result.highlights and not version_changed then
+    highlights = assert(current_result.highlights)
+  end
+
   -- convert token list to highlight ranges
   -- this could yield and run over multiple event loop iterations
-  local highlights = tokens_to_ranges(tokens, self.bufnr, client, state.active_request)
+  highlights = tokens_to_ranges(tokens, self.bufnr, client, state.active_request, highlights)
 
   -- reset active request
-  state.active_request = {}
+  state.active_request.request_id = nil
+  state.active_request.version = nil
 
   -- update the state with the new results
-  local current_result = state.current_result
   current_result.version = version
   current_result.result_id = response.resultId
   current_result.tokens = tokens
   current_result.highlights = highlights
-  current_result.namespace_cleared = false
+  if version_changed then
-
+    current_result.namespace_cleared = false
-  -- redraw all windows displaying buffer (if still valid)
-  if api.nvim_buf_is_valid(self.bufnr) then
-    api.nvim__redraw({ buf = self.bufnr, valid = true })
   end
+
+  -- redraw all windows displaying buffer
+  api.nvim__redraw({ buf = self.bufnr, valid = true })
 end
 
 --- @param bufnr integer

test/functional/plugin/lsp/semantic_tokens_spec.lua

@@ -322,6 +322,182 @@ describe('semantic token highlighting', function()
       eq(true, called_range)
     end)
 
+    it('range requests preserve highlights outside updated range', function()
+      screen:try_resize(40, 6)
+      insert(text)
+      feed('gg')
+
+      local small_range_response = [[{
+        "data": [ 2, 4, 4, 3, 8193, 2, 8, 1, 1, 1025 ]
+      }]]
+
+      local client_id, bufnr = exec_lua(function(l, resp)
+        _G.response = resp
+        _G.server2 = _G._create_server({
+          capabilities = {
+            semanticTokensProvider = {
+              range = true,
+              legend = vim.fn.json_decode(l),
+            },
+          },
+          handlers = {
+            ['textDocument/semanticTokens/range'] = function(_, _, callback)
+              callback(nil, vim.fn.json_decode(_G.response))
+            end,
+          },
+        })
+        local bufnr = vim.api.nvim_get_current_buf()
+        local client_id = assert(vim.lsp.start({ name = 'dummy', cmd = _G.server2.cmd }))
+        vim.schedule(function()
+          vim.lsp.semantic_tokens._start(bufnr, client_id, 10)
+        end)
+        return client_id, bufnr
+      end, legend, small_range_response)
+
+      screen:expect {
+        grid = [[
+        ^#include <iostream>                     |
+                                                |
+        int {8:main}()                              |
+        {                                       |
+            int {7:x};                              |
+                                                |
+      ]],
+      }
+
+      eq(
+        2,
+        exec_lua(function()
+          return #vim.lsp.semantic_tokens.__STHighlighter.active[bufnr].client_state[client_id].current_result.highlights
+        end)
+      )
+
+      small_range_response = [[{
+        "data": [ 7, 0, 5, 20, 0, 1, 0, 22, 20, 0, 1, 0, 6, 20, 0 ]
+      }]]
+
+      exec_lua(function(resp)
+        _G.response = resp
+      end, small_range_response)
+
+      feed('G')
+
+      screen:expect {
+        grid = [[
+        {6:#else}                                   |
+        {6:    printf("%d\n", x);}                  |
+        {6:#endif}                                  |
+        }                                       |
+        ^}                                       |
+                                                |
+      ]],
+      }
+
+      eq(
+        5,
+        exec_lua(function()
+          return #vim.lsp.semantic_tokens.__STHighlighter.active[bufnr].client_state[client_id].current_result.highlights
+        end)
+      )
+
+      small_range_response = [[{
+        "data": [ 2, 4, 4, 3, 8193, 2, 8, 1, 1, 1025, 1, 7, 11, 19, 8192 ]
+      }]]
+
+      exec_lua(function(resp)
+        _G.response = resp
+      end, small_range_response)
+      feed('ggLj0')
+
+      screen:expect {
+        grid = [[
+                                                |
+        int {8:main}()                              |
+        {                                       |
+            int {7:x};                              |
+        ^#ifdef {5:__cplusplus}                      |
+                                                |
+      ]],
+      }
+
+      eq(
+        6,
+        exec_lua(function()
+          return #vim.lsp.semantic_tokens.__STHighlighter.active[bufnr].client_state[client_id].current_result.highlights
+        end)
+      )
+
+      eq(
+        {
+          {
+            line = 2,
+            end_line = 2,
+            start_col = 4,
+            end_col = 8,
+            marked = true,
+            modifiers = {
+              declaration = true,
+              globalScope = true,
+            },
+            type = 'function',
+          },
+          {
+            line = 4,
+            end_line = 4,
+            start_col = 8,
+            end_col = 9,
+            marked = true,
+            modifiers = {
+              declaration = true,
+              functionScope = true,
+            },
+            type = 'variable',
+          },
+          {
+            line = 5,
+            end_line = 5,
+            start_col = 7,
+            end_col = 18,
+            marked = true,
+            modifiers = {
+              globalScope = true,
+            },
+            type = 'macro',
+          },
+          {
+            line = 7,
+            end_line = 7,
+            start_col = 0,
+            end_col = 5,
+            marked = true,
+            modifiers = {},
+            type = 'comment',
+          },
+          {
+            line = 8,
+            end_line = 8,
+            start_col = 0,
+            end_col = 22,
+            marked = true,
+            modifiers = {},
+            type = 'comment',
+          },
+          {
+            line = 9,
+            end_line = 9,
+            start_col = 0,
+            end_col = 6,
+            marked = true,
+            modifiers = {},
+            type = 'comment',
+          },
+        },
+        exec_lua(function()
+          return vim.lsp.semantic_tokens.__STHighlighter.active[bufnr].client_state[client_id].current_result.highlights
+        end)
+      )
+    end)
+
     it('use LspTokenUpdate and highlight_token', function()
       insert(text)
       exec_lua(function()

test/functional/plugin/lsp_spec.lua

@@ -1074,7 +1074,7 @@ describe('LSP', function()
         {
           { code = -32802 },
           NIL,
-          { method = 'error_code_test', bufnr = 1, client_id = 1, version = 0 },
+          { method = 'error_code_test', bufnr = 1, client_id = 1, request_id = 2, version = 0 },
         },
       }
       local client --- @type vim.lsp.Client
@@ -1107,7 +1107,7 @@ describe('LSP', function()
         {
           { code = -32801 },
           NIL,
-          { method = 'error_code_test', bufnr = 1, client_id = 1, version = 0 },
+          { method = 'error_code_test', bufnr = 1, client_id = 1, request_id = 2, version = 0 },
         },
       }
       local client --- @type vim.lsp.Client
@@ -1137,7 +1137,11 @@ describe('LSP', function()
     it('should track pending requests to the language server', function()
       local expected_handlers = {
         { NIL, {}, { method = 'finish', client_id = 1 } },
-        { NIL, {}, { method = 'slow_request', bufnr = 1, client_id = 1, version = 0 } },
+        {
+          NIL,
+          {},
+          { method = 'slow_request', bufnr = 1, client_id = 1, request_id = 2, version = 0 },
+        },
       }
       local client --- @type vim.lsp.Client
       test_rpc_server {
@@ -1212,7 +1216,11 @@ describe('LSP', function()
     it('should clear pending and cancel requests on reply', function()
       local expected_handlers = {
         { NIL, {}, { method = 'finish', client_id = 1 } },
-        { NIL, {}, { method = 'slow_request', bufnr = 1, client_id = 1, version = 0 } },
+        {
+          NIL,
+          {},
+          { method = 'slow_request', bufnr = 1, client_id = 1, request_id = 2, version = 0 },
+        },
       }
       local client --- @type vim.lsp.Client
       test_rpc_server {
@@ -1316,7 +1324,11 @@ describe('LSP', function()
     it('should trigger LspRequest autocmd when requests table changes', function()
       local expected_handlers = {
         { NIL, {}, { method = 'finish', client_id = 1 } },
-        { NIL, {}, { method = 'slow_request', bufnr = 1, client_id = 1, version = 0 } },
+        {
+          NIL,
+          {},
+          { method = 'slow_request', bufnr = 1, client_id = 1, request_id = 2, version = 0 },
+        },
       }
       local client --- @type vim.lsp.Client
       test_rpc_server {
@@ -1609,6 +1621,7 @@ describe('LSP', function()
             },
             bufnr = 2,
             client_id = 1,
+            request_id = 2,
             version = 0,
           },
         },
@@ -4515,7 +4528,7 @@ describe('LSP', function()
         name = 'prepare_rename_placeholder',
         expected_handlers = {
           { NIL, {}, { method = 'shutdown', client_id = 1 } },
-          { {}, NIL, { method = 'textDocument/rename', client_id = 1, bufnr = 1 } },
+          { {}, NIL, { method = 'textDocument/rename', client_id = 1, request_id = 3, bufnr = 1 } },
           { NIL, {}, { method = 'start', client_id = 1 } },
         },
         expected_text = 'placeholder', -- see fake lsp response
@@ -4525,7 +4538,7 @@ describe('LSP', function()
         name = 'prepare_rename_range',
         expected_handlers = {
           { NIL, {}, { method = 'shutdown', client_id = 1 } },
-          { {}, NIL, { method = 'textDocument/rename', client_id = 1, bufnr = 1 } },
+          { {}, NIL, { method = 'textDocument/rename', client_id = 1, request_id = 3, bufnr = 1 } },
           { NIL, {}, { method = 'start', client_id = 1 } },
         },
         expected_text = 'line', -- see test case and fake lsp response
@@ -4653,7 +4666,7 @@ describe('LSP', function()
         {
           NIL,
           { command = 'dummy1', title = 'Command 1' },
-          { bufnr = 1, method = 'workspace/executeCommand', client_id = 1 },
+          { bufnr = 1, method = 'workspace/executeCommand', request_id = 3, client_id = 1 },
         },
         { NIL, {}, { method = 'start', client_id = 1 } },
       }