fix(extmark): restore extmarks when completing original text

3 files changed, 83 insertions, 9 deletions

diff --git a/src/nvim/extmark.c b/src/nvim/extmark.c
index d9c1993f32..f510845ec7 100644
--- a/src/nvim/extmark.c
+++ b/src/nvim/extmark.c
@@ -311,9 +311,8 @@ void extmark_free_all(buf_T *buf)
 /// copying is useful when we cannot simply reverse the operation. This will do
 /// nothing on redo, enforces correct position when undo.
 void extmark_splice_delete(buf_T *buf, int l_row, colnr_T l_col, int u_row, colnr_T u_col,
-                           ExtmarkOp op)
+                           extmark_undo_vec_t *uvp, bool only_copy, ExtmarkOp op)
 {
-  u_header_T *uhp = u_force_get_undo_header(buf);
   MarkTreeIter itr[1] = { 0 };
   ExtmarkUndoObject undo;
 
@@ -328,7 +327,7 @@ void extmark_splice_delete(buf_T *buf, int l_row, colnr_T l_col, int u_row, coln
 
     bool invalidated = false;
     // Invalidate/delete mark
-    if (!mt_invalid(mark) && mt_invalidate(mark) && !mt_end(mark)) {
+    if (!only_copy && !mt_invalid(mark) && mt_invalidate(mark) && !mt_end(mark)) {
       MTPos endpos = marktree_get_altpos(buf->b_marktree, mark, NULL);
       if (endpos.row < 0) {
         endpos = mark.pos;
@@ -348,7 +347,7 @@ void extmark_splice_delete(buf_T *buf, int l_row, colnr_T l_col, int u_row, coln
     }
 
     // Push mark to undo header
-    if (uhp && op == kExtmarkUndo && !mt_no_undo(mark)) {
+    if (only_copy || (uvp != NULL && op == kExtmarkUndo && !mt_no_undo(mark))) {
       ExtmarkSavePos pos;
       pos.mark = mt_lookup_key(mark);
       pos.invalidated = invalidated;
@@ -359,7 +358,7 @@ void extmark_splice_delete(buf_T *buf, int l_row, colnr_T l_col, int u_row, coln
 
       undo.data.savepos = pos;
       undo.type = kExtmarkSavePos;
-      kv_push(uhp->uh_extmark, undo);
+      kv_push(*uvp, undo);
     }
 
     marktree_itr_next(buf->b_marktree, itr);
@@ -511,7 +510,9 @@ void extmark_splice_impl(buf_T *buf, int start_row, colnr_T start_col, bcount_t
     // merge!)
     int end_row = start_row + old_row;
     int end_col = (old_row ? 0 : start_col) + old_col;
-    extmark_splice_delete(buf, start_row, start_col, end_row, end_col, undo);
+    u_header_T *uhp = u_force_get_undo_header(buf);
+    extmark_undo_vec_t *uvp = uhp ? &uhp->uh_extmark : NULL;
+    extmark_splice_delete(buf, start_row, start_col, end_row, end_col, uvp, false, undo);
   }
 
   // Move the signcolumn sentinel line
diff --git a/src/nvim/insexpand.c b/src/nvim/insexpand.c
index 12543a2d42..728277e5a3 100644
--- a/src/nvim/insexpand.c
+++ b/src/nvim/insexpand.c
@@ -23,6 +23,7 @@
 #include "nvim/eval/userfunc.h"
 #include "nvim/ex_eval.h"
 #include "nvim/ex_getln.h"
+#include "nvim/extmark.h"
 #include "nvim/fileio.h"
 #include "nvim/func_attr.h"
 #include "nvim/garray.h"
@@ -252,6 +253,8 @@ static colnr_T compl_col = 0;           ///< column where the text starts
                                         ///< that is being completed
 static char *compl_orig_text = NULL;    ///< text as it was before
                                         ///< completion started
+/// Undo information to restore extmarks for original text.
+static extmark_undo_vec_t compl_orig_extmarks;
 static int compl_cont_mode = 0;
 static expand_T compl_xp;
 
@@ -1569,6 +1572,7 @@ void ins_compl_clear(void)
   XFREE_CLEAR(compl_pattern);
   XFREE_CLEAR(compl_leader);
   edit_submode_extra = NULL;
+  kv_destroy(compl_orig_extmarks);
   XFREE_CLEAR(compl_orig_text);
   compl_enter_selects = false;
   // clear v:completed_item
@@ -2019,6 +2023,7 @@ static bool ins_compl_stop(const int c, const int prev_mode, bool retval)
         ins_bytes_len(p + compl_len, (size_t)(len - compl_len));
       }
     }
+    restore_orig_extmarks();
     retval = true;
   }
 
@@ -2505,6 +2510,22 @@ static void ins_compl_add_dict(dict_T *dict)
   }
 }
 
+/// Save extmarks in "compl_orig_text" so that they may be restored when the
+/// completion is cancelled, or the original text is completed.
+static void save_orig_extmarks(void)
+{
+  extmark_splice_delete(curbuf, curwin->w_cursor.lnum - 1, compl_col, curwin->w_cursor.lnum - 1,
+                        compl_col + compl_length, &compl_orig_extmarks, true, kExtmarkUndo);
+}
+
+static void restore_orig_extmarks(void)
+{
+  for (long i = (int)kv_size(compl_orig_extmarks) - 1; i > -1; i--) {
+    ExtmarkUndoObject undo_info = kv_A(compl_orig_extmarks, i);
+    extmark_apply_undo(undo_info, true);
+  }
+}
+
 /// Start completion for the complete() function.
 ///
 /// @param startcol  where the matched text starts (1 is first column).
@@ -2526,10 +2547,10 @@ static void set_completion(colnr_T startcol, list_T *list)
     startcol = curwin->w_cursor.col;
   }
   compl_col = startcol;
-  compl_length = (int)curwin->w_cursor.col - (int)startcol;
+  compl_length = curwin->w_cursor.col - startcol;
   // compl_pattern doesn't need to be set
-  compl_orig_text = xstrnsave(get_cursor_line_ptr() + compl_col,
-                              (size_t)compl_length);
+  compl_orig_text = xstrnsave(get_cursor_line_ptr() + compl_col, (size_t)compl_length);
+  save_orig_extmarks();
   if (p_ic) {
     flags |= CP_ICASE;
   }
@@ -3689,12 +3710,16 @@ static int ins_compl_next(bool allow_get_expansion, int count, bool insert_match
   if (compl_no_insert && !started) {
     ins_bytes(compl_orig_text + get_compl_len());
     compl_used_match = false;
+    restore_orig_extmarks();
   } else if (insert_match) {
     if (!compl_get_longest || compl_used_match) {
       ins_compl_insert(in_compl_func);
     } else {
       ins_bytes(compl_leader + get_compl_len());
     }
+    if (!strcmp(compl_curr_match->cp_str, compl_orig_text)) {
+      restore_orig_extmarks();
+    }
   } else {
     compl_used_match = false;
   }
@@ -4267,7 +4292,9 @@ static int ins_compl_start(void)
 
   // Always add completion for the original text.
   xfree(compl_orig_text);
+  kv_destroy(compl_orig_extmarks);
   compl_orig_text = xstrnsave(line + compl_col, (size_t)compl_length);
+  save_orig_extmarks();
   int flags = CP_ORIGINAL_TEXT;
   if (p_ic) {
     flags |= CP_ICASE;
@@ -4276,6 +4303,7 @@ static int ins_compl_start(void)
                     flags, false) != OK) {
     XFREE_CLEAR(compl_pattern);
     XFREE_CLEAR(compl_orig_text);
+    kv_destroy(compl_orig_extmarks);
     return FAIL;
   }
 
@@ -4508,6 +4536,7 @@ static unsigned quote_meta(char *dest, char *src, int len)
 void free_insexpand_stuff(void)
 {
   XFREE_CLEAR(compl_orig_text);
+  kv_destroy(compl_orig_extmarks);
   callback_free(&cfu_cb);
   callback_free(&ofu_cb);
   callback_free(&tsrfu_cb);
diff --git a/test/functional/editor/completion_spec.lua b/test/functional/editor/completion_spec.lua
index cbaf401f06..84af90a298 100644
--- a/test/functional/editor/completion_spec.lua
+++ b/test/functional/editor/completion_spec.lua
@@ -1257,4 +1257,48 @@ describe('completion', function()
       {3:-- }{4:match 1 of 2}     |
     ]]}
   end)
