feat(extmarks): add "undo_restore" flag to opt out of undo-restoring

It is a design goal of extmarks that they allow precise tracking
of changes across undo/redo, including restore the exact positions
after a do/undo or undo/redo cycle. However this behavior is not useful
for all usecases. Many plugins won't keep marks around for long after
text changes, but uses them more like a cache until some external source
(like LSP semantic highlights) has fully updated to changed text and
then will explicitly readjust/replace extmarks as needed.

Add a "undo_restore" flag which is true by default (matches existing
behavior) but can be set to false to opt-out of this behavior.

Delete dead u_extmark_set() code.

10 files changed, 101 insertions, 54 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index b88b7d164f..e32e1aadb6 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -270,6 +270,9 @@ The following changes to existing APIs or features add new behavior.
   In addition, |nvim_buf_get_extmarks()| has gained an "overlap" option to
   return such ranges even if they started before the specified position.
 
+• Extmarks can opt-out of precise undo tracking using the new "undo_restore"
+  flag to |nvim_buf_set_extmark()|
+
 • LSP hover and signature help now use Treesitter for highlighting of Markdown
   content.
   Note that syntax highlighting of code examples requires a matching parser
diff --git a/src/nvim/api/deprecated.c b/src/nvim/api/deprecated.c
index 59b7fc18d6..906edb7b44 100644
--- a/src/nvim/api/deprecated.c
+++ b/src/nvim/api/deprecated.c
@@ -172,8 +172,7 @@ Integer nvim_buf_set_virtual_text(Buffer buffer, Integer src_id, Integer line, A
   decor.virt_text_width = width;
   decor.priority = 0;
 
-  extmark_set(buf, ns_id, NULL, (int)line, 0, -1, -1, &decor, true,
-              false, kExtmarkNoUndo, NULL);
+  extmark_set(buf, ns_id, NULL, (int)line, 0, -1, -1, &decor, true, false, false, NULL);
   return src_id;
 }
 
diff --git a/src/nvim/api/extmark.c b/src/nvim/api/extmark.c
index 0b372c57e5..3f1745df40 100644
--- a/src/nvim/api/extmark.c
+++ b/src/nvim/api/extmark.c
@@ -166,6 +166,10 @@ static Array extmark_to_array(const ExtmarkInfo *extmark, bool id, bool add_dict
       PUT(dict, "end_right_gravity", BOOLEAN_OBJ(extmark->end_right_gravity));
     }
 
+    if (extmark->no_undo) {
+      PUT(dict, "undo_restore", BOOLEAN_OBJ(false));
+    }
+
     const Decoration *decor = &extmark->decor;
     if (decor->hl_id) {
       PUT(dict, "hl_group", hl_group_name(decor->hl_id, hl_name));
@@ -519,6 +523,9 @@ Array nvim_buf_get_extmarks(Buffer buffer, Integer ns_id, Object start, Object e
 ///                   the extmark end position (if it exists) will be shifted
 ///                   in when new text is inserted (true for right, false
 ///                   for left). Defaults to false.
+///               - undo_restore : Restore the exact position of the mark
+///                   if text around the mark was deleted and then restored by undo.
+///                   Defaults to true.
 ///               - priority: a priority value for the highlight group or sign
 ///                   attribute. For example treesitter highlighting uses a
 ///                   value of 100.
@@ -825,7 +832,7 @@ Integer nvim_buf_set_extmark(Buffer buffer, Integer ns_id, Integer line, Integer
 
     extmark_set(buf, (uint32_t)ns_id, &id, (int)line, (colnr_T)col, line2, col2,
                 has_decor ? &decor : NULL, right_gravity, end_right_gravity,
-                kExtmarkNoUndo, err);
+                !GET_BOOL_OR_TRUE(opts, set_extmark, undo_restore), err);
     if (ERROR_SET(err)) {
       goto error;
     }
@@ -952,7 +959,7 @@ Integer nvim_buf_add_highlight(Buffer buffer, Integer ns_id, String hl_group, In
   extmark_set(buf, ns, NULL,
               (int)line, (colnr_T)col_start,
               end_line, (colnr_T)col_end,
-              &decor, true, false, kExtmarkNoUndo, NULL);
+              &decor, true, false, false, NULL);
   return ns_id;
 }
 
diff --git a/src/nvim/api/keysets.h b/src/nvim/api/keysets.h
index 18f6d5368b..d435262ff3 100644
--- a/src/nvim/api/keysets.h
+++ b/src/nvim/api/keysets.h
@@ -48,6 +48,7 @@ typedef struct {
   String conceal;
   Boolean spell;
   Boolean ui_watched;
+  Boolean undo_restore;
 } Dict(set_extmark);
 
 typedef struct {
diff --git a/src/nvim/decoration.c b/src/nvim/decoration.c
index 59e0711906..24f6693fe2 100644
--- a/src/nvim/decoration.c
+++ b/src/nvim/decoration.c
@@ -66,7 +66,7 @@ void bufhl_add_hl_pos_offset(buf_T *buf, int src_id, int hl_id, lpos_T pos_start
     }
     extmark_set(buf, (uint32_t)src_id, NULL,
                 (int)lnum - 1, hl_start, (int)lnum - 1 + end_off, hl_end,
-                &decor, true, false, kExtmarkNoUndo, NULL);
+                &decor, true, false, true, NULL);
   }
 }
 
