refactor(ui): use virt_lines to display added lines (#11)

Author: Xuyuanp

lua/copilot-lsp/nes/ui.lua

@@ -1,10 +1,8 @@
 local M = {}
 
 ---@param bufnr integer
----@param suggestion_ui nes.EditSuggestionUI
 ---@param ns_id integer
-local function _dismiss_suggestion_ui(bufnr, suggestion_ui, ns_id)
-    pcall(vim.api.nvim_win_close, suggestion_ui.preview_winnr, true)
+local function _dismiss_suggestion(bufnr, ns_id)
     pcall(vim.api.nvim_buf_clear_namespace, bufnr, ns_id, 0, -1)
 end
 
@@ -19,7 +17,7 @@ function M.clear_suggestion(bufnr, ns_id)
         return
     end
 
-    _dismiss_suggestion_ui(bufnr, state.ui, ns_id)
+    _dismiss_suggestion(bufnr, ns_id)
     vim.b[bufnr].nes_state = nil
 end
 
@@ -71,21 +69,13 @@ function M._calculate_lines(suggestion)
         virt_lines_count = added_lines_count,
     }
 
-    -- Calculate positions for floating window
-    ---@type nes.FloatWin
-    local float_win = {
-        height = #added_lines,
-        row = suggestion.range["end"].line + deleted_lines_count + (suggestion.range["end"].character ~= 0 and 1 or 0),
-    }
-
     return {
         deleted_lines_count = deleted_lines_count,
         added_lines = added_lines,
         added_lines_count = added_lines_count,
         same_line = same_line,
         delete_extmark = delete_extmark,
         virt_lines_extmark = virt_lines_extmark,
-        float_win = float_win,
     }
 end
 
@@ -106,7 +96,6 @@ function M._display_next_suggestion(edits, ns_id)
     local win_id = vim.fn.win_findbuf(bufnr)[1]
     local suggestion = edits[1]
 
-    local ui = {}
     local lines = M._calculate_lines(suggestion)
 
     if lines.deleted_lines_count > 0 then
@@ -117,45 +106,14 @@ function M._display_next_suggestion(edits, ns_id)
         })
     end
     if lines.added_lines_count > 0 then
-        -- Create space for float
-        local virt_lines = {}
-        for _ = 1, lines.virt_lines_extmark.virt_lines_count do
-            table.insert(virt_lines, {
-                { "", "Normal" },
-            })
-        end
+        local text = trim_end(edits[1].text)
+        local virt_lines = require("copilot-lsp.util").hl_text_to_virt_lines(text, vim.bo[bufnr].filetype)
+
         vim.api.nvim_buf_set_extmark(bufnr, ns_id, lines.virt_lines_extmark.row, 0, {
             virt_lines = virt_lines,
         })
-
-        local preview_bufnr = vim.api.nvim_create_buf(false, true)
-        vim.api.nvim_buf_set_lines(preview_bufnr, 0, -1, false, lines.added_lines)
-        vim.bo[preview_bufnr].modifiable = false
-        vim.bo[preview_bufnr].buflisted = false
-        vim.bo[preview_bufnr].bufhidden = "wipe"
-        vim.bo[preview_bufnr].filetype = vim.bo[bufnr].filetype
-
-        local cursor = vim.api.nvim_win_get_cursor(win_id)
-        local win_width = vim.api.nvim_win_get_width(win_id)
-        local offset = vim.fn.getwininfo(win_id)[1].textoff
-        local preview_winnr = vim.api.nvim_open_win(preview_bufnr, false, {
-            relative = "cursor",
-            width = win_width - offset,
-            height = lines.float_win.height,
-            row = lines.float_win.row - cursor[1],
-            col = 0,
-            style = "minimal",
-            border = "none",
-        })
-        vim.wo[preview_winnr].number = false
-        vim.wo[preview_winnr].winhighlight = "Normal:NesAdd"
-        vim.wo[preview_winnr].winblend = 0
-
-        ui.preview_winnr = preview_winnr
     end
 
