feat(ui): inline edit preview (#34)

* feat(ui): inline edit

* feat(ui): define preview for inline edits

* feat(ui): added preview

* feat: support lines insertion above and below the current line

* fix: tests

* chore: removed unused code

Author: Xuyuanp

lsp/copilot_ls.lua

@@ -30,10 +30,6 @@ return {
     }),
     root_dir = vim.uv.cwd(),
     on_init = function(client)
-        vim.api.nvim_set_hl(0, "CopilotLspNesAdd", { link = "DiffAdd", default = true })
-        vim.api.nvim_set_hl(0, "CopilotLspNesDelete", { link = "DiffDelete", default = true })
-        vim.api.nvim_set_hl(0, "CopilotLspNesApply", { link = "DiffText", default = true })
-
         local au = vim.api.nvim_create_augroup("copilotlsp.init", { clear = true })
         --NOTE: Inline Completions
         --TODO: We dont currently use this code path, so comment for now until a UI is built

lua/copilot-lsp/nes/init.lua

@@ -8,7 +8,8 @@ local nes_ns = vim.api.nvim_create_namespace("copilotlsp.nes")
 
 ---@param err lsp.ResponseError?
 ---@param result copilotlsp.copilotInlineEditResponse
-local function handle_nes_response(err, result)
+---@param ctx lsp.HandlerContext
+local function handle_nes_response(err, result, ctx)
     if err then
         -- vim.notify(err.message)
         return
@@ -17,14 +18,17 @@ local function handle_nes_response(err, result)
         --- Convert to textEdit fields
         edit.newText = edit.text
     end
-    nes_ui._display_next_suggestion(result.edits, nes_ns)
+    nes_ui._display_next_suggestion(ctx.bufnr, nes_ns, result.edits)
 end
 
 --- Requests the NextEditSuggestion from the current cursor position
----@param copilot_lss vim.lsp.Client?
+---@param copilot_lss? vim.lsp.Client|string
 function M.request_nes(copilot_lss)
     local pos_params = vim.lsp.util.make_position_params(0, "utf-16")
     local version = vim.lsp.util.buf_versions[vim.api.nvim_get_current_buf()]
+    if type(copilot_lss) == "string" then
+        copilot_lss = vim.lsp.get_clients({ name = copilot_lss })[1]
+    end
     assert(copilot_lss, errs.ErrNotStarted)
     ---@diagnostic disable-next-line: inject-field
     pos_params.textDocument.version = version
@@ -109,4 +113,10 @@ function M.apply_pending_nes(bufnr)
     return true
 end
 
+---@param bufnr? integer
+function M.clear_suggestion(bufnr)
+    bufnr = bufnr and bufnr > 0 and bufnr or vim.api.nvim_get_current_buf()
+    nes_ui.clear_suggestion(bufnr, nes_ns)
+end
+
 return M

lua/copilot-lsp/nes/state.lua

@@ -14,108 +14,159 @@ function M.clear_suggestion(bufnr, ns_id)
         vim.b[bufnr].nes_jump = false
         return
     end
-    vim.api.nvim_buf_clear_namespace(bufnr, ns_id, 0, -1)
+    _dismiss_suggestion(bufnr, ns_id)
     ---@type copilotlsp.InlineEdit
     local state = vim.b[bufnr].nes_state
     if not state then
         return
     end
 
-    _dismiss_suggestion(bufnr, ns_id)
     vim.b[bufnr].nes_state = nil
 end
 
-local function trim_end(s)
-    return s:gsub("%s+$", "")
-end
-
 ---@private
----@param suggestion copilotlsp.InlineEdit
----@return copilotlsp.nes.LineCalculationResult
-function M._calculate_lines(suggestion)
-    local deleted_lines_count = suggestion.range["end"].line - suggestion.range.start.line
-    local added_lines = vim.split(trim_end(suggestion.newText), "\n")
-    local added_lines_count = suggestion.newText == "" and 0 or #added_lines
-    local same_line = 0
-
-    if deleted_lines_count == 0 and added_lines_count == 1 then
-        ---changing within line
-        deleted_lines_count = 1
-        same_line = 1
+---@param bufnr integer
+---@param edit lsp.TextEdit
+---@return copilotlsp.nes.InlineEditPreview
+function M._calculate_preview(bufnr, edit)
+    local text = edit.newText
+    local range = edit.range
+    local start_line = range.start.line
+    local start_char = range.start.character
+    local end_line = range["end"].line
+    local end_char = range["end"].character
+
+    -- Split text by newline. Use plain=true to handle trailing newline correctly.
+    local new_lines = vim.split(text, "\n", { plain = true })
+    local num_new_lines = #new_lines
+
+    local old_lines = vim.api.nvim_buf_get_lines(bufnr, start_line, end_line + 1, false)
+    local num_old_lines = #old_lines
+
+    local is_same_line = start_line == end_line
+    local is_deletion = text == ""
+    local lines_edit = is_same_line or (start_char == 0 and end_char == 0)
+    local is_insertion = is_same_line and start_char == end_char
+
+    if is_deletion and is_insertion then
+        -- no-op
+        return {}
     end
 
