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
This commit is contained in:
jdrouhard
2026-01-27 07:56:52 -06:00
committed by GitHub
parent 8c63d84be1
commit 8ed68fda50
7 changed files with 295 additions and 38 deletions

View File

@@ -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()|

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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 = {}
state.active_request = active_request
active_request.request_id = nil
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
-- 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 })
if version_changed then
current_result.namespace_cleared = false
end
-- redraw all windows displaying buffer
api.nvim__redraw({ buf = self.bufnr, valid = true })
end
--- @param bufnr integer

View File

@@ -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()

View File

@@ -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 } },
}