diff --git a/src/nvim/extmark.c b/src/nvim/extmark.c
index a988e864db..42ca99d706 100644
--- a/src/nvim/extmark.c
+++ b/src/nvim/extmark.c
@@ -56,7 +56,7 @@
 /// must not be used during iteration!
 void extmark_set(buf_T *buf, uint32_t ns_id, uint32_t *idp, int row, colnr_T col, int end_row,
                  colnr_T end_col, Decoration *decor, bool right_gravity, bool end_right_gravity,
-                 ExtmarkOp op, Error *err)
+                 bool no_undo, Error *err)
 {
   uint32_t *ns = map_put_ref(uint32_t, uint32_t)(buf->b_extmark_ns, ns_id, NULL, NULL);
   uint32_t id = idp ? *idp : 0;
@@ -126,7 +126,7 @@ void extmark_set(buf_T *buf, uint32_t ns_id, uint32_t *idp, int row, colnr_T col
   }
 
   MTKey mark = { { row, col }, ns_id, id, 0,
-                 mt_flags(right_gravity, decor_level), 0, NULL };
+                 mt_flags(right_gravity, decor_level, no_undo), 0, NULL };
   if (decor_full) {
     mark.decor_full = decor;
   } else if (decor) {
@@ -139,15 +139,6 @@ void extmark_set(buf_T *buf, uint32_t ns_id, uint32_t *idp, int row, colnr_T col
   marktree_put(buf->b_marktree, mark, end_row, end_col, end_right_gravity);
 
 revised:
-  if (op != kExtmarkNoUndo) {
-    // TODO(bfredl): this doesn't cover all the cases and probably shouldn't
-    // be done "prematurely". Any movement in undo history might necessitate
-    // adding new marks to old undo headers. add a test case for this (doesn't
-    // fail extmark_spec.lua, and it should)
-    uint64_t mark_id = mt_lookup_id(ns_id, id, false);
-    u_extmark_set(buf, mark_id, row, col);
-  }
-
   if (decor) {
     if (kv_size(decor->virt_text) && decor->virt_text_pos == kVTInline) {
       buf->b_virt_text_inline++;
@@ -359,13 +350,14 @@ static void push_mark(ExtmarkInfoArray *array, uint32_t ns_id, ExtmarkType type_
                                    .end_col = end_pos.col,
                                    .right_gravity = mt_right(mark),
                                    .end_right_gravity = end_right,
+                                   .no_undo = mt_no_undo(mark),
                                    .decor = get_decor(mark) }));
 }
 
 /// Lookup an extmark by id
 ExtmarkInfo extmark_from_id(buf_T *buf, uint32_t ns_id, uint32_t id)
 {
-  ExtmarkInfo ret = { 0, 0, -1, -1, -1, -1, false, false, DECORATION_INIT };
+  ExtmarkInfo ret = { 0, 0, -1, -1, -1, -1, false, false, false, DECORATION_INIT };
   MTKey mark = marktree_lookup_ns(buf->b_marktree, ns_id, id, false, NULL);
   if (!mark.id) {
     return ret;
@@ -381,6 +373,7 @@ ExtmarkInfo extmark_from_id(buf_T *buf, uint32_t ns_id, uint32_t id)
   ret.end_col = end.pos.col;
   ret.right_gravity = mt_right(mark);
   ret.end_right_gravity = mt_right(end);
+  ret.no_undo = mt_no_undo(mark);
   ret.decor = get_decor(mark);
 
   return ret;
@@ -415,27 +408,6 @@ void extmark_free_all(buf_T *buf)
   *buf->b_extmark_ns = (Map(uint32_t, uint32_t)) MAP_INIT;
 }
 
-/// Save info for undo/redo of set marks
-static void u_extmark_set(buf_T *buf, uint64_t mark, int row, colnr_T col)
-{
-  u_header_T *uhp = u_force_get_undo_header(buf);
-  if (!uhp) {
-    return;
-  }
-
-  ExtmarkSavePos pos;
-  pos.mark = mark;
-  pos.old_row = -1;
-  pos.old_col = -1;
-  pos.row = row;
-  pos.col = col;
-
-  ExtmarkUndoObject undo = { .type = kExtmarkSavePos,
-                             .data.savepos = pos };
-
-  kv_push(uhp->uh_extmark, undo);
-}
-
 /// copy extmarks data between range
 ///
 /// useful when we cannot simply reverse the operation. This will do nothing on
@@ -458,16 +430,19 @@ void u_extmark_copy(buf_T *buf, int l_row, colnr_T l_col, int u_row, colnr_T u_c
         || (mark.pos.row == u_row && mark.pos.col > u_col)) {
       break;
     }
-    ExtmarkSavePos pos;
-    pos.mark = mt_lookup_key(mark);
-    pos.old_row = mark.pos.row;
-    pos.old_col = mark.pos.col;
-    pos.row = -1;
-    pos.col = -1;
 
-    undo.data.savepos = pos;
-    undo.type = kExtmarkSavePos;
-    kv_push(uhp->uh_extmark, undo);
+    if (!mt_no_undo(mark)) {
+      ExtmarkSavePos pos;
+      pos.mark = mt_lookup_key(mark);
+      pos.old_row = mark.pos.row;
+      pos.old_col = mark.pos.col;
+      pos.row = -1;
+      pos.col = -1;
+
+      undo.data.savepos = pos;
+      undo.type = kExtmarkSavePos;
+      kv_push(uhp->uh_extmark, undo);
+    }
 
     marktree_itr_next(buf->b_marktree, itr);
   }
diff --git a/src/nvim/extmark.h b/src/nvim/extmark.h
index c2103a71bf..83c7a7cb69 100644
--- a/src/nvim/extmark.h
+++ b/src/nvim/extmark.h
@@ -25,6 +25,7 @@ typedef struct {
   colnr_T end_col;
   bool right_gravity;
   bool end_right_gravity;
+  bool no_undo;
   Decoration decor;  // TODO(bfredl): CHONKY
 } ExtmarkInfo;
 
diff --git a/src/nvim/marktree.c b/src/nvim/marktree.c
index 8ab9370e6c..3df659b8e1 100644
--- a/src/nvim/marktree.c
+++ b/src/nvim/marktree.c
@@ -2007,7 +2007,7 @@ void marktree_put_test(MarkTree *b, uint32_t ns, uint32_t id, int row, int col,
                        int end_row, int end_col, bool end_right)
 {
   MTKey key = { { row, col }, ns, id, 0,
-                mt_flags(right_gravity, 0), 0, NULL };
+                mt_flags(right_gravity, 0, false), 0, NULL };
   marktree_put(b, key, end_row, end_col, end_right);
 }
 
