feat(column): apply appropriate numhl highlight to virt_lines (#32400)

Problem:  Number and statuscolumn highlighting for virtual lines does
          not take always take on numhl highlights.
Solution: Apply the appropriate numhl highlight to the number/statuscolumn
          of virtual lines, fetching the numhl highlight of the line above
          for `virt_line_above == false` lines.

5 files changed, 108 insertions, 62 deletions

diff --git a/src/nvim/decoration.c b/src/nvim/decoration.c
index 149504f424..3ca3f2904b 100644
--- a/src/nvim/decoration.c
+++ b/src/nvim/decoration.c
@@ -853,7 +853,7 @@ int sign_item_cmp(const void *p1, const void *p2)
 static const uint32_t sign_filter[4] = {[kMTMetaSignText] = kMTFilterSelect,
                                         [kMTMetaSignHL] = kMTFilterSelect };
 
-/// Return the sign attributes on the currently refreshed row.
+/// Return the signs and highest priority sign attributes on a row.
 ///
 /// @param[out] sattrs Output array for sign text and texthl id
 /// @param[out] line_id Highest priority linehl id
@@ -904,17 +904,17 @@ void decor_redraw_signs(win_T *wp, buf_T *buf, int row, SignTextAttrs sattrs[],
 
     for (size_t i = 0; i < kv_size(signs); i++) {
       DecorSignHighlight *sh = kv_A(signs, i).sh;
-      if (idx < len && sh->text[0]) {
+      if (sattrs && idx < len && sh->text[0]) {
         memcpy(sattrs[idx].text, sh->text, SIGN_WIDTH * sizeof(sattr_T));
         sattrs[idx++].hl_id = sh->hl_id;
       }
-      if (*num_id == 0) {
+      if (num_id != NULL && *num_id <= 0) {
         *num_id = sh->number_hl_id;
       }
-      if (*line_id == 0) {
+      if (line_id != NULL && *line_id <= 0) {
         *line_id = sh->line_hl_id;
       }
-      if (*cul_id == 0) {
+      if (cul_id != NULL && *cul_id <= 0) {
         *cul_id = sh->cursorline_hl_id;
       }
     }
@@ -1038,7 +1038,8 @@ bool decor_redraw_eol(win_T *wp, DecorState *state, int *eol_attr, int eol_col)
 static const uint32_t lines_filter[4] = {[kMTMetaLines] = kMTFilterSelect };
 
 /// @param apply_folds Only count virtual lines that are not in folds.
-int decor_virt_lines(win_T *wp, int start_row, int end_row, VirtLines *lines, bool apply_folds)
+int decor_virt_lines(win_T *wp, int start_row, int end_row, int *num_below, VirtLines *lines,
+                     bool apply_folds)
 {
   buf_T *buf = wp->w_buffer;
   if (!buf_meta_total(buf, kMTMetaLines)) {
@@ -1071,6 +1072,9 @@ int decor_virt_lines(win_T *wp, int start_row, int end_row, VirtLines *lines, bo
             if (lines) {
               kv_splice(*lines, vt->data.virt_lines);
             }
+            if (num_below && !above) {
+              (*num_below) += (int)kv_size(vt->data.virt_lines);
+            }
           }
         }
         vt = vt->next;
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index 6be5830f4d..e4fe50ad28 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -82,6 +82,7 @@ typedef struct {
   int line_attr;             ///< attribute for the whole line
   int line_attr_lowprio;     ///< low-priority attribute for the line
   int sign_num_attr;         ///< line number attribute (sign numhl)
+  int prev_num_attr;         ///< previous line's number attribute (sign numhl)
   int sign_cul_attr;         ///< cursorline sign attribute (sign culhl)
 
   int fromcol;               ///< start of inverting
@@ -107,6 +108,7 @@ typedef struct {
   hlf_T diff_hlf;            ///< type of diff highlighting
 
   int n_virt_lines;          ///< nr of virtual lines
+  int n_virt_below;          ///< nr of virtual lines belonging to previous line
   int filler_lines;          ///< nr of filler lines to be drawn
   int filler_todo;           ///< nr of filler lines still to do + 1
   SignTextAttrs sattrs[SIGN_SHOW_MAX];  ///< sign attributes for the sign column
@@ -564,26 +566,41 @@ static bool use_cursor_line_nr(win_T *wp, winlinevars_T *wlv)
                  && (wp->w_p_culopt_flags & kOptCuloptFlagLine)));
 }
 
+/// Return line number attribute, combining the appropriate LineNr* highlight
+/// with the highest priority sign numhl highlight, if any.
 static int get_line_number_attr(win_T *wp, winlinevars_T *wlv)
 {
+  int numhl_attr = wlv->sign_num_attr;
+
+  // Get previous sign numhl for virt_lines belonging to the previous line.
+  if ((wlv->n_virt_lines - wlv->filler_todo) < wlv->n_virt_below) {
+    if (wlv->prev_num_attr == -1) {
+      decor_redraw_signs(wp, wp->w_buffer, wlv->lnum - 2, NULL, NULL, NULL, &wlv->prev_num_attr);
+      if (wlv->prev_num_attr > 0) {
+        wlv->prev_num_attr = syn_id2attr(wlv->prev_num_attr);
+      }
+    }
+    numhl_attr = wlv->prev_num_attr;
+  }
+
   if (use_cursor_line_nr(wp, wlv)) {
     // TODO(vim): Can we use CursorLine instead of CursorLineNr
     // when CursorLineNr isn't set?
-    return win_hl_attr(wp, HLF_CLN);
+    return hl_combine_attr(win_hl_attr(wp, HLF_CLN), numhl_attr);
   }
 
   if (wp->w_p_rnu) {
     if (wlv->lnum < wp->w_cursor.lnum) {
       // Use LineNrAbove
-      return win_hl_attr(wp, HLF_LNA);
+      return hl_combine_attr(win_hl_attr(wp, HLF_LNA), numhl_attr);
     }
     if (wlv->lnum > wp->w_cursor.lnum) {
       // Use LineNrBelow
-      return win_hl_attr(wp, HLF_LNB);
+      return hl_combine_attr(win_hl_attr(wp, HLF_LNB), numhl_attr);
     }
   }
 
-  return win_hl_attr(wp, HLF_N);
+  return hl_combine_attr(win_hl_attr(wp, HLF_N), numhl_attr);
 }
 
 /// Display the absolute or relative line number.  After the first row fill with
@@ -605,8 +622,7 @@ static void draw_lnum_col(win_T *wp, winlinevars_T *wlv)
     } else {
       // Draw the line number (empty space after wrapping).
       int width = number_width(wp) + 1;
-      int attr = hl_combine_attr(get_line_number_attr(wp, wlv),
-                                 wlv->filler_todo <= 0 ? wlv->sign_num_attr : 0);
+      int attr = get_line_number_attr(wp, wlv);
       if (wlv->row == wlv->startrow + wlv->filler_lines
           && (wp->w_skipcol == 0 || wlv->row > 0 || (wp->w_p_nu && wp->w_p_rnu))) {
         char buf[32];
@@ -678,8 +694,7 @@ static void draw_statuscol(win_T *wp, winlinevars_T *wlv, linenr_T lnum, int vir
   colnr_T *fold_vcol = NULL;
   size_t len = strlen(buf);
   int scl_attr = win_hl_attr(wp, use_cursor_line_highlight(wp, wlv->lnum) ? HLF_CLS : HLF_SC);
-  int num_attr = hl_combine_attr(get_line_number_attr(wp, wlv),
-                                 wlv->filler_todo <= 0 ? wlv->sign_num_attr : 0);
+  int num_attr = get_line_number_attr(wp, wlv);
   int cur_attr = num_attr;
 
   // Draw each segment with the specified highlighting.
@@ -1073,6 +1088,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
     .tocol = MAXCOL,
     .vcol_sbr = -1,
     .old_boguscols = 0,
+    .prev_num_attr = -1,
   };
 
   buf_T *buf = wp->w_buffer;
@@ -1228,7 +1244,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
     area_highlighting = true;
   }
   VirtLines virt_lines = KV_INITIAL_VALUE;
-  wlv.n_virt_lines = decor_virt_lines(wp, lnum - 1, lnum, &virt_lines, true);
+  wlv.n_virt_lines = decor_virt_lines(wp, lnum - 1, lnum, &wlv.n_virt_below, &virt_lines, true);
   wlv.filler_lines += wlv.n_virt_lines;
   if (lnum == wp->w_topline) {
     wlv.filler_lines = wp->w_topfill;
diff --git a/src/nvim/plines.c b/src/nvim/plines.c
index 9bf486fb06..e084d107a0 100644
--- a/src/nvim/plines.c
+++ b/src/nvim/plines.c
@@ -720,7 +720,7 @@ bool win_may_fill(win_T *wp)
 /// @return Number of filler lines above lnum
 int win_get_fill(win_T *wp, linenr_T lnum)
 {
-  int virt_lines = decor_virt_lines(wp, lnum - 1, lnum, NULL, true);
+  int virt_lines = decor_virt_lines(wp, lnum - 1, lnum, NULL, NULL, true);
 
   // be quick when there are no filler lines
   if (diffopt_filler()) {
@@ -920,7 +920,7 @@ int plines_m_win(win_T *wp, linenr_T first, linenr_T last, int max)
 /// Mainly used for calculating scrolling offsets.
 int plines_m_win_fill(win_T *wp, linenr_T first, linenr_T last)
 {
-  int count = last - first + 1 + decor_virt_lines(wp, first - 1, last, NULL, false);
+  int count = last - first + 1 + decor_virt_lines(wp, first - 1, last, NULL, NULL, false);
 
   if (diffopt_filler()) {
     for (int lnum = first; lnum <= last; lnum++) {
diff --git a/test/functional/ui/sign_spec.lua b/test/functional/ui/sign_spec.lua
index ff03d86979..6ca9141a76 100644
--- a/test/functional/ui/sign_spec.lua
+++ b/test/functional/ui/sign_spec.lua
@@ -427,7 +427,7 @@ describe('Signs', function()
       feed('<C-Y>')
       -- number column on virtual lines should be empty
       screen:expect([[
-        {8:    }VIRT LINES                                       |
+        {9:    }VIRT LINES                                       |
         {101: >> }aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa|
         {9:    }aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa|
         {9:    }aa^a                                              |
diff --git a/test/functional/ui/statuscolumn_spec.lua b/test/functional/ui/statuscolumn_spec.lua
index 328e212a22..6c415b3dfa 100644
--- a/test/functional/ui/statuscolumn_spec.lua
+++ b/test/functional/ui/statuscolumn_spec.lua
@@ -18,6 +18,16 @@ describe('statuscolumn', function()
   before_each(function()
     clear('--cmd', 'set number nuw=1 | call setline(1, repeat(["aaaaa"], 16)) | norm GM')
     screen = Screen.new()
+    screen:add_extra_attr_ids {
+      [100] = { foreground = Screen.colors.Red, background = Screen.colors.LightGray },
+      [101] = { background = Screen.colors.Gray90, bold = true },
+      [102] = { foreground = Screen.colors.Brown, background = Screen.colors.Grey },
+      [103] = { bold = true, background = Screen.colors.Grey, foreground = Screen.colors.Blue1 },
+      [104] = { undercurl = true, special = Screen.colors.Red },
+      [105] = { foreground = Screen.colors.Red, underline = true },
+      [106] = { foreground = Screen.colors.Orange1 },
+      [107] = { foreground = Screen.colors.LightBlue },
+    }
     exec_lua('ns = vim.api.nvim_create_namespace("")')
   end)
 
@@ -238,12 +248,6 @@ describe('statuscolumn', function()
   end)
 
   it('works with wrapped lines, signs and folds', function()
-    screen:add_extra_attr_ids {
-      [100] = { foreground = Screen.colors.Red, background = Screen.colors.LightGray },
-      [101] = { background = Screen.colors.Gray90, bold = true },
-      [102] = { foreground = Screen.colors.Brown, background = Screen.colors.Grey },
-      [103] = { bold = true, background = Screen.colors.Grey, foreground = Screen.colors.Blue1 },
-    }
     command([[set cursorline stc=%C%s%=%{v:virtnum?'':v:lnum}│\ ]])
     command("call setline(1,repeat([repeat('aaaaa',10)],16))")
     command('hi! CursorLine gui=bold')
@@ -1007,15 +1011,12 @@ describe('statuscolumn', function()
 
   it('does not wrap multibyte characters at the end of a line', function()
     screen:try_resize(33, 4)
-    screen:add_extra_attr_ids {
-      [100] = { undercurl = true, special = Screen.colors.Red },
-    }
     command([[set spell stc=%l\ ]])
     command('call setline(8, "This is a line that contains ᶏ multibyte character.")')
     screen:expect([[
-      {8: 8 }^This is a line that contains {100:ᶏ}|
-      {8:   } {100:multibyte} character.         |
-      {8: 9 }{100:aaaaa}                         |
+      {8: 8 }^This is a line that contains {104:ᶏ}|
+      {8:   } {104:multibyte} character.         |
+      {8: 9 }{104:aaaaa}                         |
                                        |
     ]])
   end)
@@ -1044,45 +1045,70 @@ describe('statuscolumn', function()
       au InsertLeave * let g:insert = v:false | call nvim__redraw(#{statuscolumn:1, win:0})
     ]])
     feed('i')
-    screen:expect({
-      grid = [[
-        {8:insert}^aaaaa         │aaaaa              |
-        {8:insert}aaaaa         │aaaaa              |
-        {3:[No Name] [+]        }{2:[No Name] [+]      }|
-        {5:-- INSERT --}                            |
-      ]],
-    })
+    screen:expect([[
+      {8:insert}^aaaaa         │aaaaa              |
+      {8:insert}aaaaa         │aaaaa              |
+      {3:[No Name] [+]        }{2:[No Name] [+]      }|
+      {5:-- INSERT --}                            |
+    ]])
     feed('<esc>')
-    screen:expect({
-      grid = [[
-        ^aaaaa               │aaaaa              |
-        aaaaa               │aaaaa              |
-        {3:[No Name] [+]        }{2:[No Name] [+]      }|
-                                                |
-      ]],
-    })
+    screen:expect([[
+      ^aaaaa               │aaaaa              |
+      aaaaa               │aaaaa              |
+      {3:[No Name] [+]        }{2:[No Name] [+]      }|
+                                              |
+    ]])
     -- All windows
     command([[
       au! InsertEnter * let g:insert = v:true | call nvim__redraw(#{statuscolumn:1})
       au! InsertLeave * let g:insert = v:false | call nvim__redraw(#{statuscolumn:1})
     ]])
     feed('i')
-    screen:expect({
-      grid = [[
-        {8:insert}^aaaaa         │{8:insert}aaaaa        |
-        {8:insert}aaaaa         │{8:insert}aaaaa        |
-        {3:[No Name] [+]        }{2:[No Name] [+]      }|
-        {5:-- INSERT --}                            |
-      ]],
-    })
+    screen:expect([[
+      {8:insert}^aaaaa         │{8:insert}aaaaa        |
+      {8:insert}aaaaa         │{8:insert}aaaaa        |
+      {3:[No Name] [+]        }{2:[No Name] [+]      }|
+      {5:-- INSERT --}                            |
+    ]])
     feed('<esc>')
-    screen:expect({
-      grid = [[
-        ^aaaaa               │aaaaa              |
-        aaaaa               │aaaaa              |
-        {3:[No Name] [+]        }{2:[No Name] [+]      }|
-                                                |
-      ]],
-    })
+    screen:expect([[
+      ^aaaaa               │aaaaa              |
+      aaaaa               │aaaaa              |
+      {3:[No Name] [+]        }{2:[No Name] [+]      }|
+                                              |
+    ]])
+  end)
+
+  it('applies numhl highlight to virtual lines', function()
+    exec_lua([[
+      vim.o.statuscolumn = '%=%{%v:virtnum==0?"%l":v:virtnum>0?"↳":"•"%}│'
+      vim.o.cursorline = true
+      vim.api.nvim_set_hl(0, 'CursorLineNr', { underline = true })
+
+      vim.api.nvim_buf_set_extmark(0, ns, 0, 0, { number_hl_group = 'DiagnosticError' })
+
+      local opts_1 = { number_hl_group = 'DiagnosticWarn', virt_lines = { { { 'Hello' } }, { { 'Hello' } } }, virt_lines_above = true }
+      vim.api.nvim_buf_set_extmark(0, ns, 1, 0, opts_1)
+      opts_1.virt_lines_above = nil
+      vim.api.nvim_buf_set_extmark(0, ns, 1, 0, opts_1)
+
+      local opts_2 = { number_hl_group = 'DiagnosticInfo', virt_lines = { { { 'World' } }, { { 'World' } } }, virt_lines_above = true }
+      vim.api.nvim_buf_set_extmark(0, ns, 2, 0, opts_2)
+      opts_2.virt_lines_above = nil
+      vim.api.nvim_buf_set_extmark(0, ns, 2, 0, opts_2)
+      vim.cmd.norm('gg')
+    ]])
+    screen:expect([[
+      {105: 1│}{21:^aaaaa                                             }|
+      {106: •│}Hello                                             |*2
+      {106: 2│}aaaaa                                             |
+      {106: •│}Hello                                             |*2
+      {107: •│}World                                             |*2
+      {107: 3│}aaaaa                                             |
+      {107: •│}World                                             |*2
+      {8: 4│}aaaaa                                             |
+      {8: 5│}aaaaa                                             |
+                                                           |
+    ]])
   end)
 end)