-    suggestion.ui = ui
-
     vim.b[bufnr].nes_state = suggestion
 
     vim.api.nvim_create_autocmd({ "CursorMoved", "CursorMovedI" }, {

lua/copilot-lsp/types.lua

@@ -4,7 +4,6 @@
 ---@field text string
 ---@field newText string
 ---@field textDocument lsp.VersionedTextDocumentIdentifier
----@field ui nes.EditSuggestionUI?
 
 ---@class copilotlsp.copilotInlineEditResponse
 ---@field edits copilotlsp.InlineEdit[]
@@ -22,11 +21,6 @@
 ---@field row number
 ---@field virt_lines_count number
 
----@class nes.FloatWin
---- Defines dimensions and position for the floating window.
----@field height number
----@field row number
-
 ---@class nes.LineCalculationResult
 --- The result of calculating lines for inline suggestion UI.
 ---@field deleted_lines_count number
@@ -35,4 +29,3 @@
 ---@field same_line number
 ---@field delete_extmark nes.DeleteExtmark
 ---@field virt_lines_extmark nes.VirtLinesExtmark
----@field float_win nes.FloatWin

lua/copilot-lsp/util.lua

@@ -22,4 +22,185 @@ function M.debounce(fn, delay)
     end
 end
 
+---@private
+---@class Capture
+---@field hl string|string[]
+---range(0-based)
+---@field start_row integer
+---@field start_col integer
+---@field end_row integer
+---@field end_col integer
+
+---@param text string
+---@param lang string
+---@return Capture[]?
+local function parse_text(text, lang)
+    local ok, trees = pcall(vim.treesitter.get_string_parser, text, lang)
+    if not ok then
+        return
+    end
+    trees:parse(true)
+
+    local captures = {}
+
+    trees:for_each_tree(function(tree, _ltree)
+        local hl_query = vim.treesitter.query.get(lang, "highlights")
+        if not hl_query then
+            return
+        end
+
+        local iter = hl_query:iter_captures(tree:root(), text)
+        vim.iter(iter):each(function(id, node)
+            local name = hl_query.captures[id]
+            local hl = "Normal"
+            if not vim.startswith(name, "_") then
+                hl = "@" .. name .. "." .. lang
+            end
+            local start_row, start_col, end_row, end_col = node:range()
+
+            -- Ignore zero-width captures if they cause issues (sometimes happen at EOF)
+            if start_row == end_row and start_col == end_col then
+                return
+            end
+
+            table.insert(captures, {
+                hl = hl,
+                start_row = start_row,
+                start_col = start_col,
+                end_row = end_row,
+                end_col = end_col,
+            })
+        end)
+    end)
+    return captures
+end
+
+local function merge_captures(captures)
+    table.sort(captures, function(a, b)
+        if a.start_row == b.start_row then
+            return a.start_col < b.start_col
+        end
+        return a.start_row < b.start_row
+    end)
+    local merged_captures = {}
+    for i = 2, #captures do
+        local prev = captures[i - 1]
+        local curr = captures[i]
+        if
+            prev.start_row == curr.start_row
+            and prev.start_col == curr.start_col
+            and prev.end_row == curr.end_row
+            and prev.end_col == curr.end_col
+        then
+            local prev_hl = type(prev.hl) == "table" and prev.hl or { prev.hl }
+            local curr_hl = type(curr.hl) == "table" and curr.hl or { curr.hl }
+            ---@diagnostic disable-next-line: param-type-mismatch
+            vim.list_extend(prev_hl, curr_hl)
+            curr.hl = prev_hl
+        else
+            table.insert(merged_captures, prev)
+        end
+    end
+    table.insert(merged_captures, captures[#captures])
+
+    return merged_captures
+end
+
+function M.hl_text_to_virt_lines(text, lang)
+    local lines = vim.split(text, "\n")
+    local normal_hl = "Normal"
+    local bg_hl = "NesAdd"
+
+    local function hl_chunk(chunk, hl)
+        if not hl then
+            return { chunk, { normal_hl, bg_hl } }
+        end
+        if type(hl) == "string" then
+            return { chunk, { hl, bg_hl } }
+        end
+        hl = vim.deepcopy(hl)
+        table.insert(hl, bg_hl)
+        return { chunk, hl }
+    end
+
+    local captures = parse_text(text, lang)
+    if not captures or #captures == 0 then
+        return vim.iter(lines)
+            :map(function(line)
+                return { hl_chunk(line) }
+            end)
+            :totable()
+    end
+
+    captures = merge_captures(captures)
+
+    local virt_lines = {}
+
+    local curr_row = 0
+    local curr_col = 0
+    local curr_virt_line = {}
+
+    vim.iter(captures):each(function(cap)
+        -- skip if the capture is before the current position
+        if cap.end_row < curr_row or (cap.end_row == curr_row and cap.end_col < curr_col) then
+            return
+        end
+
+        if cap.start_row > curr_row then
+            -- add the rest of the line
+            local chunk_text = string.sub(lines[curr_row + 1], curr_col + 1)
+            table.insert(curr_virt_line, hl_chunk(chunk_text))
+            table.insert(virt_lines, curr_virt_line)
+
+            for i = curr_row + 1, cap.start_row - 1 do
+                local line_text = lines[i + 1]
+                table.insert(virt_lines, { hl_chunk(line_text) })
+            end
+
+            curr_row = cap.start_row
+            curr_col = 0
+            curr_virt_line = {}
+        end
+
+        assert(curr_row == cap.start_row, "Unexpected start row")
+
+        if cap.start_col > curr_col then
+            local chunk_text = string.sub(lines[curr_row + 1], curr_col + 1, cap.start_col)
+            table.insert(curr_virt_line, hl_chunk(chunk_text))
+            curr_col = cap.start_col
+        end
+
+        assert(curr_col == cap.start_col, "Unexpected start column")
+
+        if cap.end_row > curr_row then
+            local chunk_text = string.sub(lines[curr_row + 1], curr_col + 1)
+            table.insert(curr_virt_line, hl_chunk(chunk_text, cap.hl))
+            table.insert(virt_lines, curr_virt_line)
+
+            for i in curr_row + 1, cap.end_row - 1 do
+                local line_text = lines[i + 1]
+                table.insert(virt_lines, { hl_chunk(line_text, cap.hl) })
+            end
+
+            curr_row = cap.end_row
+            curr_col = 0
+            curr_virt_line = {}
+        end
+
+        assert(curr_row == cap.end_row, "Unexpected end row")
+
+        if cap.end_col > curr_col then
+            local chunk_text = string.sub(lines[curr_row + 1], curr_col + 1, cap.end_col)
+            table.insert(curr_virt_line, hl_chunk(chunk_text, cap.hl))
+            curr_col = cap.end_col
+        end
+    end)
+
+    if #curr_virt_line > 0 then
+        table.insert(virt_lines, curr_virt_line)
+    end
+
+    return virt_lines
+end
+
 return M

tests/nes/test_ui.lua

@@ -35,10 +35,6 @@ T["diff placement calculations"] = MiniTest.new_set({
                     row = 0,
                 },
                 deleted_lines_count = 1,
-                float_win = {
-                    height = 1,
-                    row = 2,
-                },
                 same_line = 1,
                 virt_lines_extmark = {
                     row = 0,
@@ -75,10 +71,6 @@ T["diff placement calculations"] = MiniTest.new_set({
                     row = 0,
                 },
                 deleted_lines_count = 1,
-                float_win = {
-                    height = 1,
-                    row = 2,
-                },
                 same_line = 0,
                 virt_lines_extmark = {
                     row = 0,
@@ -115,10 +107,6 @@ T["diff placement calculations"] = MiniTest.new_set({
                     row = 0,
                 },
                 deleted_lines_count = 1,
-                float_win = {
-                    height = 2,
-                    row = 2,
-                },
                 same_line = 0,
                 virt_lines_extmark = {
                     row = 0,
@@ -149,10 +137,6 @@ T["diff placement calculations"] = MiniTest.new_set({
                     end_row = 2,
                     row = 2,
                 },
-                float_win = {
-                    height = 1,
-                    row = 3,
-                },
                 same_line = 1,
                 virt_lines_extmark = {
                     row = 1,

tests/screenshots/tests-nes-test_nes.lua---nes---add-only-edit-002

@@ -27,7 +27,7 @@
 --|---------|---------|---------|---------|---------|---------|---------|---------|
 01|00000000000000000000000000000000000000000000000000000000000000000000000000000000
 02|00000000000000000000000000000000000000000000000000000000000000000000000000000000
-03|11111111111111111111111111111111111111111111111111111111111111111111111111111111
+03|11111100000000000000000000000000000000000000000000000000000000000000000000000000
 04|00000000000000000000000000000000000000000000000000000000000000000000000000000000
 05|22222222222222222222222222222222222222222222222222222222222222222222222222222222
 06|22222222222222222222222222222222222222222222222222222222222222222222222222222222

tests/screenshots/tests-nes-test_nes.lua---nes---multi-line-edit-002

@@ -27,8 +27,8 @@
 --|---------|---------|---------|---------|---------|---------|---------|---------|
 01|00000000111111111111111111111111111111111111111111111111111111111111111111111111
 02|00000000111111111111111111111111111111111111111111111111111111111111111111111111
-03|22222222222222222222222222222222222222222222222222222222222222222222222222222222
-04|22222222222222222222222222222222222222222222222222222222222222222222222222222222
+03|22222222222211111111111111111111111111111111111111111111111111111111111111111111
+04|22222222222211111111111111111111111111111111111111111111111111111111111111111111
 05|11111111111111111111111111111111111111111111111111111111111111111111111111111111
 06|33333333333333333333333333333333333333333333333333333333333333333333333333333333
 07|33333333333333333333333333333333333333333333333333333333333333333333333333333333

tests/screenshots/tests-nes-test_nes.lua---nes---same-line-edit-002

@@ -26,7 +26,7 @@
 
 --|---------|---------|---------|---------|---------|---------|---------|---------|
 01|00011111111111111111111111111111111111111111111111111111111111111111111111111111
-02|22222222222222222222222222222222222222222222222222222222222222222222222222222222
+02|22211111111111111111111111111111111111111111111111111111111111111111111111111111
 03|11111111111111111111111111111111111111111111111111111111111111111111111111111111
 04|11111111111111111111111111111111111111111111111111111111111111111111111111111111
 05|33333333333333333333333333333333333333333333333333333333333333333333333333333333