diff --git a/src/nvim/marktree.h b/src/nvim/marktree.h
index 6f157936cb..79e8f8f50e 100644
--- a/src/nvim/marktree.h
+++ b/src/nvim/marktree.h
@@ -77,18 +77,18 @@ typedef struct {
 // orphaned: the other side of this paired mark was deleted. this mark must be deleted very soon!
 #define MT_FLAG_ORPHANED (((uint16_t)1) << 3)
 #define MT_FLAG_HL_EOL (((uint16_t)1) << 4)
+#define MT_FLAG_NO_UNDO (((uint16_t)1) << 5)
 
 #define DECOR_LEVELS 4
-#define MT_FLAG_DECOR_OFFSET 5
+#define MT_FLAG_DECOR_OFFSET 6
 #define MT_FLAG_DECOR_MASK (((uint16_t)(DECOR_LEVELS - 1)) << MT_FLAG_DECOR_OFFSET)
 
-// next flag is (((uint16_t)1) << 6)
-
 // These _must_ be last to preserve ordering of marks
 #define MT_FLAG_RIGHT_GRAVITY (((uint16_t)1) << 14)
 #define MT_FLAG_LAST (((uint16_t)1) << 15)
 
-#define MT_FLAG_EXTERNAL_MASK (MT_FLAG_DECOR_MASK | MT_FLAG_RIGHT_GRAVITY | MT_FLAG_HL_EOL)
+#define MT_FLAG_EXTERNAL_MASK (MT_FLAG_DECOR_MASK | MT_FLAG_RIGHT_GRAVITY | \
+                               MT_FLAG_NO_UNDO | MT_FLAG_HL_EOL)
 
 // this is defined so that start and end of the same range have adjacent ids
 #define MARKTREE_END_FLAG ((uint64_t)1)
