fix(lsp): restore marks after apply_text_edits() #14630

PROBLEM:
Whenever any text edits are applied to the buffer, the `marks` part of those
lines will be lost. This is mostly problematic for code formatters that format
the whole buffer like `prettier`, `luafmt`, ...

When doing atomic changes inside a vim doc, vim keeps track of those changes and
can update the positions of marks accordingly, but in this case we have a whole
doc that changed. There's no simple way to update the positions of all marks
from the previous document state to the new document state.

SOLUTION:
* save marks right before `nvim_buf_set_lines` is called inside `apply_text_edits`
* check if any marks were lost after doing `nvim_buf_set_lines`
* restore those marks to the previous positions

TEST CASE:
* have a formatter enabled
* open any file
* create a couple of marks
* indent the whole file to the right
* save the file
Before this change: all marks will be removed.
After this change: they will be preserved.

Fixes #14307

3 files changed, 73 insertions, 0 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index f8757af2be..e42be43461 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -41,6 +41,9 @@ ADDED FEATURES                                                     *news-added*
 
 The following new APIs or features were added.
 
+• Neovim's LSP client now always saves and restores named buffer marks when
+  applying text edits.
+
 • Nvim's LSP client now advertises the general.positionEncodings client
   capability to indicate to servers that it supports utf-8, utf-16, and utf-32
   encodings. If the server responds with the positionEncoding capability in
diff --git a/runtime/lua/vim/lsp/util.lua b/runtime/lua/vim/lsp/util.lua
index 53f8dba814..ba8c72128e 100644
--- a/runtime/lua/vim/lsp/util.lua
+++ b/runtime/lua/vim/lsp/util.lua
@@ -451,6 +451,14 @@ function M.apply_text_edits(text_edits, bufnr, offset_encoding)
     }
   end)()
 
+  -- save and restore local marks since they get deleted by nvim_buf_set_lines
+  local marks = {}
+  for _, m in pairs(vim.fn.getmarklist(bufnr or vim.api.nvim_get_current_buf())) do
+    if m.mark:match("^'[a-z]$") then
+      marks[m.mark:sub(2, 2)] = { m.pos[2], m.pos[3] - 1 } -- api-indexed
+    end
+  end
+
   -- Apply text edits.
   local is_cursor_fixed = false
   local has_eol_text_edit = false
@@ -518,6 +526,20 @@ function M.apply_text_edits(text_edits, bufnr, offset_encoding)
 
   local max = api.nvim_buf_line_count(bufnr)
 
+  -- no need to restore marks that still exist
+  for _, m in pairs(vim.fn.getmarklist(bufnr or vim.api.nvim_get_current_buf())) do
+    marks[m.mark:sub(2, 2)] = nil
+  end
+  -- restore marks
+  for mark, pos in pairs(marks) do
+    if pos then
+      -- make sure we don't go out of bounds
+      pos[1] = math.min(pos[1], max)
+      pos[2] = math.min(pos[2], #(get_line(bufnr, pos[1] - 1) or ''))
+      vim.api.nvim_buf_set_mark(bufnr or 0, mark, pos[1], pos[2], {})
+    end
+  end
+
   -- Apply fixed cursor position.
   if is_cursor_fixed then
     local is_valid_cursor = true
diff --git a/test/functional/plugin/lsp_spec.lua b/test/functional/plugin/lsp_spec.lua
index 26fd2276fd..72a9bc45a2 100644
--- a/test/functional/plugin/lsp_spec.lua
+++ b/test/functional/plugin/lsp_spec.lua
@@ -1679,6 +1679,54 @@ describe('LSP', function()
         'foobar';
       }, buf_lines(1))
     end)
+    it('it restores marks', function()
+      local edits = {
+        make_edit(1, 0, 2, 5, "foobar");
+        make_edit(4, 0, 5, 0, "barfoo");
+      }
+      eq(true, exec_lua('return vim.api.nvim_buf_set_mark(1, "a", 2, 1, {})'))
+      exec_lua('vim.lsp.util.apply_text_edits(...)', edits, 1, "utf-16")
+      eq({
+        'First line of text';
+        'foobar line of text';
+        'Fourth line of text';
+        'barfoo';
+      }, buf_lines(1))
+      local mark = exec_lua('return vim.api.nvim_buf_get_mark(1, "a")')
+      eq({ 2, 1 }, mark)
+    end)
+
+    it('it restores marks to last valid col', function()
+      local edits = {
+        make_edit(1, 0, 2, 15, "foobar");
+        make_edit(4, 0, 5, 0, "barfoo");
+      }
+      eq(true, exec_lua('return vim.api.nvim_buf_set_mark(1, "a", 2, 10, {})'))
+      exec_lua('vim.lsp.util.apply_text_edits(...)', edits, 1, "utf-16")
+      eq({
+        'First line of text';
+        'foobarext';
+        'Fourth line of text';
+        'barfoo';
+      }, buf_lines(1))
+      local mark = exec_lua('return vim.api.nvim_buf_get_mark(1, "a")')
+      eq({ 2, 9 }, mark)
+    end)
+
+    it('it restores marks to last valid line', function()
+      local edits = {
+        make_edit(1, 0, 4, 5, "foobar");
+        make_edit(4, 0, 5, 0, "barfoo");
+      }
+      eq(true, exec_lua('return vim.api.nvim_buf_set_mark(1, "a", 4, 1, {})'))
+      exec_lua('vim.lsp.util.apply_text_edits(...)', edits, 1, "utf-16")
+      eq({
+        'First line of text';
+        'foobaro';
+      }, buf_lines(1))
+      local mark = exec_lua('return vim.api.nvim_buf_get_mark(1, "a")')
+      eq({ 2, 1 }, mark)
+    end)
 
     describe('cursor position', function()
       it('don\'t fix the cursor if the range contains the cursor', function()