feat(lsp): custom codelens as virtual text

lua/core/keymap.lua

@@ -161,6 +161,9 @@ vim.keymap.set({ "i", "s" }, "<C-s>", function()
 end)
 vim.keymap.set({ "n", "i" }, "<C-h>", vim.lsp.buf.document_highlight)
 vim.keymap.set({ "n", "x" }, "<leader>lf", vim.lsp.buf.format)
+vim.keymap.set("n", "grl", function()
+    require("lsp.codelens").toggle()
+end)
 vim.keymap.set("n", "<leader>ld", function()
     vim.diagnostic.enable(
         not vim.diagnostic.is_enabled({ bufnr = 0 }),

lua/lsp/codelens.lua

@@ -0,0 +1,283 @@
+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 = {}
+
+---@type "eol" | "right_align" | "inline"
+M.virt_text_pos = "eol"
+M.separator = " | "
+
+---@class ow.lsp.codelens.Row
+---@field ready lsp.CodeLens[]  resolved lenses with `command`, ready to paint
+---@field pending integer       count of unresolved lenses still in flight
+
+---@class ow.lsp.codelens.State
+---@field enabled boolean
+---@field attached boolean
+---@field generation integer
+---@field rows table<integer, ow.lsp.codelens.Row>
+
+---@type table<integer, ow.lsp.codelens.State>
+local state_by_buf = {}
+
+---@param buf integer
+---@return ow.lsp.codelens.State
+local function get_state(buf)
+    state_by_buf[buf] = state_by_buf[buf]
+        or {
+            enabled = false,
+            attached = false,
+            generation = 0,
+            rows = {},
+        }
+    return state_by_buf[buf]
+end
+
+---@param buf integer
+---@param row integer
+local function render_row(buf, row)
+    if not vim.api.nvim_buf_is_valid(buf) then
+        return
+    end
+    local state = state_by_buf[buf]
+    if not state or not state.enabled then
+        return
+    end
+
+    local entry = state.rows[row]
+    -- Preserve existing extmark while any lens on the row is still resolving
+    -- so the user doesn't see a partial row (e.g. "refs" without "impls").
+    if not entry or entry.pending > 0 or #entry.ready == 0 then
+        return
+    end
+
+    table.sort(entry.ready, function(a, b)
+        return a.range.start.character < b.range.start.character
+    end)
+    local parts = {}
+    for i, lens in ipairs(entry.ready) do
+        table.insert(parts, { lens.command.title, "LspCodeLens" })
+        if i < #entry.ready then
+            table.insert(parts, { M.separator, "LspCodeLensSeparator" })
+        end
+    end
+    local col = M.virt_text_pos == "inline"
+            and entry.ready[1].range.start.character
+        or 0
+
+    -- 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,
+        { row, 0 },
+        { 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
+        and primary.details.virt_text_pos == M.virt_text_pos
+        and vim.deep_equal(primary.details.virt_text, parts)
+    then
+        return
+    end
+
+    vim.api.nvim_buf_set_extmark(buf, NS, row, col, {
+        id = primary and primary.id or nil,
+        virt_text = parts,
+        virt_text_pos = M.virt_text_pos,
+        hl_mode = "combine",
+    })
+end
+
+---@param buf integer
+local function render_all(buf)
+    if not vim.api.nvim_buf_is_valid(buf) then
+        return
+    end
+    local state = state_by_buf[buf]
+    if not state or not state.enabled then
+        vim.api.nvim_buf_clear_namespace(buf, NS, 0, -1)
+        return
+    end
+
+    for row in pairs(state.rows) do
+        render_row(buf, row)
+    end
+
+    -- Drop extmarks on rows that are no longer in state (lens removed).
+    for _, mark in ipairs(vim.api.nvim_buf_get_extmarks(buf, NS, 0, -1, {})) do
+        if not state.rows[mark[2]] then
+            vim.api.nvim_buf_del_extmark(buf, NS, mark[1])
+        end
+    end
+end
+
+---@param buf integer
+local function refresh(buf)
+    local state = get_state(buf)
+    if not state.enabled then
+        return
+    end
+
+    state.generation = state.generation + 1
+    local gen = state.generation
+    local params =
+        { textDocument = vim.lsp.util.make_text_document_params(buf) }
+    local new_rows = {}
+
+    vim.lsp.buf_request_all(
+        buf,
+        vim.lsp.protocol.Methods.textDocument_codeLens,
+        params,
+        function(results)
+            if state.generation ~= gen then
+                return
+            end
+            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
+                        local row = lens.range.start.line
+                        new_rows[row] = new_rows[row]
+                            or { ready = {}, pending = 0 }
+                        if lens.command then
+                            table.insert(new_rows[row].ready, lens)
+                        elseif
+                            client
+                            and client:supports_method(
+                                vim.lsp.protocol.Methods.codeLens_resolve
+                            )
+                        then
+                            new_rows[row].pending = new_rows[row].pending + 1
+                            client:request(
+                                vim.lsp.protocol.Methods.codeLens_resolve,
+                                lens,
+                                function(_, resolved)
+                                    if state.generation ~= gen then
+                                        return
+                                    end
+                                    local entry = new_rows[row]
+                                    entry.pending = entry.pending - 1
+                                    if type(resolved) == "table" then
+                                        table.insert(entry.ready, resolved)
+                                    end
+                                    render_row(buf, row)
+                                end,
+                                buf
+                            )
+                        end
+                    end
+                end
+            end
+            -- Swap in the new row map and paint everything once. Rows with
+            -- pending resolves keep their prior extmark in place; each
+            -- resolve callback above repaints its own row when it arrives.
+            state.rows = new_rows
+            render_all(buf)
+        end
+    )
+end
+
+local refresh_debounced = util.debounce(refresh, REFRESH_DEBOUNCE_MS)
+
+---@param buf integer
+local function attach_buf(buf)
+    local state = get_state(buf)
+    if state.attached then
+        return
+    end
+    state.attached = true
+    vim.api.nvim_buf_attach(buf, false, {
+        on_lines = function(_, b)
+            local s = state_by_buf[b]
+            if not s or not s.enabled then
+                if s then
+                    s.attached = false
+                end
+                return true
+            end
+            refresh_debounced:call(b)
+        end,
+        on_reload = function(_, b)
+            local s = state_by_buf[b]
+            if s and s.enabled then
+                refresh_debounced:call(b)
+            end
+        end,
+        on_detach = function(_, b)
+            local s = state_by_buf[b]
+            if s then
+                s.attached = false
+            end
+        end,
+    })
+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)
+    buf = buf or vim.api.nvim_get_current_buf()
+    if value == nil then
+        value = true
+    end
+    local state = get_state(buf)
+    state.enabled = value
+    if value then
+        attach_buf(buf)
+        refresh(buf)
+    else
+        refresh_debounced:cancel(buf)
+        render_all(buf)
+    end
+end
+
+---@param buf? integer
+function M.toggle(buf)
+    M.enable(not M.is_enabled(buf), buf)
+end
+
+function M.setup()
+    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
+                refresh(ev.buf)
+            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

lua/lsp/init.lua

@@ -1,3 +1,4 @@
+local codelens = require("lsp.codelens")
 local completion = require("lsp.completion")
 local diagnostic = require("lsp.diagnostic")
 local log = require("log")
@@ -102,6 +103,7 @@ function M.setup()
     })
 
     completion.setup()
+    codelens.setup()
 
     vim.lsp.log.set_level(vim.log.levels.WARN)
 end