fix(column): estimate 'statuscolumn' width appropriately

Problem:    The 'statuscolumn' width is being estimated without the
            proper context. In particular, this resulted in the fact
            that a custom fold column could be included in the estimated
            `number_width()`, and doubly added when actually drawing the
            statuscolumn due to `win_col_off()` also adding the
            `'foldcolumn'` width. Resulting in a status column that is
            `'foldcolumn'` cells wider than necessary.
Solution:   Estimate 'statuscolumn' width in `get_statuscol_str()` when
            a buffer's line count has changed.

4 files changed, 36 insertions, 10 deletions

diff --git a/src/nvim/buffer_defs.h b/src/nvim/buffer_defs.h
index f01edd1ad2..4c99191170 100644
--- a/src/nvim/buffer_defs.h
+++ b/src/nvim/buffer_defs.h
@@ -1392,6 +1392,7 @@ struct window_S {
   int w_prev_fraction_row;
 
   linenr_T w_nrwidth_line_count;        // line count when ml_nrwidth_width was computed.
+  linenr_T w_statuscol_line_count;      // line count when 'statuscolumn' width was computed.
   int w_nrwidth_width;                  // nr of chars to print line count.
 
   qf_info_T *w_llist;                 // Location list for this window
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index e24d86b353..01ff207c2b 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -412,7 +412,6 @@ static void get_statuscol_str(win_T *wp, linenr_T lnum, int row, int startrow, i
   bool use_cul = use_cursor_line_sign(wp, lnum);
   int virtnum = row - startrow - filler_lines;
 
-  set_vim_var_nr(VV_VIRTNUM, virtnum);
   // When called the first time for line "lnum" set num_attr
   if (stcp->num_attr == 0) {
     stcp->num_attr = sign_num_attr ? sign_num_attr
@@ -437,6 +436,18 @@ static void get_statuscol_str(win_T *wp, linenr_T lnum, int row, int startrow, i
   }
   stcp->sign_text[i] = NULL;
 
+  // When a buffer's line count has changed, make a best estimate for the full
+  // width of the status column by building with "w_nrwidth_line_count". Add
+  // potentially truncated width and rebuild before drawing anything.
+  if (wp->w_statuscol_line_count != wp->w_nrwidth_line_count) {
+    wp->w_statuscol_line_count = wp->w_nrwidth_line_count;
+    set_vim_var_nr(VV_VIRTNUM, 0);
+    build_statuscol_str(wp, wp->w_nrwidth_line_count, 0, stcp->width,
+                        ' ', stcp->text, &stcp->hlrec, stcp);
+    stcp->width += stcp->truncate;
+  }
+  set_vim_var_nr(VV_VIRTNUM, virtnum);
+
   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
diff --git a/src/nvim/screen.c b/src/nvim/screen.c
index ebff52cd69..05da6e0ef1 100644
--- a/src/nvim/screen.c
+++ b/src/nvim/screen.c
@@ -768,7 +768,6 @@ void comp_col(void)
 /// Otherwise it depends on 'numberwidth' and the line count.
 int number_width(win_T *wp)
 {
-  int n;
   linenr_T lnum;
 
   if (wp->w_p_rnu && !wp->w_p_nu) {
@@ -784,17 +783,13 @@ int number_width(win_T *wp)
   }
   wp->w_nrwidth_line_count = lnum;
 
-  // make best estimate for 'statuscolumn'
+  // reset for 'statuscolumn'
   if (*wp->w_p_stc != NUL) {
-    char buf[MAXPATHL];
-    wp->w_nrwidth_width = 0;
-    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);
+    wp->w_nrwidth_width = (wp->w_p_nu || wp->w_p_rnu) * (int)wp->w_p_nuw;
     return wp->w_nrwidth_width;
   }
 
-  n = 0;
+  int n = 0;
   do {
     lnum /= 10;
     n++;
diff --git a/test/functional/ui/statuscolumn_spec.lua b/test/functional/ui/statuscolumn_spec.lua
index ae3b95fb0f..3233e6cd19 100644
--- a/test/functional/ui/statuscolumn_spec.lua
+++ b/test/functional/ui/statuscolumn_spec.lua
@@ -439,7 +439,7 @@ describe('statuscolumn', function()
       vim.api.nvim_buf_set_extmark(0, ns, 7, 0, {
         virt_lines_leftcol = true, virt_lines = {{{"virt", ""}}} })
     ]])
-    feed('lh')  -- force update wcol/row
+    feed('lh')  -- force update cursor row
     screen:expect([[
                 4 aaaaa                                    |
                 5 aaaaa                                    |
@@ -458,5 +458,24 @@ describe('statuscolumn', function()
     ]])
     command('set stc=')  -- also for the default sign column
     screen:expect_unchanged()
+    -- 'statuscolumn' is not too wide with custom (bogus) fold column
+    command([[set stc=%{foldlevel(v:lnum)>0?repeat('-',foldlevel(v:lnum)):''}%=%l\ ]])
+    feed('Gd10Ggg<C-l>')
+    screen:expect([[
+               1 ^aaaaa                                     |
+               2 aaaaa                                     |
+               3 aaaaa                                     |
+               4 aaaaa                                     |
+               5 aaaaa                                     |
+               6 aaaaa                                     |
+               7 aaaaa                                     |
+      virt                                                 |
+      ---------8 aaaaa                                     |
+      virt                                                 |
+      ---------9 aaaaa                                     |
+      ~                                                    |
+      ~                                                    |
+                                                           |
+    ]])
   end)
 end)