From 20dcd208731bbacf15aeb0020fd9e12031faa9b3 Mon Sep 17 00:00:00 2001 From: Oscar Wallberg Date: Sun, 19 Apr 2026 00:56:34 +0200 Subject: [PATCH] refactor(codelens): split into modules and add cancellation --- lua/lsp/codelens.lua | 354 ------------------------------ lua/lsp/codelens/init.lua | 73 ++++++ lua/lsp/codelens/refresh_task.lua | 134 +++++++++++ lua/lsp/codelens/row.lua | 155 +++++++++++++ lua/lsp/codelens/session.lua | 109 +++++++++ 5 files changed, 471 insertions(+), 354 deletions(-) delete mode 100644 lua/lsp/codelens.lua create mode 100644 lua/lsp/codelens/init.lua create mode 100644 lua/lsp/codelens/refresh_task.lua create mode 100644 lua/lsp/codelens/row.lua create mode 100644 lua/lsp/codelens/session.lua diff --git a/lua/lsp/codelens.lua b/lua/lsp/codelens.lua deleted file mode 100644 index 09e5519..0000000 --- a/lua/lsp/codelens.lua +++ /dev/null @@ -1,354 +0,0 @@ -local util = require("util") - -local GROUP = vim.api.nvim_create_augroup("ow.lsp.codelens", { clear = true }) -local NS = vim.api.nvim_create_namespace("ow.lsp.codelens") -local REFRESH_DEBOUNCE_MS = 200 - -local M = {} - ----@alias ow.lsp.codelens.Position "eol" | "right_align" | "inline" | "above" - ----@type ow.lsp.codelens.Position -local position = "above" -local separator = " | " - ----@type table -local state_by_buf = {} - -local refresh_debounced = util.debounce(function(buf) - local state = state_by_buf[buf] - if state then - state:refresh() - end -end, REFRESH_DEBOUNCE_MS) - ----@class ow.lsp.codelens.Row ----@field row integer ----@field ready lsp.CodeLens[] ----@field pending integer -local Row = {} -Row.__index = Row - ----@param row integer ----@return ow.lsp.codelens.Row -function Row.new(row) - return setmetatable({ row = row, ready = {}, pending = 0 }, Row) -end - ----@param lens lsp.CodeLens -function Row:add(lens) - table.insert(self.ready, lens) -end - -function Row:await() - self.pending = self.pending + 1 -end - ----@param resolved? lsp.CodeLens -function Row:resolved(resolved) - self.pending = self.pending - 1 - if resolved then - table.insert(self.ready, resolved) - end -end - ----@param buf integer -function Row:render(buf) - if not vim.api.nvim_buf_is_valid(buf) then - return - end - if self.pending > 0 or #self.ready == 0 then - return - end - - table.sort(self.ready, function(a, b) - return a.range.start.character < b.range.start.character - end) - local parts = {} - for i, lens in ipairs(self.ready) do - table.insert(parts, { lens.command.title, "LspCodeLens" }) - if i < #self.ready then - table.insert(parts, { separator, "LspCodeLensSeparator" }) - end - end - - local col, opts - if position == "above" then - local line = vim.api.nvim_buf_get_lines( - buf, - self.row, - self.row + 1, - false - )[1] or "" - local indent = line:match("^%s*") or "" - local chunks = parts - if indent ~= "" then - chunks = { { indent, "LspCodeLensSeparator" } } - vim.list_extend(chunks, parts) - end - col = 0 - opts = { - virt_lines = { chunks }, - virt_lines_above = true, - hl_mode = "combine", - } - else - col = position == "inline" and self.ready[1].range.start.character or 0 - opts = { - virt_text = parts, - virt_text_pos = position, - hl_mode = "combine", - } - end - - -- One extmark per row. Extmarks auto-shift with edits, so multiple from - -- prior refreshes can end up on the same row; pick the first and drop - -- the rest, then reuse its id for an atomic update. - local marks = vim.api.nvim_buf_get_extmarks( - buf, - NS, - { self.row, 0 }, - { self.row, -1 }, - { details = true } - ) - local primary - for i, mark in ipairs(marks) do - if i == 1 then - primary = { id = mark[1], col = mark[3], details = mark[4] } - else - vim.api.nvim_buf_del_extmark(buf, NS, mark[1]) - end - end - - if primary and primary.col == col and primary.details then - local d = primary.details --[[@as vim.api.keyset.extmark_details]] - local same = opts.virt_lines - and vim.deep_equal(d.virt_lines, opts.virt_lines) - and d.virt_lines_above == opts.virt_lines_above - or not opts.virt_lines - and d.virt_text_pos == opts.virt_text_pos - and vim.deep_equal(d.virt_text, opts.virt_text) - if same then - return - end - end - - opts.id = primary and primary.id or nil - vim.api.nvim_buf_set_extmark(buf, NS, self.row, col, opts) -end - ----@class ow.lsp.codelens.State ----@field buf integer ----@field enabled boolean ----@field attached boolean ----@field generation integer ----@field rows table -local State = {} -State.__index = State - ----@param buf integer ----@return ow.lsp.codelens.State -function State.new(buf) - return setmetatable({ - buf = buf, - enabled = false, - attached = false, - generation = 0, - rows = {}, - }, State) -end - -function State:render() - if not vim.api.nvim_buf_is_valid(self.buf) then - return - end - if not self.enabled then - vim.api.nvim_buf_clear_namespace(self.buf, NS, 0, -1) - return - end - - for _, row in pairs(self.rows) do - row:render(self.buf) - end - - for _, mark in - ipairs(vim.api.nvim_buf_get_extmarks(self.buf, NS, 0, -1, {})) - do - if not self.rows[mark[2]] then - vim.api.nvim_buf_del_extmark(self.buf, NS, mark[1]) - end - end -end - -function State:refresh() - if not self.enabled then - return - end - - self.generation = self.generation + 1 - local gen = self.generation - local params = { - textDocument = vim.lsp.util.make_text_document_params(self.buf), - } - - local function collect(new_rows, client, lens) - local row = new_rows[lens.range.start.line] - if not row then - row = Row.new(lens.range.start.line) - new_rows[lens.range.start.line] = row - end - if lens.command then - row:add(lens) - return - end - if - not client - or not client:supports_method( - vim.lsp.protocol.Methods.codeLens_resolve - ) - then - return - end - row:await() - client:request( - vim.lsp.protocol.Methods.codeLens_resolve, - lens, - function(_, resolved) - if self.generation ~= gen then - return - end - row:resolved(type(resolved) == "table" and resolved or nil) - row:render(self.buf) - end, - self.buf - ) - end - - vim.lsp.buf_request_all( - self.buf, - vim.lsp.protocol.Methods.textDocument_codeLens, - params, - function(results) - if self.generation ~= gen then - return - end - local new_rows = {} - for client_id, response in pairs(results) do - if not response.err and type(response.result) == "table" then - local client = vim.lsp.get_client_by_id(client_id) - for _, lens in ipairs(response.result) do - collect(new_rows, client, lens) - end - end - end - self.rows = new_rows - self:render() - end - ) -end - -function State:attach() - if self.attached then - return - end - self.attached = true - vim.api.nvim_buf_attach(self.buf, false, { - on_lines = function() - if not self.enabled then - self.attached = false - return true - end - refresh_debounced:call(self.buf) - end, - on_reload = function() - if self.enabled then - refresh_debounced:call(self.buf) - end - end, - on_detach = function() - self.attached = false - end, - }) -end - ----@param value boolean -function State:enable(value) - self.enabled = value - if value then - self:attach() - self:refresh() - else - refresh_debounced:cancel(self.buf) - self:render() - end -end - -function State:toggle() - self:enable(not self.enabled) -end - ----@param buf integer ----@return ow.lsp.codelens.State -local function get_state(buf) - state_by_buf[buf] = state_by_buf[buf] or State.new(buf) - return state_by_buf[buf] -end - ----@param buf? integer ----@return boolean -function M.is_enabled(buf) - buf = buf or vim.api.nvim_get_current_buf() - local state = state_by_buf[buf] - return state ~= nil and state.enabled -end - ----@param value? boolean ----@param buf? integer -function M.enable(value, buf) - if value == nil then - value = true - end - buf = buf or vim.api.nvim_get_current_buf() - get_state(buf):enable(value) -end - ----@param buf? integer -function M.toggle(buf) - buf = buf or vim.api.nvim_get_current_buf() - get_state(buf):toggle() -end - ----@class ow.lsp.codelens.SetupOpts ----@field position? ow.lsp.codelens.Position ----@field separator? string - ----@param opts? ow.lsp.codelens.SetupOpts -function M.setup(opts) - opts = opts or {} - if opts.position ~= nil then - position = opts.position - end - if opts.separator ~= nil then - separator = opts.separator - end - - vim.api.nvim_create_autocmd({ "BufEnter", "LspAttach" }, { - group = GROUP, - callback = function(ev) - local state = state_by_buf[ev.buf] - if state and state.enabled then - state:refresh() - end - end, - }) - - vim.api.nvim_create_autocmd({ "BufDelete", "BufWipeout" }, { - group = GROUP, - callback = function(ev) - refresh_debounced:cancel(ev.buf) - state_by_buf[ev.buf] = nil - end, - }) -end - -return M diff --git a/lua/lsp/codelens/init.lua b/lua/lsp/codelens/init.lua new file mode 100644 index 0000000..f0856da --- /dev/null +++ b/lua/lsp/codelens/init.lua @@ -0,0 +1,73 @@ +---@class ow.lsp.CodeLensResponse : ow.lsp.Response +---@field result lsp.CodeLens[] + +local Row = require("lsp.codelens.row") +local Session = require("lsp.codelens.session") + +local GROUP = vim.api.nvim_create_augroup("ow.lsp.codelens", { clear = true }) + +---@type table +local session_by_buf = {} + +local M = {} + +---@param buf integer +---@return ow.lsp.codelens.Session +local function get_session(buf) + session_by_buf[buf] = session_by_buf[buf] or Session.new(buf) + return session_by_buf[buf] +end + +---@param buf? integer +---@return boolean +function M.is_enabled(buf) + buf = buf or vim.api.nvim_get_current_buf() + local session = session_by_buf[buf] + return session ~= nil and session.enabled +end + +---@param value? boolean +---@param buf? integer +function M.enable(value, buf) + if value == nil then + value = true + end + buf = buf or vim.api.nvim_get_current_buf() + get_session(buf):enable(value) +end + +---@param buf? integer +function M.toggle(buf) + buf = buf or vim.api.nvim_get_current_buf() + get_session(buf):toggle() +end + +---@class ow.lsp.codelens.SetupOpts : ow.lsp.codelens.row.ConfigureOpts + +---@param opts? ow.lsp.codelens.SetupOpts +function M.setup(opts) + Row.configure(opts or {}) + + vim.api.nvim_create_autocmd({ "BufEnter", "LspAttach" }, { + group = GROUP, + callback = function(ev) + local session = session_by_buf[ev.buf] + if session and session.enabled then + session:refresh() + end + end, + }) + + vim.api.nvim_create_autocmd({ "BufDelete", "BufWipeout" }, { + group = GROUP, + callback = function(ev) + local session = session_by_buf[ev.buf] + if session then + session:abort() + end + session_by_buf[ev.buf] = nil + end, + }) +end + +return M diff --git a/lua/lsp/codelens/refresh_task.lua b/lua/lsp/codelens/refresh_task.lua new file mode 100644 index 0000000..004a4bf --- /dev/null +++ b/lua/lsp/codelens/refresh_task.lua @@ -0,0 +1,134 @@ +local Row = require("lsp.codelens.row") +local log = require("log") + +---@class ow.lsp.codelens.RefreshTask +---@field session ow.lsp.codelens.Session +---@field aborted boolean +---@field private cancels fun()[] +local RefreshTask = {} +RefreshTask.__index = RefreshTask + +---@param session ow.lsp.codelens.Session +---@return ow.lsp.codelens.RefreshTask +function RefreshTask.new(session) + return setmetatable({ + session = session, + aborted = false, + cancels = {}, + }, RefreshTask) +end + +---@param cancel fun() +function RefreshTask:track(cancel) + table.insert(self.cancels, cancel) +end + +function RefreshTask:abort() + self.aborted = true + for _, cancel in ipairs(self.cancels) do + pcall(cancel) + end + self.cancels = {} +end + +---@param client vim.lsp.Client +---@param lens lsp.CodeLens +---@param row ow.lsp.codelens.Row +function RefreshTask:resolve(client, lens, row) + local method = vim.lsp.protocol.Methods.codeLens_resolve + local _, req_id = client:request( + method, + lens, + ---@param resolved lsp.CodeLens? + function(err, resolved) + if self.aborted then + return + end + if err then + log.warning( + "client %d: %s failed: %s", + client.id, + method, + err.message + ) + end + row:resolve(resolved) + row:render(self.session.buf) + end, + self.session.buf + ) + if req_id then + self:track(function() + client:cancel_request(req_id) + end) + end +end + +---@param rows table +---@param client vim.lsp.Client? +---@param lens lsp.CodeLens +function RefreshTask:process_lens(rows, client, lens) + local row = rows[lens.range.start.line] + if not row then + row = Row.new(lens.range.start.line) + rows[lens.range.start.line] = row + end + if lens.command then + row:add(lens) + return + end + if + not client + or not client:supports_method(vim.lsp.protocol.Methods.codeLens_resolve) + then + return + end + row:expect() + self:resolve(client, lens, row) +end + +---@param responses table +function RefreshTask:process_responses(responses) + local method = vim.lsp.protocol.Methods.textDocument_codeLens + local session = self.session + local new_rows = {} + for client_id, response in pairs(responses) do + if response.err then + log.warning( + "client %d: %s failed: %s", + client_id, + method, + response.err.message + ) + end + if not response.err and type(response.result) == "table" then + local client = vim.lsp.get_client_by_id(client_id) + for _, lens in ipairs(response.result) do + self:process_lens(new_rows, client, lens) + end + end + end + session.rows = new_rows + session:render() +end + +function RefreshTask:run() + local method = vim.lsp.protocol.Methods.textDocument_codeLens + local params = { + textDocument = vim.lsp.util.make_text_document_params(self.session.buf), + } + local cancel = vim.lsp.buf_request_all( + self.session.buf, + method, + params, + function(responses) + if self.aborted then + return + end + self:process_responses(responses) + end + ) + self:track(cancel) +end + +return RefreshTask diff --git a/lua/lsp/codelens/row.lua b/lua/lsp/codelens/row.lua new file mode 100644 index 0000000..1431738 --- /dev/null +++ b/lua/lsp/codelens/row.lua @@ -0,0 +1,155 @@ +local NS = vim.api.nvim_create_namespace("ow.lsp.codelens") + +---@alias ow.lsp.codelens.Position "eol" | "right_align" | "inline" | "above" + +---@type ow.lsp.codelens.Position +local position = "above" +local separator = " | " + +---@class ow.lsp.codelens.Row +---@field row integer +---@field lenses lsp.CodeLens[] +---@field pending integer +local Row = {} +Row.__index = Row + +---@param row integer +---@return ow.lsp.codelens.Row +function Row.new(row) + return setmetatable({ row = row, lenses = {}, pending = 0 }, Row) +end + +---@param lens lsp.CodeLens +function Row:add(lens) + table.insert(self.lenses, lens) +end + +function Row:expect() + self.pending = self.pending + 1 +end + +---@param lens? lsp.CodeLens +function Row:resolve(lens) + self.pending = self.pending - 1 + if lens then + table.insert(self.lenses, lens) + end +end + +---@return boolean +function Row:is_ready() + return self.pending == 0 and #self.lenses > 0 +end + +---@param buf integer +function Row:render(buf) + if not self:is_ready() or not vim.api.nvim_buf_is_valid(buf) then + return + end + + table.sort(self.lenses, function(a, b) + return a.range.start.character < b.range.start.character + end) + local parts = {} + for i, lens in ipairs(self.lenses) do + table.insert(parts, { lens.command.title, "LspCodeLens" }) + if i < #self.lenses then + table.insert(parts, { separator, "LspCodeLensSeparator" }) + end + end + + local col, opts + if position == "above" then + local line = vim.api.nvim_buf_get_lines( + buf, + self.row, + self.row + 1, + false + )[1] or "" + local indent = line:match("^%s*") or "" + local chunks = parts + if indent ~= "" then + chunks = { { indent, "LspCodeLensSeparator" } } + vim.list_extend(chunks, parts) + end + col = 0 + opts = { + virt_lines = { chunks }, + virt_lines_above = true, + hl_mode = "combine", + } + else + col = position == "inline" and self.lenses[1].range.start.character or 0 + opts = { + virt_text = parts, + virt_text_pos = position, + hl_mode = "combine", + } + end + + -- One extmark per row. Extmarks auto-shift with edits, so multiple from + -- prior refreshes can end up on the same row; pick the first and drop + -- the rest, then reuse its id for an atomic update. + local marks = vim.api.nvim_buf_get_extmarks( + buf, + NS, + { self.row, 0 }, + { self.row, -1 }, + { details = true } + ) + local primary + for i, mark in ipairs(marks) do + if i == 1 then + primary = { id = mark[1], col = mark[3], details = mark[4] } + else + vim.api.nvim_buf_del_extmark(buf, NS, mark[1]) + end + end + + if primary and primary.col == col and primary.details then + local d = primary.details --[[@as vim.api.keyset.extmark_details]] + local same = opts.virt_lines + and vim.deep_equal(d.virt_lines, opts.virt_lines) + and d.virt_lines_above == opts.virt_lines_above + or not opts.virt_lines + and d.virt_text_pos == opts.virt_text_pos + and vim.deep_equal(d.virt_text, opts.virt_text) + if same then + return + end + end + + opts.id = primary and primary.id or nil + vim.api.nvim_buf_set_extmark(buf, NS, self.row, col, opts) +end + +---@param buf integer +function Row.clear_all(buf) + vim.api.nvim_buf_clear_namespace(buf, NS, 0, -1) +end + +---@param buf integer +---@param rows table +function Row.prune(buf, rows) + for _, mark in ipairs(vim.api.nvim_buf_get_extmarks(buf, NS, 0, -1, {})) do + if not rows[mark[2]] then + vim.api.nvim_buf_del_extmark(buf, NS, mark[1]) + end + end +end + +---@class ow.lsp.codelens.row.ConfigureOpts +---@field position? ow.lsp.codelens.Position +---@field separator? string + +---@param opts ow.lsp.codelens.row.ConfigureOpts +function Row.configure(opts) + if opts.position ~= nil then + position = opts.position + end + if opts.separator ~= nil then + separator = opts.separator + end +end + +return Row diff --git a/lua/lsp/codelens/session.lua b/lua/lsp/codelens/session.lua new file mode 100644 index 0000000..f03ec79 --- /dev/null +++ b/lua/lsp/codelens/session.lua @@ -0,0 +1,109 @@ +local RefreshTask = require("lsp.codelens.refresh_task") +local Row = require("lsp.codelens.row") +local util = require("util") + +local REFRESH_DEBOUNCE_MS = 200 + +---@class ow.lsp.codelens.Session +---@field buf integer +---@field enabled boolean +---@field attached boolean +---@field rows table +---@field current? ow.lsp.codelens.RefreshTask +local Session = {} +Session.__index = Session + +---@param session ow.lsp.codelens.Session +local function do_refresh(session) + if not session.enabled then + return + end + session:abort() + local task = RefreshTask.new(session) + session.current = task + task:run() +end + +local refresher = util.keyed_debounce(do_refresh, REFRESH_DEBOUNCE_MS) + +---@param buf integer +---@return ow.lsp.codelens.Session +function Session.new(buf) + return setmetatable({ + buf = buf, + enabled = false, + attached = false, + rows = {}, + }, Session) +end + +function Session:abort() + refresher:cancel(self) + if self.current then + self.current:abort() + self.current = nil + end +end + +function Session:refresh() + refresher(self) +end + +function Session:render() + if not vim.api.nvim_buf_is_valid(self.buf) then + return + end + if not self.enabled then + Row.clear_all(self.buf) + return + end + + for _, row in pairs(self.rows) do + row:render(self.buf) + end + + Row.prune(self.buf, self.rows) +end + +function Session:attach() + if self.attached then + return + end + self.attached = true + vim.api.nvim_buf_attach(self.buf, false, { + on_lines = function() + if not self.enabled then + self.attached = false + return true + end + self:refresh() + end, + on_reload = function() + if self.enabled then + self:refresh() + end + end, + on_detach = function() + self.attached = false + end, + }) +end + +---@param value boolean +function Session:enable(value) + self.enabled = value + if value then + self:attach() + self:refresh() + else + self:abort() + self.rows = {} + self:render() + end +end + +function Session:toggle() + self:enable(not self.enabled) +end + +return Session