feat(lsp): enhance copilot inline suggestions

- Set custom highlight groups for NesAdd, NesDelete, and NesApply.
- Add keymap to apply pending inline suggestions on <leader>xa.
- Comment out extraneous debug output in completion handler.
- Refactor inline edit application using LSP util function.

Co-authored-by: Xuyaun Pang <[email protected]>

Author: tris203

lsp/copilot.lua

@@ -19,6 +19,10 @@ return {
     },
     root_markers = { ".git" },
     on_init = function(client)
+        vim.api.nvim_set_hl(0, "NesAdd", { link = "DiffAdd", default = true })
+        vim.api.nvim_set_hl(0, "NesDelete", { link = "DiffDelete", default = true })
+        vim.api.nvim_set_hl(0, "NesApply", { link = "DiffText", default = true })
+
         local nes = require("copilot-lsp.nes")
         local inline_completion = require("copilot-lsp.completion")
 
@@ -30,6 +34,10 @@ return {
             nes.request_nes(client)
         end)
 
+        vim.keymap.set("n", "<leader>xa", function()
+            nes.apply_pending_nes(nil, { trigger = true, jump = true }, client)
+        end)
+
         local au = vim.api.nvim_create_augroup("copilot-language-server", { clear = true })
         vim.api.nvim_create_autocmd("TextChangedI", {
             callback = function()

lua/copilot-lsp/completion/init.lua

@@ -18,7 +18,7 @@ local function handle_inlineCompletion_response(results, _ctx, _config)
 
     for _, result in pairs(results1) do
         -- This is where we show the completion results
-        dd(result)
+        -- dd(result)
     end
 end
 

lua/copilot-lsp/nes/init.lua

@@ -1,55 +1,23 @@
 local errs = require("copilot-lsp.errors")
+local nes_ui = require("copilot-lsp.nes.ui")
 local utils = require("copilot-lsp.util")
 
 local M = {}
 
-local nes_ext
 local nes_ns = vim.api.nvim_create_namespace("copilot-nes")
 
----@param edit copilotInlineEdit
-local function display_nes(edit)
-    dd("trying to display")
-    local bufnr = vim.uri_to_bufnr(edit.textDocument.uri)
-    if edit.text:match("\n") then
-        assert(false, "multi line edits not supported yet")
-    end
-
-    nes_ext = vim.api.nvim_buf_set_extmark(bufnr, nes_ns, edit.range.start.line, edit.range.start.character, {
-        id = nes_ext,
-        virt_lines = { { { edit.text, "Comment" } } },
-    })
-
-    --create accept and decline keymaps
-    vim.keymap.set("n", "<leader>xa", function()
-        utils.apply_inline_edit(edit)
-        vim.api.nvim_buf_del_extmark(bufnr, nes_ns, nes_ext)
-    end, { buffer = bufnr })
-
-    vim.keymap.set("n", "<leader>xd", function()
-        vim.api.nvim_buf_del_extmark(bufnr, nes_ns, nes_ext)
-    end, { buffer = bufnr })
-end
-
 ---@param err lsp.ResponseError?
----@param result copilotInlineEditResponse
+---@param result copilotlsp.copilotInlineEditResponse
 local function handle_nes_response(err, result)
     if err then
         vim.notify(err.message)
         return
     end
-    if #result.edits > 1 then
-        vim.notify("more than 1 edit, dont know what to do yet")
-        return
-    end
-
-    if #result.edits == 0 then
-        vim.notify("no edits")
-        return
+    for _, edit in ipairs(result.edits) do
+        --- Convert to textEdit fields
+        edit.newText = edit.text
     end
-
-    local edit = result.edits[1]
-    display_nes(edit)
-    -- utils.apply_inline_edit(edit)
+    nes_ui._display_next_suggestion(result.edits, nes_ns)
 end
 
 ---@param copilot_lss vim.lsp.Client?
@@ -62,4 +30,26 @@ function M.request_nes(copilot_lss)
     copilot_lss:request("textDocument/copilotInlineEdit", pos_params, handle_nes_response)
 end
 
+---@param bufnr? integer
+---@param opts? nes.Apply.Opts
+---@param client vim.lsp.Client
+function M.apply_pending_nes(bufnr, opts, client)
+    opts = opts or {}
+
+    bufnr = bufnr and bufnr > 0 and bufnr or vim.api.nvim_get_current_buf()
+
+    ---@type copilotlsp.InlineEdit
+    local state = vim.b[bufnr].nes_state
+    if not state then
+        return
+    end
+    utils.apply_inline_edit(state)
+    nes_ui.clear_suggestion(bufnr, nes_ns)
+    if opts.trigger then
+        vim.schedule(function()
+            M.request_nes(client)
+        end)
+    end
+end
+
 return M

lua/copilot-lsp/nes/ui.lua

@@ -0,0 +1,122 @@
+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)
+    pcall(vim.api.nvim_buf_clear_namespace, bufnr, ns_id, 0, -1)
+end
+
+---@param bufnr? integer
+---@param ns_id integer
+function M.clear_suggestion(bufnr, ns_id)
+    bufnr = bufnr and bufnr > 0 and bufnr or vim.api.nvim_get_current_buf()
+    vim.api.nvim_buf_clear_namespace(bufnr, ns_id, 0, -1)
+    ---@type copilotlsp.InlineEdit
+    local state = vim.b[bufnr].nes_state
+    if not state then
+        return
+    end
+
+    _dismiss_suggestion_ui(bufnr, state.ui, ns_id)
+    vim.b[bufnr].nes_state = nil
+end
+
+---@private
+---@param edits copilotlsp.InlineEdit[]
+---@param ns_id integer
+function M._display_next_suggestion(edits, ns_id)
+    if not edits or #edits == 0 then
+        vim.notify("No suggestion available", vim.log.levels.INFO)
+        return
+    end
+    local bufnr = vim.uri_to_bufnr(edits[1].textDocument.uri)
+    local win_id = vim.fn.win_findbuf(bufnr)[1]
+    local suggestion = edits[1]
+
+    local ui = {}
+    local deleted_lines_count = suggestion.range["end"].line - suggestion.range.start.line
+    local added_lines = vim.split(suggestion.newText, "\n")
+    local added_lines_count = suggestion.newText == "" and 0 or #added_lines - 1
+    local same_line = 0
+
+    if deleted_lines_count == 0 and added_lines_count == 0 then
+        ---changing within line
+        deleted_lines_count = 1
+        added_lines_count = 1
+        same_line = 1
+        added_lines = { suggestion.newText }
+    end
+
+    if deleted_lines_count > 0 then
+        vim.api.nvim_buf_set_extmark(bufnr, ns_id, suggestion.range.start.line, 0, {
+            hl_group = "NesDelete",
+            end_line = suggestion.range["end"].line,
+        })
+    end
+    if added_lines_count > 0 then
+        local virt_lines = {}
+        for _ = 1, added_lines_count do
+            table.insert(virt_lines, {
+                { "", "Normal" },
+            })
+        end
+        local line = suggestion.range.start.line + deleted_lines_count - 1 + same_line
+
+        vim.api.nvim_buf_set_extmark(bufnr, ns_id, line, 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, 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 = (#added_lines + same_line) - 1,
+            row = (suggestion.range["end"].line + same_line) - cursor[1] + 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" }, {
+        buffer = bufnr,
+        callback = function()
+            if not vim.b.nes_state then
+                return true
+            end
+
+            local accepted_cursor = vim.b.nes_state.accepted_cursor
+            if accepted_cursor then
+                local cursor = vim.api.nvim_win_get_cursor(win_id)
+                if cursor[1] == accepted_cursor[1] and cursor[2] == accepted_cursor[2] then
+                    return
+                end
+            end
+
+            M.clear_suggestion(bufnr, ns_id)
+            return true
+        end,
+    })
+end
+
+return M

lua/copilot-lsp/types.lua

@@ -1,8 +1,17 @@
----@class copilotInlineEdit
+---@class copilotlsp.InlineEdit
 ---@field command lsp.Command
 ---@field range lsp.Range
 ---@field text string
+---@field newText string
 ---@field textDocument lsp.VersionedTextDocumentIdentifier
+---@field ui nes.EditSuggestionUI?
 
----@class copilotInlineEditResponse
----@field edits copilotInlineEdit[]
+---@class copilotlsp.copilotInlineEditResponse
+---@field edits copilotlsp.InlineEdit[]
+
+---@class nes.EditSuggestionUI
+---@field preview_winnr? integer
+
+---@class nes.Apply.Opts
+---@field jump? boolean | { hl_timeout: integer? } auto jump to the end of the new edit
+---@field trigger? boolean auto trigger the next edit suggestion

lua/copilot-lsp/util.lua

@@ -1,20 +1,10 @@
 local M = {}
----@param edit copilotInlineEdit
+---@param edit copilotlsp.InlineEdit
 function M.apply_inline_edit(edit)
     local bufnr = vim.uri_to_bufnr(edit.textDocument.uri)
-    local multi_line
-    if edit.text:match("\n") then
-        multi_line = vim.split(edit.text, "\n")
-    end
 
-    vim.api.nvim_buf_set_text(
-        bufnr,
-        edit.range.start.line,
-        edit.range.start.character,
-        edit.range["end"].line,
-        edit.range["end"].character,
-        multi_line or { edit.text }
-    )
+    ---@diagnostic disable-next-line: assign-type-mismatch
+    vim.lsp.util.apply_text_edits({ edit }, bufnr, "utf-8")
 end
 
 return M