refactor(lsp): turn codelens module into Row/State classes

lua/lsp/codelens.lua

@@ -10,62 +10,70 @@ local M = {}
 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
+local refresh_debounced = util.debounce(function(buf)
----@return ow.lsp.codelens.State
+    local state = state_by_buf[buf]
-local function get_state(buf)
+    if state then
-    state_by_buf[buf] = state_by_buf[buf]
+        state:refresh()
-        or {
+    end
-            enabled = false,
+end, REFRESH_DEBOUNCE_MS)
-            attached = false,
+
-            generation = 0,
+---@class ow.lsp.codelens.Row
-            rows = {},
+---@field row integer
-        }
+---@field ready lsp.CodeLens[]
-    return state_by_buf[buf]
+---@field pending integer  count of unresolved lenses still in flight
+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
+
+--- Paint this row's extmark. Preserves the previous extmark while any lens
+--- on the row is still resolving so the user doesn't see a partial row
+--- (e.g. "refs" without "impls").
 ---@param buf integer
----@param row integer
+function Row:render(buf)
-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 self.pending > 0 or #self.ready == 0 then
-    if not state or not state.enabled then
         return
     end
 
-    local entry = state.rows[row]
+    table.sort(self.ready, function(a, b)
-    -- 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
+    for i, lens in ipairs(self.ready) do
         table.insert(parts, { lens.command.title, "LspCodeLens" })
-        if i < #entry.ready then
+        if i < #self.ready then
             table.insert(parts, { M.separator, "LspCodeLensSeparator" })
         end
     end
     local col = M.virt_text_pos == "inline"
-            and entry.ready[1].range.start.character
+            and self.ready[1].range.start.character
         or 0
 
     -- One extmark per row. Extmarks auto-shift with edits, so multiple from
@@ -74,8 +82,8 @@ local function render_row(buf, row)
     local marks = vim.api.nvim_buf_get_extmarks(
         buf,
         NS,
-        { row, 0 },
+        { self.row, 0 },
-        { row, -1 },
+        { self.row, -1 },
         { details = true }
     )
     local primary
@@ -97,7 +105,7 @@ local function render_row(buf, row)
         return
     end
 
-    vim.api.nvim_buf_set_extmark(buf, NS, row, col, {
+    vim.api.nvim_buf_set_extmark(buf, NS, self.row, col, {
         id = primary and primary.id or nil,
         virt_text = parts,
         virt_text_pos = M.virt_text_pos,
@@ -105,130 +113,161 @@ local function render_row(buf, row)
     })
 end
 
+---@class ow.lsp.codelens.State
+---@field buf integer
+---@field enabled boolean
+---@field attached boolean
+---@field generation integer
+---@field rows table<integer, ow.lsp.codelens.Row>
+local State = {}
+State.__index = State
+
 ---@param buf integer
-local function render_all(buf)
+---@return ow.lsp.codelens.State
-    if not vim.api.nvim_buf_is_valid(buf) then
+function State.new(buf)
+    return setmetatable({
+        buf = buf,
+        enabled = false,
+        attached = false,
+        generation = 0,
+        rows = {},
+    }, State)
+end
+
+--- Paint every row and drop extmarks on rows that are no longer present.
+function State:render()
+    if not vim.api.nvim_buf_is_valid(self.buf) then
         return
     end
-    local state = state_by_buf[buf]
+    if not self.enabled then
-    if not state or not state.enabled then
+        vim.api.nvim_buf_clear_namespace(self.buf, NS, 0, -1)
-        vim.api.nvim_buf_clear_namespace(buf, NS, 0, -1)
         return
     end
 
-    for row in pairs(state.rows) do
+    for _, row in pairs(self.rows) do
-        render_row(buf, row)
+        row:render(self.buf)
     end
 
-    -- Drop extmarks on rows that are no longer in state (lens removed).
+    for _, mark in
-    for _, mark in ipairs(vim.api.nvim_buf_get_extmarks(buf, NS, 0, -1, {})) do
+        ipairs(vim.api.nvim_buf_get_extmarks(self.buf, NS, 0, -1, {}))
-        if not state.rows[mark[2]] then
+    do
-            vim.api.nvim_buf_del_extmark(buf, NS, mark[1])
+        if not self.rows[mark[2]] then
+            vim.api.nvim_buf_del_extmark(self.buf, NS, mark[1])
         end
     end
 end
 
----@param buf integer
+function State:refresh()
-local function refresh(buf)
+    if not self.enabled then
-    local state = get_state(buf)
-    if not state.enabled then
         return
     end
 
-    state.generation = state.generation + 1
+    self.generation = self.generation + 1
-    local gen = state.generation
+    local gen = self.generation
-    local params =
+    local params = {
-        { textDocument = vim.lsp.util.make_text_document_params(buf) }
+        textDocument = vim.lsp.util.make_text_document_params(self.buf),
+    }
     local new_rows = {}
 
     vim.lsp.buf_request_all(
-        buf,
+        self.buf,
         vim.lsp.protocol.Methods.textDocument_codeLens,
         params,
         function(results)
-            if state.generation ~= gen then
+            if self.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
+                        local row_num = lens.range.start.line
-                        new_rows[row] = new_rows[row]
+                        local row = new_rows[row_num]
-                            or { ready = {}, pending = 0 }
+                        if not row then
+                            row = Row.new(row_num)
+                            new_rows[row_num] = row
+                        end
                         if lens.command then
-                            table.insert(new_rows[row].ready, lens)
+                            row:add(lens)
                         elseif
                             client
                             and client:supports_method(
                                 vim.lsp.protocol.Methods.codeLens_resolve
                             )
                         then
-                            new_rows[row].pending = new_rows[row].pending + 1
+                            row:await()
                             client:request(
                                 vim.lsp.protocol.Methods.codeLens_resolve,
                                 lens,
                                 function(_, resolved)
-                                    if state.generation ~= gen then
+                                    if self.generation ~= gen then
                                         return
                                     end
-                                    local entry = new_rows[row]
+                                    row:resolved(
-                                    entry.pending = entry.pending - 1
+                                        type(resolved) == "table" and resolved
-                                    if type(resolved) == "table" then
+                                            or nil
-                                        table.insert(entry.ready, resolved)
+                                    )
-                                    end
+                                    row:render(self.buf)
-                                    render_row(buf, row)
                                 end,
-                                buf
+                                self.buf
                             )
                         end
                     end
                 end
             end
-            -- Swap in the new row map and paint everything once. Rows with
+            self.rows = new_rows
-            -- pending resolves keep their prior extmark in place; each
+            self:render()
-            -- 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)
+function State:attach()
-
+    if self.attached then
----@param buf integer
-local function attach_buf(buf)
-    local state = get_state(buf)
-    if state.attached then
         return
     end
-    state.attached = true
+    self.attached = true
-    vim.api.nvim_buf_attach(buf, false, {
+    vim.api.nvim_buf_attach(self.buf, false, {
-        on_lines = function(_, b)
+        on_lines = function()
-            local s = state_by_buf[b]
+            if not self.enabled then
-            if not s or not s.enabled then
+                self.attached = false
-                if s then
-                    s.attached = false
-                end
                 return true
             end
-            refresh_debounced:call(b)
+            refresh_debounced:call(self.buf)
         end,
-        on_reload = function(_, b)
+        on_reload = function()
-            local s = state_by_buf[b]
+            if self.enabled then
-            if s and s.enabled then
+                refresh_debounced:call(self.buf)
-                refresh_debounced:call(b)
             end
         end,
-        on_detach = function(_, b)
+        on_detach = function()
-            local s = state_by_buf[b]
+            self.attached = false
-            if s then
-                s.attached = false
-            end
         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)
@@ -240,24 +279,17 @@ 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)
+    buf = buf or vim.api.nvim_get_current_buf()
-    state.enabled = value
+    get_state(buf):enable(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)
+    buf = buf or vim.api.nvim_get_current_buf()
+    get_state(buf):toggle()
 end
 
 function M.setup()
@@ -266,7 +298,7 @@ function M.setup()
         callback = function(ev)
             local state = state_by_buf[ev.buf]
             if state and state.enabled then
-                refresh(ev.buf)
+                state:refresh()
             end
         end,
     })