+
+  it('restores extmarks if original text is restored #23653', function()
+    screen:try_resize(screen._width, 4)
+    command([[
+      call setline(1, ['aaaa'])
+      let ns_id = nvim_create_namespace('extmark')
+      let mark_id = nvim_buf_set_extmark(0, ns_id, 0, 0, { 'end_col':2, 'hl_group':'Error'})
+      let mark = nvim_buf_get_extmark_by_id(0, ns_id, mark_id, { 'details':1 })
+      inoremap <C-x> <C-r>=Complete()<CR>
+      function Complete() abort
+        call complete(1, [{ 'word': 'aaaaa' }])
+        return ''
+      endfunction
+    ]])
+    feed('A<C-X><C-E><Esc>')
+    eq(eval('mark'), eval("nvim_buf_get_extmark_by_id(0, ns_id, mark_id, { 'details':1 })"))
+    feed('A<C-N>')
+    eq(eval('mark'), eval("nvim_buf_get_extmark_by_id(0, ns_id, mark_id, { 'details':1 })"))
+    feed('<Esc>0Yppia<Esc>ggI<C-N>')
+    screen:expect([[
+      aaaa{7:^aa}aa                                                    |
+      {2:aaaa           }                                             |
+      {1:aaaaa          }                                             |
+      {3:-- Keyword completion (^N^P) }{4:match 1 of 2}                   |
+    ]])
+    feed('<C-N><C-N><Esc>')
+    eq(eval('mark'), eval("nvim_buf_get_extmark_by_id(0, ns_id, mark_id, { 'details':1 })"))
+    feed('A<C-N>')
+    eq(eval('mark'), eval("nvim_buf_get_extmark_by_id(0, ns_id, mark_id, { 'details':1 })"))
+    feed('<C-N>')
+    screen:expect([[
+      aaaaa^                                                       |
+      {1:aaaa           }                                             |
+      {2:aaaaa          }                                             |
+      {3:-- Keyword completion (^N^P) }{4:match 2 of 2}                   |
+    ]])
+    feed('<C-E>')
+    screen:expect([[
+      {7:aa}aa^                                                        |
+      aaaa                                                        |
+      aaaaa                                                       |
+      {3:-- INSERT --}                                                |
+    ]])
+  end)
 end)