-    -- if
-    --     suggestion.range.start.line == suggestion.range["end"].line
-    --     and suggestion.range.start.character == suggestion.range["end"].character
-    -- then
-    --     --add only
-    --     TODO: Do we need to position specifically for add only?
-    --     UI tests seem to say no
-    -- end
-
-    -- Calculate positions for delete highlight extmark
-    ---@type copilotlsp.nes.DeleteExtmark
-    local delete_extmark = {
-        row = suggestion.range.start.line,
-        end_row = (
-            suggestion.range["end"].character ~= 0 and suggestion.range["end"].line + 1
-            or suggestion.range["end"].line
-        ),
-    }
+    if is_deletion and lines_edit then
+        return {
+            deletion = {
+                range = edit.range,
+            },
+        }
+    end
 
-    -- Calculate positions for virtual lines extmark
-    ---@type copilotlsp.nes.AddExtmark
-    local virt_lines_extmark = {
-        row = (
-            suggestion.range["end"].character ~= 0 and suggestion.range["end"].line
-            or suggestion.range["end"].line - 1
-        ),
-        virt_lines_count = added_lines_count,
-    }
+    if is_insertion and num_new_lines == 1 and text ~= "" then
+        -- inline insertion
+        return {
+            inline_insertion = {
+                text = text,
+                line = start_line,
+                character = start_char,
+            },
+        }
+    end
+
+    if is_insertion and num_new_lines > 1 then
+        if start_char == #old_lines[1] and new_lines[1] == "" then
+            -- insert lines after the start line
+            return {
+                lines_insertion = {
+                    text = table.concat(vim.list_slice(new_lines, 2), "\n"),
+                    line = start_line,
+                },
+            }
+        end
+
+        if end_char == 0 and new_lines[num_new_lines] == "" then
+            -- insert lines before the end line
+            return {
+                lines_insertion = {
+                    text = table.concat(vim.list_slice(new_lines, 1, num_new_lines - 1), "\n"),
+                    line = start_line,
+                    above = true,
+                },
+            }
+        end
+    end
+
+    -- insert lines in the middle
+    local prefix = old_lines[1]:sub(1, start_char)
+    local suffix = old_lines[num_old_lines]:sub(end_char + 1)
+    local new_lines_extend = vim.deepcopy(new_lines)
+    new_lines_extend[1] = prefix .. new_lines_extend[1]
+    new_lines_extend[num_new_lines] = new_lines_extend[num_new_lines] .. suffix
+    local insertion = table.concat(new_lines_extend, "\n")
 
     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,
+        deletion = {
+            range = {
+                start = { line = start_line, character = 0 },
+                ["end"] = { line = end_line, character = #old_lines[num_old_lines] },
+            },
+        },
+        lines_insertion = {
+            text = insertion,
+            line = end_line,
+        },
     }
 end
 
 ---@private
----@param edits copilotlsp.InlineEdit[]
+---@param bufnr integer
 ---@param ns_id integer
-function M._display_next_suggestion(edits, ns_id)
-    local state = vim.b[vim.api.nvim_get_current_buf()].nes_state
-    if state then
-        M.clear_suggestion(vim.api.nvim_get_current_buf(), ns_id)
-    end
-
-    if not edits or #edits == 0 then
-        -- vim.notify("No suggestion available", vim.log.levels.INFO)
-        return
+---@param preview copilotlsp.nes.InlineEditPreview
+function M._display_preview(bufnr, ns_id, preview)
+    if preview.deletion then
+        local range = preview.deletion.range
+        vim.api.nvim_buf_set_extmark(bufnr, ns_id, range.start.line, range.start.character, {
+            hl_group = "CopilotLspNesDelete",
+            end_row = range["end"].line,
+            end_col = range["end"].character,
+        })
     end
-    local bufnr = vim.uri_to_bufnr(edits[1].textDocument.uri)
-    local suggestion = edits[1]
 
