fix(column)!: ensure 'statuscolumn' works with virtual and wrapped lines

Problem:    The `'statuscolumn'` was not re-evaluated for wrapped lines,
            when preceded by virtual/filler lines. There was also no way
            to distinguish virtual and wrapped lines in the status column.
Solution:   Make sure to rebuild the statuscolumn, and replace variable
            `v:wrap` with `v:virtnum`. `v:virtnum` is negative when drawing
            virtual lines, zero when drawing the actual buffer line, and
            positive when drawing the wrapped part of a buffer line.

8 files changed, 80 insertions, 44 deletions

diff --git a/runtime/doc/eval.txt b/runtime/doc/eval.txt
index 6feb5cbb49..0ac8152144 100644
--- a/runtime/doc/eval.txt
+++ b/runtime/doc/eval.txt
@@ -2295,13 +2295,15 @@ v:version	Vim version number: major version times 100 plus minor
 			:if has("nvim-0.2.1")
 <
 
-				*v:vim_did_enter* *vim_did_enter-variable*
-v:vim_did_enter	0 during startup, 1 just before |VimEnter|.
+						*v:virtnum* *virtnum-variable*
+v:virtnum	Virtual line number for the 'statuscolumn' expression.
+	        Negative when drawing the status column for virtual lines, zero
+		when drawing an actual buffer line, and positive when drawing
+		the wrapped part of a buffer line.
 		Read-only.
 
-						*v:wrap* *wrap-variable*
-v:wrap		Boolean indicating whether 'statuscolumn' is being evaluated
-		for the wrapped part of a line.
+				*v:vim_did_enter* *vim_did_enter-variable*
+v:vim_did_enter	0 during startup, 1 just before |VimEnter|.
 		Read-only.
 
 					*v:warningmsg* *warningmsg-variable*
diff --git a/runtime/doc/options.txt b/runtime/doc/options.txt
index 4498dda300..efeeba983e 100644
--- a/runtime/doc/options.txt
+++ b/runtime/doc/options.txt
@@ -6017,12 +6017,15 @@ A jump table for the options with a short description can be found at |Q_op|.
 	%s	sign column for currently drawn line
 	%C	fold column for currently drawn line
 
-	To draw the sign and fold columns, they must be included in
-	'statuscolumn'.
+	NOTE: To draw the sign and fold columns, their items must be included in
+	'statuscolumn'. Even when they are not included, the status column width
+	will adapt to the 'signcolumn' and 'foldcolumn' width.
 
-	The |v:lnum|   variable holds the line number to be drawn.
-	The |v:relnum| variable holds the relative line number to be drawn.
-	The |v:wrap|   variable holds true for the wrapped part of a line.
+	The |v:lnum|    variable holds the line number to be drawn.
+	The |v:relnum|  variable holds the relative line number to be drawn.
+	The |v:virtnum| variable is negative when drawing virtual lines, zero
+		      when drawing the actual buffer line, and positive when
+		      drawing the wrapped part of a buffer line.
 
 	Examples: >vim
 		" Relative number with bar separator and click handlers:
@@ -6032,7 +6035,7 @@ A jump table for the options with a short description can be found at |Q_op|.
 		:let &stc='%=%{v:relnum?v:relnum:v:lnum} '
 
 		" Line numbers in hexadecimal for non wrapped part of lines:
-		:let &stc='%=%{v:wrap?"":printf("%x",v:lnum)} '
+		:let &stc='%=%{v:virtnum>0?"":printf("%x",v:lnum)} '
 
 		" Human readable line numbers with thousands separator:
 		:let &stc='%{substitute(v:lnum,"\\d\\zs\\ze\\'
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index 77158e233a..c1a78cc775 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -399,40 +399,45 @@ static int get_sign_attrs(buf_T *buf, linenr_T lnum, SignTextAttrs *sattrs, int
 
 /// Prepare and build the 'statuscolumn' string for line "lnum" in window "wp".
 /// Fill "stcp" with the built status column string and attributes.
+/// This can be called three times per win_line(), once for virt_lines, once for
+/// the start of the buffer line "lnum" and once for the wrapped lines.
 ///
 /// @param[out] stcp  Status column attributes
 static void get_statuscol_str(win_T *wp, linenr_T lnum, int row, int startrow, int filler_lines,
-                              int cul_attr, int sign_num_attr, int sign_cul_attr, char_u *extra,
-                              foldinfo_T foldinfo, SignTextAttrs *sattrs, statuscol_T *stcp)
+                              int cul_attr, int sign_num_attr, int sign_cul_attr, statuscol_T *stcp,
+                              foldinfo_T foldinfo, SignTextAttrs *sattrs)
 {
-  long relnum = 0;
-  bool wrapped = row != startrow + filler_lines;
+  long relnum = -1;
   bool use_cul = use_cursor_line_sign(wp, lnum);
+  int virtnum = row - startrow - filler_lines;
 
-  // Set num, fold and sign text and attrs, empty when wrapped
+  set_vim_var_nr(VV_VIRTNUM, virtnum);
+  // When called the first time for line "lnum" set num_attr
   if (row == startrow) {
-    relnum = labs(get_cursor_rel_lnum(wp, lnum));
     stcp->num_attr = sign_num_attr ? sign_num_attr
                      : get_line_number_attr(wp, lnum, row, startrow, filler_lines);
-
+  }
+  // When called for the first non-filler row of line "lnum" set num v:vars and fold column
+  if (virtnum == 0) {
+    relnum = labs(get_cursor_rel_lnum(wp, lnum));
     if (compute_foldcolumn(wp, 0)) {
       size_t n = fill_foldcolumn(stcp->fold_text, wp, foldinfo, lnum);
       stcp->fold_text[n] = NUL;
       stcp->fold_attr = win_hl_attr(wp, use_cul ? HLF_CLF : HLF_FC);
     }
   }
-
+  // Make sure to clear->set->clear sign column for filler->first->wrapped lines
   int i = 0;
   for (; i < wp->w_scwidth; i++) {
-    SignTextAttrs *sattr = wrapped ? NULL : sign_get_attr(i, sattrs, wp->w_scwidth);
+    SignTextAttrs *sattr = virtnum ? NULL : sign_get_attr(i, sattrs, wp->w_scwidth);
     stcp->sign_text[i] = sattr && sattr->text ? sattr->text : "  ";
     stcp->sign_attr[i] = sattr ? (use_cul && sign_cul_attr ? sign_cul_attr : sattr->hl_attr_id)
                                : win_hl_attr(wp, use_cul ? HLF_CLS : HLF_SC);
   }
   stcp->sign_text[i] = NULL;
 
-  int width = build_statuscol_str(wp, row == startrow, wrapped, lnum, relnum,
-                                  stcp->width, ' ', stcp->text, &stcp->hlrec, stcp);
+  int width = build_statuscol_str(wp, lnum, relnum, stcp->width,
+                                  ' ', stcp->text, &stcp->hlrec, stcp);
   // Force a redraw in case of error or when truncated
   if (*wp->w_p_stc == NUL || (stcp->truncate > 0 && wp->w_nrwidth < MAX_NUMBERWIDTH)) {
     if (stcp->truncate) {  // Avoid truncating 'statuscolumn'
@@ -465,9 +470,8 @@ static void get_statuscol_str(win_T *wp, linenr_T lnum, int row, int startrow, i
 ///
 /// @param stcp  Status column attributes
 /// @param[out] draw_state  Current draw state in win_line()
-static void get_statuscol_display_info(LineDrawState *draw_state, int *char_attr, int *n_extrap,
-                                       int *c_extrap, int *c_finalp, char_u *extra, char **pp_extra,
-                                       statuscol_T *stcp)
+static void get_statuscol_display_info(statuscol_T *stcp, LineDrawState *draw_state, int *char_attr,
+                                       int *n_extrap, int *c_extrap, int *c_finalp, char **pp_extra)
 {
   *c_extrap = NUL;
   *c_finalp = NUL;
@@ -1326,14 +1330,14 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool nochange,
         // Draw the 'statuscolumn' if option is set.
         if (statuscol.draw) {
           if (statuscol.textp == NULL) {
-            get_statuscol_str(wp, lnum, row, startrow, filler_lines, cul_attr, sign_num_attr,
-                              sign_cul_attr, extra, foldinfo, sattrs, &statuscol);
+            get_statuscol_str(wp, lnum, row, startrow, filler_lines, cul_attr,
+                              sign_num_attr, sign_cul_attr, &statuscol, foldinfo, sattrs);
             if (wp->w_redr_statuscol) {
-              return 0;
+              break;
             }
           }
-          get_statuscol_display_info(&draw_state, &char_attr, &n_extra, &c_extra,
-                                     &c_final, extra, &p_extra, &statuscol);
+          get_statuscol_display_info(&statuscol, &draw_state, &char_attr,
+                                     &n_extra, &c_extra, &c_final, &p_extra);
         }
       }
 
@@ -2818,7 +2822,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool nochange,
         need_showbreak = true;
       }
       if (statuscol.draw) {
-        if (row == startrow + 1 || row == startrow + filler_lines) {
+        if (row == startrow + filler_lines + 1 || row == startrow + filler_lines) {
           // Re-evaluate 'statuscolumn' for the first wrapped row and non filler line
           statuscol.textp = NULL;
         } else {  // Otherwise just reset the text/hlrec pointers
diff --git a/src/nvim/eval.c b/src/nvim/eval.c
index f9825496a5..6d32b71016 100644
--- a/src/nvim/eval.c
+++ b/src/nvim/eval.c
@@ -268,7 +268,7 @@ static struct vimvar {
   VV(VV__NULL_BLOB,       "_null_blob",       VAR_BLOB, VV_RO),
   VV(VV_LUA,              "lua",              VAR_PARTIAL, VV_RO),
   VV(VV_RELNUM,           "relnum",           VAR_NUMBER, VV_RO),
-  VV(VV_WRAP,             "wrap",             VAR_BOOL, VV_RO),
+  VV(VV_VIRTNUM,          "virtnum",          VAR_NUMBER, VV_RO),
 };
 #undef VV
 
diff --git a/src/nvim/eval.h b/src/nvim/eval.h
index d67414f12f..86bc76e793 100644
--- a/src/nvim/eval.h
+++ b/src/nvim/eval.h
@@ -166,7 +166,7 @@ typedef enum {
   VV__NULL_BLOB,  // Blob with NULL value. For test purposes only.
   VV_LUA,
   VV_RELNUM,
-  VV_WRAP,
+  VV_VIRTNUM,
 } VimVarIndex;
 
 /// All recognized msgpack types
diff --git a/src/nvim/screen.c b/src/nvim/screen.c
index b18bf7ed6a..6a3103123b 100644
--- a/src/nvim/screen.c
+++ b/src/nvim/screen.c
@@ -784,7 +784,7 @@ int number_width(win_T *wp)
   if (*wp->w_p_stc != NUL) {
     char buf[MAXPATHL];
     wp->w_nrwidth_width = 0;
-    n = build_statuscol_str(wp, true, false, lnum, 0, 0, NUL, buf, NULL, NULL);
+    n = build_statuscol_str(wp, lnum, 0, 0, NUL, buf, NULL, NULL);
     n = MAX(n, (wp->w_p_nu || wp->w_p_rnu) * (int)wp->w_p_nuw);
     wp->w_nrwidth_width = MIN(n, MAX_NUMBERWIDTH);
     return wp->w_nrwidth_width;
diff --git a/src/nvim/statusline.c b/src/nvim/statusline.c
index b75b3c313f..cb5dde1013 100644
--- a/src/nvim/statusline.c
+++ b/src/nvim/statusline.c
@@ -876,14 +876,13 @@ void draw_tabline(void)
 /// @param hlrec  HL attributes (can be NULL)
 /// @param stcp  Status column attributes (can be NULL)
 /// @return  The width of the built status column string for line "lnum"
-int build_statuscol_str(win_T *wp, bool setnum, bool wrap, linenr_T lnum, long relnum, int maxwidth,
-                        int fillchar, char *buf, stl_hlrec_t **hlrec, statuscol_T *stcp)
+int build_statuscol_str(win_T *wp, linenr_T lnum, long relnum, int maxwidth, int fillchar,
+                        char *buf, stl_hlrec_t **hlrec, statuscol_T *stcp)
 {
-  if (setnum) {
+  if (relnum >= 0) {
     set_vim_var_nr(VV_LNUM, lnum);
     set_vim_var_nr(VV_RELNUM, relnum);
   }
-  set_vim_var_bool(VV_WRAP, wrap);
 
   StlClickRecord *clickrec;
   char *stc = xstrdup(wp->w_p_stc);
@@ -1506,13 +1505,12 @@ int build_stl_str_hl(win_T *wp, char *out, size_t outlen, char *fmt, char *opt_n
     case STL_LINE:
       // Overload %l with v:lnum for 'statuscolumn'
       if (opt_name != NULL && strcmp(opt_name, "statuscolumn") == 0) {
-        if (wp->w_p_nu) {
+        if (wp->w_p_nu && !get_vim_var_nr(VV_VIRTNUM)) {
           num = get_vim_var_nr(VV_LNUM);
         }
       } else {
         num = (wp->w_buffer->b_ml.ml_flags & ML_EMPTY) ? 0L : (long)(wp->w_cursor.lnum);
       }
-
       break;
 
     case STL_NUMLINES:
@@ -1612,7 +1610,7 @@ int build_stl_str_hl(win_T *wp, char *out, size_t outlen, char *fmt, char *opt_n
     case STL_ROFLAG_ALT:
       // Overload %r with v:relnum for 'statuscolumn'
       if (opt_name != NULL && strcmp(opt_name, "statuscolumn") == 0) {
-        if (wp->w_p_rnu) {
+        if (wp->w_p_rnu && !get_vim_var_nr(VV_VIRTNUM)) {
           num = get_vim_var_nr(VV_RELNUM);
         }
       } else {
diff --git a/test/functional/ui/statuscolumn_spec.lua b/test/functional/ui/statuscolumn_spec.lua
index a9d796c10b..fd57d5dfba 100644
--- a/test/functional/ui/statuscolumn_spec.lua
+++ b/test/functional/ui/statuscolumn_spec.lua
@@ -4,6 +4,7 @@ local clear = helpers.clear
 local command = helpers.command
 local eq = helpers.eq
 local eval = helpers.eval
+local exec = helpers.exec_lua
 local meths = helpers.meths
 local pcall_err = helpers.pcall_err
 
@@ -183,7 +184,7 @@ describe('statuscolumn', function()
   end)
 
   it('works with wrapped lines, signs and folds', function()
-    command("set stc=%C%s%=%{v:wrap?'':v:lnum}│\\ ")
+    command("set stc=%C%s%=%{v:virtnum?'':v:lnum}│\\ ")
     command("call setline(1,repeat([repeat('aaaaa',10)],16))")
     screen:set_default_attr_ids({
       [0] = {bold = true, foreground = Screen.colors.Blue},
@@ -234,7 +235,7 @@ describe('statuscolumn', function()
     ]])
     command('norm zf$')
     -- Check that alignment works properly with signs after %=
-    command("set stc=%C%=%{v:wrap?'':v:lnum}│%s\\ ")
+    command("set stc=%C%=%{v:virtnum?'':v:lnum}│%s\\ ")
     screen:expect([[
       {2: }{1: 4│>>}{2:  }{1: }aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa|
       {2: }{1:  │}{2:    }{1: }aaaaaa                                      |
@@ -304,7 +305,7 @@ describe('statuscolumn', function()
       {2: }{1: 2│}{2:    }{1: }aaaaaa                                      |
                                                            |
     ]])
-    command("set stc=%C%=\\ %{v:wrap?'':v:relnum}│%s\\ ")
+    command([[set stc=%C%=\ %{v:virtnum?'':v:relnum}│%s\ ]])
     screen:expect([[
       {2: }{1: 4│>>}{2:  }{1: }aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa|
       {2: }{1:  │}{2:    }{1: }aaaaaa                                      |
@@ -346,6 +347,34 @@ describe('statuscolumn', function()
       {2: }{1:  │}{2:                  }{1: }aaaaaaaaaaaaaaaaaaaa          |
                                                            |
     ]])
+    -- Status column is re-evaluated for virt_lines, buffer line, and wrapped line
+    exec([[
+      local ns = vim.api.nvim_create_namespace("ns")
+      vim.api.nvim_buf_set_extmark(0, ns, 4, 0, {
+        virt_lines = {{{"virt_line", ""}}, {{"virt_line", ""}}}
+      })
+      vim.api.nvim_buf_set_extmark(0, ns, 5, 0, {
+        virt_lines_above = true, virt_lines = {{{"virt_line above", ""}}, {{"virt_line above", ""}}}
+      })
+    ]])
+    command('set foldcolumn=0 signcolumn=no')
+    command([[set stc=%{v:virtnum<0?'virtual':(!v:virtnum?'buffer':'wrapped')}%=%{'\ '.v:virtnum.'\ '.v:lnum}]])
+    screen:expect([[
+      {1:buffer  0 4}aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa|
+      {1:wrapped 1 4}aaaaaaaa                                  |
+      {1:buffer  0 5}aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa|
+      {1:wrapped 1 5}aaaaaaaa                                  |
+      {1:virtual-4 5}virt_line                                 |
+      {1:virtual-4 5}virt_line                                 |
+      {1:virtual-4 5}virt_line above                           |
+      {1:virtual-4 5}virt_line above                           |
+      {1:buffer  0 6}aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa|
+      {1:wrapped 1 6}aaaaaaaa                                  |
+      {1:buffer  0 7}aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa|
+      {1:wrapped 1 7}aaaaaaaa                                  |
+      {4:buffer  0 8}{5:^+--  1 line: aaaaaaaaaaaaaaaaaaaaaaaaaaaaa}|
+                                                           |
+    ]])
   end)
 
   it('works with \'statuscolumn\' clicks', function()