From 995dbd2ca6c112dc745f226690a99d79b9dc51c9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bj=C3=B6rn=20Linse?= <bjorn.linse@gmail.com>
Date: Mon, 1 Nov 2021 22:53:22 +0100
Subject: [PATCH] fix(extmark): fix missing virt_lines when using id param of
 set_extmark

---
 src/nvim/extmark.c                      | 19 +++++++++----------
 src/nvim/marktree.c                     |  4 ++--
 test/functional/ui/decorations_spec.lua |  6 +++---
 3 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/src/nvim/extmark.c b/src/nvim/extmark.c
index 8726abc59d..c4d8f75a21 100644
--- a/src/nvim/extmark.c
+++ b/src/nvim/extmark.c
@@ -68,6 +68,14 @@ uint64_t extmark_set(buf_T *buf, uint64_t ns_id, uint64_t *idp, int row, colnr_T
   uint64_t mark = 0;
   uint64_t id = idp ? *idp : 0;
 
+  uint8_t decor_level = kDecorLevelNone;  // no decor
+  if (decor) {
+    decor_level = kDecorLevelVisible;  // decor affects redraw
+    if (kv_size(decor->virt_lines)) {
+      decor_level = kDecorLevelVirtLine;  // decor affects horizontal size
+    }
+  }
+
   if (id == 0) {
     id = ns->free_id++;
   } else {
@@ -76,7 +84,6 @@ uint64_t extmark_set(buf_T *buf, uint64_t ns_id, uint64_t *idp, int row, colnr_T
       if (old_mark & MARKTREE_PAIRED_FLAG || end_row > -1) {
         extmark_del(buf, ns_id, id);
       } else {
-        // TODO(bfredl): we need to do more if "revising" a decoration mark.
         MarkTreeIter itr[1] = { 0 };
         old_pos = marktree_lookup(buf->b_marktree, old_mark, itr);
         assert(itr->node);
@@ -86,7 +93,7 @@ uint64_t extmark_set(buf_T *buf, uint64_t ns_id, uint64_t *idp, int row, colnr_T
           if (it.decor) {
             decor_remove(buf, row, row, it.decor);
           }
-          mark = marktree_revise(buf->b_marktree, itr);
+          mark = marktree_revise(buf->b_marktree, itr, decor_level);
           goto revised;
         }
         marktree_del_itr(buf->b_marktree, itr, false);
@@ -96,14 +103,6 @@ uint64_t extmark_set(buf_T *buf, uint64_t ns_id, uint64_t *idp, int row, colnr_T
     }
   }
 
-  uint8_t decor_level = kDecorLevelNone;  // no decor
-  if (decor) {
-    decor_level = kDecorLevelVisible;  // decor affects redraw
-    if (kv_size(decor->virt_lines)) {
-      decor_level = kDecorLevelVirtLine;  // decor affects horizontal size
-    }
-  }
-
   if (end_row > -1) {
     mark = marktree_put_pair(buf->b_marktree,
                              row, col, right_gravity,
diff --git a/src/nvim/marktree.c b/src/nvim/marktree.c
index a7f540c748..38014ab375 100644
--- a/src/nvim/marktree.c
+++ b/src/nvim/marktree.c
@@ -570,11 +570,11 @@ void marktree_free_node(mtnode_t *x)
 }
 
 /// NB: caller must check not pair!
-uint64_t marktree_revise(MarkTree *b, MarkTreeIter *itr)
+uint64_t marktree_revise(MarkTree *b, MarkTreeIter *itr, uint8_t decor_level)
 {
   uint64_t old_id = rawkey(itr).id;
   pmap_del(uint64_t)(b->id2node, ANTIGRAVITY(old_id));
-  uint64_t new_id = (b->next_id += ID_INCR);
+  uint64_t new_id = (b->next_id += ID_INCR) + ((uint64_t)decor_level << DECOR_OFFSET);
   rawkey(itr).id = new_id + (RIGHT_GRAVITY&old_id);
   refkey(b, itr->node, itr->i);
   return new_id;
diff --git a/test/functional/ui/decorations_spec.lua b/test/functional/ui/decorations_spec.lua
index 40c33d9a4d..dce6384b9b 100644
--- a/test/functional/ui/decorations_spec.lua
+++ b/test/functional/ui/decorations_spec.lua
@@ -992,7 +992,7 @@ if (h->n_buckets < new_n_buckets) { // expand
     ]]}
   end)
 
-  it('works with text et the end of the buffer', function()
+  it('works with text at the end of the buffer', function()
     insert(example_text)
     feed 'G'
 
@@ -1213,7 +1213,7 @@ if (h->n_buckets < new_n_buckets) { // expand
                                                         |
     ]]}
 
-    meths.buf_set_extmark(0, ns, 2, 0, {
+    local markid = meths.buf_set_extmark(0, ns, 2, 0, {
       virt_lines={
         {{"Some special", "Special"}};
         {{"remark about codes", "Comment"}};
@@ -1235,13 +1235,13 @@ if (h->n_buckets < new_n_buckets) { // expand
                                                         |
     ]]}
 
-    meths.buf_clear_namespace(0, ns, 0, -1)
     meths.buf_set_extmark(0, ns, 2, 0, {
       virt_lines={
         {{"Some special", "Special"}};
         {{"remark about codes", "Comment"}};
       };
       virt_lines_leftcol=true;
+      id=markid;
     })
     screen:expect{grid=[[
       {8:  }{9:  1 }^if (h->n_buckets < new_n_buckets) { // expan|