-    local lines = M._calculate_lines(suggestion)
-
-    if lines.deleted_lines_count > 0 then
-        -- Deleted range red highlight
-        vim.api.nvim_buf_set_extmark(bufnr, ns_id, lines.delete_extmark.row, 0, {
-            hl_group = "CopilotLspNesDelete",
-            end_row = lines.delete_extmark.end_row,
+    local inline_insertion = preview.inline_insertion
+    if inline_insertion then
+        local virt_lines =
+            require("copilot-lsp.util").hl_text_to_virt_lines(inline_insertion.text, vim.bo[bufnr].filetype)
+        vim.api.nvim_buf_set_extmark(bufnr, ns_id, inline_insertion.line, inline_insertion.character, {
+            virt_text = virt_lines[1],
+            virt_text_pos = "inline",
         })
     end
-    if lines.added_lines_count > 0 then
-        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, {
+    local lines_insertion = preview.lines_insertion
+    if lines_insertion then
+        local virt_lines =
+            require("copilot-lsp.util").hl_text_to_virt_lines(lines_insertion.text, vim.bo[bufnr].filetype)
+        vim.api.nvim_buf_set_extmark(bufnr, ns_id, lines_insertion.line, 0, {
             virt_lines = virt_lines,
+            virt_lines_above = lines_insertion.above,
         })
     end
+end
+
+---@private
+---@param bufnr integer
+---@param ns_id integer
+---@param edits copilotlsp.InlineEdit[]
+function M._display_next_suggestion(bufnr, ns_id, edits)
+    M.clear_suggestion(bufnr, ns_id)
+    if not edits or #edits == 0 then
+        -- vim.notify("No suggestion available", vim.log.levels.INFO)
+        return
+    end
+
+    local suggestion = edits[1]
+
+    local preview = M._calculate_preview(bufnr, suggestion)
+    M._display_preview(bufnr, ns_id, preview)
 
     vim.b[bufnr].nes_state = suggestion
 

lua/copilot-lsp/nes/ui.lua

@@ -1,31 +1,25 @@
----@class copilotlsp.InlineEdit
+---@class copilotlsp.InlineEdit : lsp.TextEdit
 ---@field command lsp.Command
----@field range lsp.Range
 ---@field text string
----@field newText string
 ---@field textDocument lsp.VersionedTextDocumentIdentifier
 
 ---@class copilotlsp.copilotInlineEditResponse
 ---@field edits copilotlsp.InlineEdit[]
 
----@class copilotlsp.nes.EditSuggestionUI
----@field preview_winnr? integer
+---@class copilotlsp.nes.TextDeletion
+---@field range lsp.Range
 
----@class copilotlsp.nes.DeleteExtmark
---- Holds row information for delete highlight extmark.
----@field row number
----@field end_row number
+---@class copilotlsp.nes.InlineInsertion
+---@field text string
+---@field line integer
+---@field character integer
 
----@class copilotlsp.nes.AddExtmark
--- Holds row and virtual lines count for virtual lines extmark.
----@field row number
----@field virt_lines_count number
+---@class copilotlsp.nes.TextInsertion
+---@field text string
+---@field line integer insert lines at this line
+---@field above? boolean above the line
 
----@class copilotlsp.nes.LineCalculationResult
---- The result of calculating lines for inline suggestion UI.
----@field deleted_lines_count number
----@field added_lines string[]
----@field added_lines_count number
----@field same_line number
----@field delete_extmark copilotlsp.nes.DeleteExtmark
----@field virt_lines_extmark copilotlsp.nes.AddExtmark
+---@class copilotlsp.nes.InlineEditPreview
+---@field deletion? copilotlsp.nes.TextDeletion
+---@field inline_insertion? copilotlsp.nes.InlineInsertion
+---@field lines_insertion? copilotlsp.nes.TextInsertion

lua/copilot-lsp/types.lua

@@ -0,0 +1,3 @@
+vim.api.nvim_set_hl(0, "CopilotLspNesAdd", { link = "DiffAdd", default = true })
+vim.api.nvim_set_hl(0, "CopilotLspNesDelete", { link = "DiffDelete", default = true })
+vim.api.nvim_set_hl(0, "CopilotLspNesApply", { link = "DiffText", default = true })

plugin/copilot-lsp.lua

@@ -74,7 +74,7 @@ local function getNesResponse(td)
                     command = { title = "mock", command = "mock" },
                     range = {
                         start = { line = 4, character = 0 },
-                        ["end"] = { line = 4, character = 30 },
+                        ["end"] = { line = 4, character = 31 },
                     },
                     textDocument = td,
                     text = [[  printf("Goodb, %s!\n", name);]],