4 files changed, 29 insertions, 13 deletions

diff --git a/src/nvim/decoration.c b/src/nvim/decoration.c
index 4e80528c74..931c3fa20a 100644
--- a/src/nvim/decoration.c
+++ b/src/nvim/decoration.c
@@ -118,6 +118,8 @@ Decoration *decor_find_virttext(buf_T *buf, int row, uint64_t ns_id)
     mtmark_t mark = marktree_itr_current(itr);
     if (mark.row < 0 || mark.row > row) {
       break;
+    } else if (mt_decor_level(mark.id) < 1) {
+      goto next_mark;
     }
     ExtmarkItem *item = map_ref(uint64_t, ExtmarkItem)(buf->b_extmark_index,
                                                        mark.id, false);
@@ -125,6 +127,7 @@ Decoration *decor_find_virttext(buf_T *buf, int row, uint64_t ns_id)
         && item->decor && kv_size(item->decor->virt_text)) {
       return item->decor;
     }
+next_mark:
     marktree_itr_next(buf->b_marktree, itr);
   }
   return NULL;
@@ -158,21 +161,22 @@ bool decor_redraw_start(buf_T *buf, int top_row, DecorState *state)
     if (mark.row < 0) {  // || mark.row > end_row
       break;
     }
-    if ((mark.row < top_row && mark.id&MARKTREE_END_FLAG)) {
+    if ((mark.row < top_row && mark.id&MARKTREE_END_FLAG)
+        || mt_decor_level(mark.id) < 1) {
       goto next_mark;
     }
-    mtpos_t altpos = marktree_lookup(buf->b_marktree,
-                                     mark.id^MARKTREE_END_FLAG, NULL);
 
     uint64_t start_id = mark.id & ~MARKTREE_END_FLAG;
     ExtmarkItem *item = map_ref(uint64_t, ExtmarkItem)(buf->b_extmark_index,
                                                        start_id, false);
     if (!item || !item->decor) {
-      // TODO(bfredl): dedicated flag for being a decoration?
       goto next_mark;
     }
     Decoration *decor = item->decor;
 
+    mtpos_t altpos = marktree_lookup(buf->b_marktree,
+                                     mark.id^MARKTREE_END_FLAG, NULL);
+
     if ((!(mark.id&MARKTREE_END_FLAG) && altpos.row < top_row
          && !kv_size(decor->virt_text))
         || ((mark.id&MARKTREE_END_FLAG) && altpos.row >= top_row)) {
@@ -251,21 +255,20 @@ int decor_redraw_col(buf_T *buf, int col, int win_col, bool hidden, DecorState *
       break;
     }
 
-    if ((mark.id&MARKTREE_END_FLAG)) {
-      // TODO(bfredl): check decoration flag
+    if ((mark.id&MARKTREE_END_FLAG) || mt_decor_level(mark.id) < 1) {
       goto next_mark;
     }
-    mtpos_t endpos = marktree_lookup(buf->b_marktree,
-                                     mark.id|MARKTREE_END_FLAG, NULL);
 
     ExtmarkItem *item = map_ref(uint64_t, ExtmarkItem)(buf->b_extmark_index,
                                                        mark.id, false);
     if (!item || !item->decor) {
-      // TODO(bfredl): dedicated flag for being a decoration?
       goto next_mark;
     }
     Decoration *decor = item->decor;
 
+    mtpos_t endpos = marktree_lookup(buf->b_marktree,
+                                     mark.id|MARKTREE_END_FLAG, NULL);
+
     if (endpos.row == -1) {
       endpos.row = mark.row;
       endpos.col = mark.col;
diff --git a/src/nvim/extmark.c b/src/nvim/extmark.c
index dc73e34111..dc3dc4268d 100644
--- a/src/nvim/extmark.c
+++ b/src/nvim/extmark.c
@@ -96,12 +96,14 @@ uint64_t extmark_set(buf_T *buf, uint64_t ns_id, uint64_t *idp, int row, colnr_T
     }
   }
 
+  uint8_t decor_level = (decor != NULL) ? 1 : 0;
+
   if (end_row > -1) {
     mark = marktree_put_pair(buf->b_marktree,
                              row, col, right_gravity,
-                             end_row, end_col, end_right_gravity);
+                             end_row, end_col, end_right_gravity, decor_level);
   } else {
-    mark = marktree_put(buf->b_marktree, row, col, right_gravity);
+    mark = marktree_put(buf->b_marktree, row, col, right_gravity, decor_level);
   }
 
 revised:
diff --git a/src/nvim/marktree.c b/src/nvim/marktree.c
index 39c1e36147..a7f540c748 100644
--- a/src/nvim/marktree.c
+++ b/src/nvim/marktree.c
@@ -221,9 +221,11 @@ static inline void marktree_putp_aux(MarkTree *b, mtnode_t *x, mtkey_t k)
   }
 }
 
-uint64_t marktree_put(MarkTree *b, int row, int col, bool right_gravity)
+uint64_t marktree_put(MarkTree *b, int row, int col, bool right_gravity, uint8_t decor_level)
 {
   uint64_t id = (b->next_id+=ID_INCR);
+  assert(decor_level < DECOR_LEVELS);
+  id = id | ((uint64_t)decor_level << DECOR_OFFSET);
   uint64_t keyid = id;
   if (right_gravity) {
     // order all right gravity keys after the left ones, for effortless
@@ -235,9 +237,11 @@ uint64_t marktree_put(MarkTree *b, int row, int col, bool right_gravity)
 }
 
 uint64_t marktree_put_pair(MarkTree *b, int start_row, int start_col, bool start_right, int end_row,
-                           int end_col, bool end_right)
+                           int end_col, bool end_right, uint8_t decor_level)
 {
   uint64_t id = (b->next_id+=ID_INCR)|PAIRED;
+  assert(decor_level < DECOR_LEVELS);
+  id = id | ((uint64_t)decor_level << DECOR_OFFSET);
   uint64_t start_id = id|(start_right?RIGHT_GRAVITY:0);
   uint64_t end_id = id|END_FLAG|(end_right?RIGHT_GRAVITY:0);
   marktree_put_key(b, start_row, start_col, start_id);
diff --git a/src/nvim/marktree.h b/src/nvim/marktree.h
index 02f73b8052..aec5def879 100644
--- a/src/nvim/marktree.h
+++ b/src/nvim/marktree.h
@@ -75,4 +75,11 @@ typedef struct {
 #define MARKTREE_PAIRED_FLAG (((uint64_t)1) << 1)
 #define MARKTREE_END_FLAG (((uint64_t)1) << 0)
 
+#define DECOR_LEVELS 4
+#define DECOR_OFFSET 61
+#define DECOR_MASK (((uint64_t)(DECOR_LEVELS-1)) << DECOR_OFFSET)
+static inline uint8_t mt_decor_level(uint64_t id) {
+  return (uint8_t)((id&DECOR_MASK) >> DECOR_OFFSET);
+}
+
 #endif  // NVIM_MARKTREE_H