@@ -127,16 +127,22 @@ static inline bool mt_right(MTKey key)
   return key.flags & MT_FLAG_RIGHT_GRAVITY;
 }
 
+static inline bool mt_no_undo(MTKey key)
+{
+  return key.flags & MT_FLAG_NO_UNDO;
+}
+
 static inline uint8_t marktree_decor_level(MTKey key)
 {
   return (uint8_t)((key.flags&MT_FLAG_DECOR_MASK) >> MT_FLAG_DECOR_OFFSET);
 }
 
-static inline uint16_t mt_flags(bool right_gravity, uint8_t decor_level)
+static inline uint16_t mt_flags(bool right_gravity, uint8_t decor_level, bool no_undo)
 {
   assert(decor_level < DECOR_LEVELS);
   return (uint16_t)((right_gravity ? MT_FLAG_RIGHT_GRAVITY : 0)
-                    | (decor_level << MT_FLAG_DECOR_OFFSET));
+                    | (decor_level << MT_FLAG_DECOR_OFFSET)
+                    | (no_undo ? MT_FLAG_NO_UNDO : 0));
 }
 
 typedef kvec_withinit_t(uint64_t, 4) Intersection;
diff --git a/test/functional/ui/decorations_spec.lua b/test/functional/ui/decorations_spec.lua
index e56f82bd9f..68c3e73e61 100644
--- a/test/functional/ui/decorations_spec.lua
+++ b/test/functional/ui/decorations_spec.lua
@@ -11,6 +11,7 @@ local meths = helpers.meths
 local funcs = helpers.funcs
 local curbufmeths = helpers.curbufmeths
 local command = helpers.command
+local eq = helpers.eq
 local assert_alive = helpers.assert_alive
 
 describe('decorations providers', function()
@@ -1997,6 +1998,60 @@ describe('extmark decorations', function()
     ]]}
   end)
 
+  local function with_undo_restore(val)
+    screen:try_resize(50, 5)
+    insert(example_text)
+    feed'gg'
+    meths.buf_set_extmark(0, ns, 0, 6, { end_col=13, hl_group = 'NonText', undo_restore=val})
+    screen:expect{grid=[[
+      ^for _,{1:item in} ipairs(items) do                    |
+          local text, hl_id_cell, count = unpack(item)  |
+          if hl_id_cell ~= nil then                     |
+              hl_id = hl_id_cell                        |
+                                                        |
+    ]]}
+
+    meths.buf_set_text(0, 0, 4, 0, 8, {''})
+    screen:expect{grid=[[
+      ^for {1:em in} ipairs(items) do                        |
+          local text, hl_id_cell, count = unpack(item)  |
+          if hl_id_cell ~= nil then                     |
+              hl_id = hl_id_cell                        |
+                                                        |
+    ]]}
+  end
+
+  it("highlights do reapply to restored text after delete", function()
+    with_undo_restore(true) -- also default behavior
+
+    feed 'u'
+    screen:expect{grid=[[
+      ^for _,{1:item in} ipairs(items) do                    |
+          local text, hl_id_cell, count = unpack(item)  |
+          if hl_id_cell ~= nil then                     |
+              hl_id = hl_id_cell                        |
+      1 change; before #2  0 seconds ago                |
+    ]]}
+  end)
+
+  it("highlights don't reapply to restored text after delete with no_undo_restore", function()
+    with_undo_restore(false)
+
+    feed 'u'
+    screen:expect{grid=[[
+      ^for _,it{1:em in} ipairs(items) do                    |
+          local text, hl_id_cell, count = unpack(item)  |
+          if hl_id_cell ~= nil then                     |
+              hl_id = hl_id_cell                        |
+      1 change; before #2  0 seconds ago                |
+    ]]}
+
+    eq({ { 1, 0, 8, { end_col = 13, end_right_gravity = false, end_row = 0,
+                       hl_eol = false, hl_group = "NonText", undo_restore = false,
+                       ns_id = 1, priority = 4096, right_gravity = true } } },
+       meths.buf_get_extmarks(0, ns, {0,0}, {0, -1}, {details=true}))
+  end)
+
   it('virtual text works with rightleft', function()
     screen:try_resize(50, 3)
     insert('abcdefghijklmn')