feat(lsp): improve vim.lsp.util.apply_text_edits (#15561)

- Fix the cursor position after applying TextEdits
- Support reversed range of TextEdit
- Invoke nvim_buf_set_text one by one

2 files changed, 206 insertions, 40 deletions

diff --git a/runtime/lua/vim/lsp/util.lua b/runtime/lua/vim/lsp/util.lua
index 9a3ce185a0..e95f170427 100644
--- a/runtime/lua/vim/lsp/util.lua
+++ b/runtime/lua/vim/lsp/util.lua
@@ -146,10 +146,6 @@ local function sort_by_key(fn)
     return false
   end
 end
----@private
-local edit_sort_key = sort_by_key(function(e)
-  return {e.A[1], e.A[2], e.i}
-end)
 
 ---@private
 --- Position is a https://microsoft.github.io/language-server-protocol/specifications/specification-current/#position
@@ -174,6 +170,7 @@ local function get_line_byte_from_position(bufnr, position)
       if ok then
         return result
       end
+      return math.min(#lines[1], col)
     end
   end
   return col
@@ -237,8 +234,8 @@ function M.get_progress_messages()
 end
 
 --- Applies a list of text edits to a buffer.
----@param text_edits (table) list of `TextEdit` objects
----@param buf_nr (number) Buffer id
+---@param text_edits table list of `TextEdit` objects
+---@param bufnr number Buffer id
 ---@see https://microsoft.github.io/language-server-protocol/specifications/specification-current/#textEdit
 function M.apply_text_edits(text_edits, bufnr)
   if not next(text_edits) then return end
@@ -246,45 +243,110 @@ function M.apply_text_edits(text_edits, bufnr)
     vim.fn.bufload(bufnr)
   end
   api.nvim_buf_set_option(bufnr, 'buflisted', true)
-  local start_line, finish_line = math.huge, -1
-  local cleaned = {}
-  for i, e in ipairs(text_edits) do
-    -- adjust start and end column for UTF-16 encoding of non-ASCII characters
-    local start_row = e.range.start.line
-    local start_col = get_line_byte_from_position(bufnr, e.range.start)
-    local end_row = e.range["end"].line
-    local end_col = get_line_byte_from_position(bufnr, e.range['end'])
-    start_line = math.min(e.range.start.line, start_line)
-    finish_line = math.max(e.range["end"].line, finish_line)
-    -- TODO(ashkan) sanity check ranges for overlap.
-    table.insert(cleaned, {
-      i = i;
-      A = {start_row; start_col};
-      B = {end_row; end_col};
-      lines = vim.split(e.newText, '\n', true);
-    })
-  end
 
-  -- Reverse sort the orders so we can apply them without interfering with
-  -- eachother. Also add i as a sort key to mimic a stable sort.
-  table.sort(cleaned, edit_sort_key)
-  local lines = api.nvim_buf_get_lines(bufnr, start_line, finish_line + 1, false)
-  local fix_eol = api.nvim_buf_get_option(bufnr, 'fixeol')
-  local set_eol = fix_eol and api.nvim_buf_line_count(bufnr) <= finish_line + 1
-  if set_eol and (#lines == 0 or #lines[#lines] ~= 0) then
-    table.insert(lines, '')
+  -- Fix reversed range and indexing each text_edits
+  local index = 0
+  text_edits = vim.tbl_map(function(text_edit)
+    index = index + 1
+    text_edit._index = index
+
+    if text_edit.range.start.line > text_edit.range['end'].line or text_edit.range.start.line == text_edit.range['end'].line and text_edit.range.start.character > text_edit.range['end'].character then
+      local start = text_edit.range.start
+      text_edit.range.start = text_edit.range['end']
+      text_edit.range['end'] = start
+    end
+    return text_edit
+  end, text_edits)
+
+  -- Sort text_edits
+  table.sort(text_edits, function(a, b)
+    if a.range.start.line ~= b.range.start.line then
+      return a.range.start.line > b.range.start.line
+    end
+    if a.range.start.character ~= b.range.start.character then
+      return a.range.start.character > b.range.start.character
+    end
+    if a._index ~= b._index then
+      return a._index > b._index
+    end
+  end)
+
+  -- Some LSP servers may return +1 range of the buffer content but nvim_buf_set_text can't accept it so we should fix it here.
+  local has_eol_text_edit = false
+  local max = vim.api.nvim_buf_line_count(bufnr)
+  local len = vim.str_utfindex(vim.api.nvim_buf_get_lines(bufnr, -2, -1, false)[1] or '')
+  text_edits = vim.tbl_map(function(text_edit)
+    if max <= text_edit.range.start.line then
+      text_edit.range.start.line = max - 1
+      text_edit.range.start.character = len
+      text_edit.newText = '\n' .. text_edit.newText
+      has_eol_text_edit = true
+    end
+    if max <= text_edit.range['end'].line then
+      text_edit.range['end'].line = max - 1
+      text_edit.range['end'].character = len
+      has_eol_text_edit = true
+    end
+    return text_edit
+  end, text_edits)
+
+  -- Some LSP servers are depending on the VSCode behavior.
+  -- The VSCode will re-locate the cursor position after applying TextEdit so we also do it.
+  local is_current_buf = vim.api.nvim_get_current_buf() == bufnr
+  local cursor = (function()
+    if not is_current_buf then
+      return {
+        row = -1,
+        col = -1,
+      }
+    end
+    local cursor = vim.api.nvim_win_get_cursor(0)
+    return {
+      row = cursor[1] - 1,
+      col = cursor[2],
+    }
+  end)()
+
+  -- Apply text edits.
+  local is_cursor_fixed = false
+  for _, text_edit in ipairs(text_edits) do
+    local e = {
+      start_row = text_edit.range.start.line,
+      start_col = get_line_byte_from_position(bufnr, text_edit.range.start),
+      end_row = text_edit.range['end'].line,
+      end_col  = get_line_byte_from_position(bufnr, text_edit.range['end']),
+      text = vim.split(text_edit.newText, '\n', true),
+    }
+    vim.api.nvim_buf_set_text(bufnr, e.start_row, e.start_col, e.end_row, e.end_col, e.text)
+
+    local row_count = (e.end_row - e.start_row) + 1
+    if e.end_row < cursor.row then
+      cursor.row = cursor.row + (#e.text - row_count)
+      is_cursor_fixed = true
+    elseif e.end_row == cursor.row and e.end_col <= cursor.col then
+      cursor.row = cursor.row + (#e.text - row_count)
+      cursor.col = #e.text[#e.text] + (cursor.col - e.end_col)
+      if #e.text == 1 then
+        cursor.col = cursor.col + e.start_col
+      end
+      is_cursor_fixed = true
+    end
   end
 
-  for i = #cleaned, 1, -1 do
-    local e = cleaned[i]
-    local A = {e.A[1] - start_line, e.A[2]}
-    local B = {e.B[1] - start_line, e.B[2]}
-    lines = M.set_lines(lines, A, B, e.lines)
+  if is_cursor_fixed then
+    vim.api.nvim_win_set_cursor(0, {
+      cursor.row + 1,
+      math.min(cursor.col, #(vim.api.nvim_buf_get_lines(bufnr, cursor.row, cursor.row + 1, false)[1] or ''))
+    })
   end
-  if set_eol and #lines[#lines] == 0 then
-    table.remove(lines)
+
+  -- Remove final line if needed
+  local fix_eol = has_eol_text_edit
+  fix_eol = fix_eol and api.nvim_buf_get_option(bufnr, 'fixeol')
+  fix_eol = fix_eol and (vim.api.nvim_buf_get_lines(bufnr, -2, -1, false)[1] or '') == ''
+  if fix_eol then
+    vim.api.nvim_buf_set_lines(bufnr, -2, -1, false, {})
   end
-  api.nvim_buf_set_lines(bufnr, start_line, finish_line + 1, false, lines)
 end
 
 -- local valid_windows_path_characters = "[^<>:\"/\\|?*]"
diff --git a/test/functional/plugin/lsp_spec.lua b/test/functional/plugin/lsp_spec.lua
index ef78c8db4d..6ad37110c7 100644
--- a/test/functional/plugin/lsp_spec.lua
+++ b/test/functional/plugin/lsp_spec.lua
@@ -1066,6 +1066,30 @@ describe('LSP', function()
         'å å ɧ 汉语 ↥ 🤦 🦄';
       }, buf_lines(1))
     end)
+    it('applies complex edits (reversed range)', function()
+      local edits = {
+        make_edit(0, 0, 0, 0, {"", "12"});
+        make_edit(0, 0, 0, 0, {"3", "foo"});
+        make_edit(0, 1, 0, 1, {"bar", "123"});
+        make_edit(0, #"First line of text", 0, #"First ", {"guy"});
+        make_edit(1, #'Second', 1, 0, {"baz"});
+        make_edit(2, #"Third", 2, #'Th', {"e next"});
+        make_edit(3, #"Fourth", 3, #'', {"another line of text", "before this"});
+        make_edit(3, #"Fourth line of text", 3, #'Fourth', {"!"});
+      }
+      exec_lua('vim.lsp.util.apply_text_edits(...)', edits, 1)
+      eq({
+        '';
+        '123';
+        'fooFbar';
+        '123irst guy';
+        'baz line of text';
+        'The next line of text';
+        'another line of text';
+        'before this!';
+        'å å ɧ 汉语 ↥ 🤦 🦄';
+      }, buf_lines(1))
+    end)
     it('applies non-ASCII characters edits', function()
       local edits = {
         make_edit(4, 3, 4, 4, {"ä"});
@@ -1094,6 +1118,86 @@ describe('LSP', function()
       }, buf_lines(1))
     end)
 
+    describe('cursor position', function()
+      it('don\'t fix the cursor if the range contains the cursor', function()
+        funcs.nvim_win_set_cursor(0, { 2, 6 })
+        local edits = {
+          make_edit(1, 0, 1, 19, 'Second line of text')
+        }
+        exec_lua('vim.lsp.util.apply_text_edits(...)', edits, 1)
+        eq({
+          'First line of text';
+          'Second line of text';
+          'Third line of text';
+          'Fourth line of text';
+          'å å ɧ 汉语 ↥ 🤦 🦄';
+        }, buf_lines(1))
+        eq({ 2, 6 }, funcs.nvim_win_get_cursor(0))
+      end)
+
+      it('fix the cursor to the valid column if the content was removed', function()
+        funcs.nvim_win_set_cursor(0, { 2, 6 })
+        local edits = {
+          make_edit(1, 0, 1, 19, '')
+        }
+        exec_lua('vim.lsp.util.apply_text_edits(...)', edits, 1)
+        eq({
+          'First line of text';
+          '';
+          'Third line of text';
+          'Fourth line of text';
+          'å å ɧ 汉语 ↥ 🤦 🦄';
+        }, buf_lines(1))
+        eq({ 2, 0 }, funcs.nvim_win_get_cursor(0))
+      end)
+
+      it('fix the cursor row', function()
+        funcs.nvim_win_set_cursor(0, { 3, 0 })
+        local edits = {
+          make_edit(1, 0, 2, 0, '')
+        }
+        exec_lua('vim.lsp.util.apply_text_edits(...)', edits, 1)
+        eq({
+          'First line of text';
+          'Third line of text';
+          'Fourth line of text';
+          'å å ɧ 汉语 ↥ 🤦 🦄';
+        }, buf_lines(1))
+        eq({ 2, 0 }, funcs.nvim_win_get_cursor(0))
+      end)
+
+      it('fix the cursor col', function()
+        funcs.nvim_win_set_cursor(0, { 2, 11 })
+        local edits = {
+          make_edit(1, 7, 1, 11, '')
+        }
+        exec_lua('vim.lsp.util.apply_text_edits(...)', edits, 1)
+        eq({
+          'First line of text';
+          'Second  of text';
+          'Third line of text';
+          'Fourth line of text';
+          'å å ɧ 汉语 ↥ 🤦 🦄';
+        }, buf_lines(1))
+        eq({ 2, 7 }, funcs.nvim_win_get_cursor(0))
+      end)
+
+      it('fix the cursor row and col', function()
+        funcs.nvim_win_set_cursor(0, { 2, 12 })
+        local edits = {
+          make_edit(0, 11, 1, 12, '')
+        }
+        exec_lua('vim.lsp.util.apply_text_edits(...)', edits, 1)
+        eq({
+          'First line of text';
+          'Third line of text';
+          'Fourth line of text';
+          'å å ɧ 汉语 ↥ 🤦 🦄';
+        }, buf_lines(1))
+        eq({ 1, 11 }, funcs.nvim_win_get_cursor(0))
+      end)
+    end)
+
     describe('with LSP end line after what Vim considers to be the end line', function()
       it('applies edits when the last linebreak is considered a new line', function()
         local edits = {