vim-patch:9.1.0079: LineNrAbove/Below highlighting wrong on wrapped lines (#27363)

Problem:  LineNrAbove and LineNrBelow background wrong on wrapped lines.
Solution: Update number column also for wrapped part of a line.
          (zeertzjq)

closes: vim/vim#13974

https://github.com/vim/vim/commit/ebfd856cfdf6ea0b16c8d115000961c998ce97da

Cherry-pick test_number.vim changes from patch 9.0.0626.

3 files changed, 261 insertions, 78 deletions

diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index b5c0755d44..174543242f 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -929,7 +929,8 @@ static int get_rightmost_vcol(win_T *wp, const int *color_cols)
 /// @param lnum         line to display
 /// @param startrow     first row relative to window grid
 /// @param endrow       last grid row to be redrawn
-/// @param col_rows     only update the columns for this amount of rows
+/// @param col_rows     set to the height of the line when only updating the columns,
+///                     otherwise set to 0
 /// @param spv          'spell' related variables kept between calls for "wp"
 /// @param foldinfo     fold info for this line
 /// @param[in, out] providers  decoration providers active this line
@@ -1556,6 +1557,36 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
 
       win_col_offset = wlv.off;
 
+      // When only updating the columns and that's done, stop here.
+      if (col_rows > 0) {
+        win_put_linebuf(wp, wlv.row, 0, wlv.off, wlv.off, bg_attr, false);
+        // Need to update more screen lines if:
+        // - 'statuscolumn' needs to be drawn, or
+        // - LineNrAbove or LineNrBelow is used, or
+        // - still drawing filler lines.
+        if ((wlv.row + 1 - wlv.startrow < col_rows
+             && (statuscol.draw
+                 || win_hl_attr(wp, HLF_LNA) != win_hl_attr(wp, HLF_N)
+                 || win_hl_attr(wp, HLF_LNB) != win_hl_attr(wp, HLF_N)))
+            || wlv.filler_todo > 0) {
+          wlv.row++;
+          if (wlv.row == endrow) {
+            break;
+          }
+          wlv.filler_todo--;
+          if (wlv.filler_todo == 0 && (wp->w_botfill || end_fill)) {
+            break;
+          }
+          // win_line_start(wp, &wlv);
+          wlv.col = 0;
+          wlv.off = 0;
+          draw_cols = true;
+          continue;
+        } else {
+          break;
+        }
+      }
+
       // Check if 'breakindent' applies and show it.
       if (!wp->w_briopt_sbr) {
         handle_breakindent(wp, &wlv);
@@ -1580,24 +1611,12 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, int col_rows, s
       apply_cursorline_highlight(wp, &wlv);
     }
 
-    // When still displaying '$' of change command, stop at cursor
-    if (((dollar_vcol >= 0
-          && wp == curwin
-          && lnum == wp->w_cursor.lnum
-          && wlv.vcol >= wp->w_virtcol)
-         || col_rows > 0)
-        && wlv.filler_todo <= 0) {
-      if (col_rows == 0) {
-        draw_virt_text(wp, buf, win_col_offset, &wlv.col, wlv.row);
-      }
+    // When still displaying '$' of change command, stop at cursor.
+    if (dollar_vcol >= 0 && wp == curwin
+        && lnum == wp->w_cursor.lnum && wlv.vcol >= wp->w_virtcol) {
+      draw_virt_text(wp, buf, win_col_offset, &wlv.col, wlv.row);
       // don't clear anything after wlv.col
       win_put_linebuf(wp, wlv.row, 0, wlv.col, wlv.col, bg_attr, false);
-      // update the 'statuscolumn' for the entire line size
-      if (col_rows > 0 && statuscol.draw && ++wlv.row - wlv.startrow < col_rows) {
-        draw_cols = true;
-        wlv.off = 0;
-        continue;
-      }
       // Pretend we have finished updating the window.  Except when
       // 'cursorcolumn' is set.
       if (wp->w_p_cuc) {
diff --git a/test/functional/ui/highlight_spec.lua b/test/functional/ui/highlight_spec.lua
index 2ec10d149f..d9bb4d3a98 100644
--- a/test/functional/ui/highlight_spec.lua
+++ b/test/functional/ui/highlight_spec.lua
@@ -1708,68 +1708,206 @@ end)
 describe("'number' and 'relativenumber' highlight", function()
   before_each(clear)
 
+  -- oldtest: Test_relativenumber_colors()
   it('LineNr, LineNrAbove and LineNrBelow', function()
-    local screen = Screen.new(20, 10)
+    local screen = Screen.new(50, 10)
     screen:set_default_attr_ids({
       [1] = { foreground = Screen.colors.Red },
       [2] = { foreground = Screen.colors.Blue },
       [3] = { foreground = Screen.colors.Green },
     })
     screen:attach()
-    command('set number relativenumber')
-    command('call setline(1, range(50))')
-    command('highlight LineNr guifg=Red')
-    feed('4j')
-    screen:expect([[
-      {1:  4 }0               |
-      {1:  3 }1               |
-      {1:  2 }2               |
-      {1:  1 }3               |
-      {1:5   }^4               |
-      {1:  1 }5               |
-      {1:  2 }6               |
-      {1:  3 }7               |
-      {1:  4 }8               |
-                          |
-    ]])
-    command('highlight LineNrAbove guifg=Blue')
-    screen:expect([[
-      {2:  4 }0               |
-      {2:  3 }1               |
-      {2:  2 }2               |
-      {2:  1 }3               |
-      {1:5   }^4               |
-      {1:  1 }5               |
-      {1:  2 }6               |
-      {1:  3 }7               |
-      {1:  4 }8               |
-                          |
-    ]])
-    command('highlight LineNrBelow guifg=Green')
-    screen:expect([[
-      {2:  4 }0               |
-      {2:  3 }1               |
-      {2:  2 }2               |
-      {2:  1 }3               |
-      {1:5   }^4               |
-      {3:  1 }5               |
-      {3:  2 }6               |
-      {3:  3 }7               |
-      {3:  4 }8               |
-                          |
-    ]])
-    feed('3j')
-    screen:expect([[
-      {2:  7 }0               |
-      {2:  6 }1               |
-      {2:  5 }2               |
-      {2:  4 }3               |
-      {2:  3 }4               |
-      {2:  2 }5               |
-      {2:  1 }6               |
-      {1:8   }^7               |
-      {3:  1 }8               |
-                          |
+    exec([[
+      call setline(1, range(200))
+      111
+      set number relativenumber
+      hi LineNr guifg=red
+    ]])
+    screen:expect([[
+      {1:  4 }106                                           |
+      {1:  3 }107                                           |
+      {1:  2 }108                                           |
+      {1:  1 }109                                           |
+      {1:111 }^110                                           |
+      {1:  1 }111                                           |
+      {1:  2 }112                                           |
+      {1:  3 }113                                           |
+      {1:  4 }114                                           |
+                                                        |
+    ]])
+    command('hi LineNrAbove guifg=blue')
+    screen:expect([[
+      {2:  4 }106                                           |
+      {2:  3 }107                                           |
+      {2:  2 }108                                           |
+      {2:  1 }109                                           |
+      {1:111 }^110                                           |
+      {1:  1 }111                                           |
+      {1:  2 }112                                           |
+      {1:  3 }113                                           |
+      {1:  4 }114                                           |
+                                                        |
+    ]])
+    command('hi LineNrBelow guifg=green')
+    screen:expect([[
+      {2:  4 }106                                           |
+      {2:  3 }107                                           |
+      {2:  2 }108                                           |
+      {2:  1 }109                                           |
+      {1:111 }^110                                           |
+      {3:  1 }111                                           |
+      {3:  2 }112                                           |
+      {3:  3 }113                                           |
+      {3:  4 }114                                           |
+                                                        |
+    ]])
+    command('hi clear LineNrAbove')
+    screen:expect([[
+      {1:  4 }106                                           |
+      {1:  3 }107                                           |
+      {1:  2 }108                                           |
+      {1:  1 }109                                           |
+      {1:111 }^110                                           |
+      {3:  1 }111                                           |
+      {3:  2 }112                                           |
+      {3:  3 }113                                           |
+      {3:  4 }114                                           |
+                                                        |
+    ]])
+  end)
+
+  -- oldtest: Test_relativenumber_colors_wrapped()
+  it('LineNr, LineNrAbove and LineNrBelow with wrapped lines', function()
+    local screen = Screen.new(50, 20)
+    screen:set_default_attr_ids({
+      [1] = { background = Screen.colors.Red, foreground = Screen.colors.Black },
+      [2] = { background = Screen.colors.Blue, foreground = Screen.colors.Black },
+      [3] = { background = Screen.colors.Green, foreground = Screen.colors.Black },
+      [4] = { bold = true, foreground = Screen.colors.Blue },
+    })
+    screen:attach()
+    exec([[
+      set display=lastline scrolloff=0
+      call setline(1, range(200)->map('v:val->string()->repeat(40)'))
+      111
+      set number relativenumber
+      hi LineNr guibg=red guifg=black
+      hi LineNrAbove guibg=blue guifg=black
+      hi LineNrBelow guibg=green guifg=black
+    ]])
+    screen:expect([[
+      {2:  2 }1081081081081081081081081081081081081081081081|
+      {2:    }0810810810810810810810810810810810810810810810|
+      {2:    }8108108108108108108108108108                  |
+      {2:  1 }1091091091091091091091091091091091091091091091|
+      {2:    }0910910910910910910910910910910910910910910910|
+      {2:    }9109109109109109109109109109                  |
+      {1:111 }^1101101101101101101101101101101101101101101101|
+      {1:    }1011011011011011011011011011011011011011011011|
+      {1:    }0110110110110110110110110110                  |
+      {3:  1 }1111111111111111111111111111111111111111111111|
+      {3:    }1111111111111111111111111111111111111111111111|
+      {3:    }1111111111111111111111111111                  |
+      {3:  2 }1121121121121121121121121121121121121121121121|
+      {3:    }1211211211211211211211211211211211211211211211|
+      {3:    }2112112112112112112112112112                  |
+      {3:  3 }1131131131131131131131131131131131131131131131|
+      {3:    }1311311311311311311311311311311311311311311311|
+      {3:    }3113113113113113113113113113                  |
+      {3:  4 }1141141141141141141141141141141141141141141{4:@@@}|
+                                                        |
+    ]])
+    feed('k')
+    screen:expect([[
+      {2:  1 }1081081081081081081081081081081081081081081081|
+      {2:    }0810810810810810810810810810810810810810810810|
+      {2:    }8108108108108108108108108108                  |
+      {1:110 }^1091091091091091091091091091091091091091091091|
+      {1:    }0910910910910910910910910910910910910910910910|
+      {1:    }9109109109109109109109109109                  |
+      {3:  1 }1101101101101101101101101101101101101101101101|
+      {3:    }1011011011011011011011011011011011011011011011|
+      {3:    }0110110110110110110110110110                  |
+      {3:  2 }1111111111111111111111111111111111111111111111|
+      {3:    }1111111111111111111111111111111111111111111111|
+      {3:    }1111111111111111111111111111                  |
+      {3:  3 }1121121121121121121121121121121121121121121121|
+      {3:    }1211211211211211211211211211211211211211211211|
+      {3:    }2112112112112112112112112112                  |
+      {3:  4 }1131131131131131131131131131131131131131131131|
+      {3:    }1311311311311311311311311311311311311311311311|
+      {3:    }3113113113113113113113113113                  |
+      {3:  5 }1141141141141141141141141141141141141141141{4:@@@}|
+                                                        |
+    ]])
+    feed('2j')
+    screen:expect([[
+      {2:  3 }1081081081081081081081081081081081081081081081|
+      {2:    }0810810810810810810810810810810810810810810810|
+      {2:    }8108108108108108108108108108                  |
+      {2:  2 }1091091091091091091091091091091091091091091091|
+      {2:    }0910910910910910910910910910910910910910910910|
+      {2:    }9109109109109109109109109109                  |
+      {2:  1 }1101101101101101101101101101101101101101101101|
+      {2:    }1011011011011011011011011011011011011011011011|
+      {2:    }0110110110110110110110110110                  |
+      {1:112 }^1111111111111111111111111111111111111111111111|
+      {1:    }1111111111111111111111111111111111111111111111|
+      {1:    }1111111111111111111111111111                  |
+      {3:  1 }1121121121121121121121121121121121121121121121|
+      {3:    }1211211211211211211211211211211211211211211211|
+      {3:    }2112112112112112112112112112                  |
+      {3:  2 }1131131131131131131131131131131131131131131131|
+      {3:    }1311311311311311311311311311311311311311311311|
+      {3:    }3113113113113113113113113113                  |
+      {3:  3 }1141141141141141141141141141141141141141141{4:@@@}|
+                                                        |
+    ]])
+    feed('2j')
+    screen:expect([[
+      {2:  5 }1081081081081081081081081081081081081081081081|
+      {2:    }0810810810810810810810810810810810810810810810|
+      {2:    }8108108108108108108108108108                  |
+      {2:  4 }1091091091091091091091091091091091091091091091|
+      {2:    }0910910910910910910910910910910910910910910910|
+      {2:    }9109109109109109109109109109                  |
+      {2:  3 }1101101101101101101101101101101101101101101101|
+      {2:    }1011011011011011011011011011011011011011011011|
+      {2:    }0110110110110110110110110110                  |
+      {2:  2 }1111111111111111111111111111111111111111111111|
+      {2:    }1111111111111111111111111111111111111111111111|
+      {2:    }1111111111111111111111111111                  |
+      {2:  1 }1121121121121121121121121121121121121121121121|
+      {2:    }1211211211211211211211211211211211211211211211|
+      {2:    }2112112112112112112112112112                  |
+      {1:114 }^1131131131131131131131131131131131131131131131|
+      {1:    }1311311311311311311311311311311311311311311311|
+      {1:    }3113113113113113113113113113                  |
+      {3:  1 }1141141141141141141141141141141141141141141{4:@@@}|
+                                                        |
+    ]])
+    feed('k')
+    screen:expect([[
+      {2:  4 }1081081081081081081081081081081081081081081081|
+      {2:    }0810810810810810810810810810810810810810810810|
+      {2:    }8108108108108108108108108108                  |
+      {2:  3 }1091091091091091091091091091091091091091091091|
+      {2:    }0910910910910910910910910910910910910910910910|
+      {2:    }9109109109109109109109109109                  |
+      {2:  2 }1101101101101101101101101101101101101101101101|
+      {2:    }1011011011011011011011011011011011011011011011|
+      {2:    }0110110110110110110110110110                  |
+      {2:  1 }1111111111111111111111111111111111111111111111|
+      {2:    }1111111111111111111111111111111111111111111111|
+      {2:    }1111111111111111111111111111                  |
+      {1:113 }^1121121121121121121121121121121121121121121121|
+      {1:    }1211211211211211211211211211211211211211211211|
+      {1:    }2112112112112112112112112112                  |
+      {3:  1 }1131131131131131131131131131131131131131131131|
+      {3:    }1311311311311311311311311311311311311311311311|
+      {3:    }3113113113113113113113113113                  |
+      {3:  2 }1141141141141141141141141141141141141141141{4:@@@}|
+                                                        |
     ]])
   end)
 
diff --git a/test/old/testdir/test_number.vim b/test/old/testdir/test_number.vim
index 1232db80c5..6ac3c4cfe4 100644
--- a/test/old/testdir/test_number.vim
+++ b/test/old/testdir/test_number.vim
@@ -276,11 +276,9 @@ func Test_relativenumber_colors()
     set number relativenumber
     hi LineNr ctermfg=red
   [CODE]
-  call writefile(lines, 'XTest_relnr')
+  call writefile(lines, 'XTest_relnr', 'D')
 
-  " Check that the balloon shows up after a mouse move
   let buf = RunVimInTerminal('-S XTest_relnr', {'rows': 10, 'cols': 50})
-  call TermWait(buf, 50)
   " Default colors
   call VerifyScreenDump(buf, 'Test_relnr_colors_1', {})
 
@@ -295,7 +293,36 @@ func Test_relativenumber_colors()
 
   " clean up
   call StopVimInTerminal(buf)
-  call delete('XTest_relnr')
+endfunc
+
+func Test_relativenumber_colors_wrapped()
+  CheckScreendump
+
+  let lines =<< trim [CODE]
+    set display=lastline scrolloff=0
+    call setline(1, range(200)->map('v:val->string()->repeat(40)'))
+    111
+    set number relativenumber
+    hi LineNr ctermbg=red ctermfg=black
+    hi LineNrAbove ctermbg=blue ctermfg=black
+    hi LineNrBelow ctermbg=green ctermfg=black
+  [CODE]
+  call writefile(lines, 'XTest_relnr_wrap', 'D')
+
+  let buf = RunVimInTerminal('-S XTest_relnr_wrap', {'rows': 20, 'cols': 50})
+
+  call VerifyScreenDump(buf, 'Test_relnr_colors_wrapped_1', {})
+  call term_sendkeys(buf, "k")
+  call VerifyScreenDump(buf, 'Test_relnr_colors_wrapped_2', {})
+  call term_sendkeys(buf, "2j")
+  call VerifyScreenDump(buf, 'Test_relnr_colors_wrapped_3', {})
+  call term_sendkeys(buf, "2j")
+  call VerifyScreenDump(buf, 'Test_relnr_colors_wrapped_4', {})
+  call term_sendkeys(buf, "k")
+  call VerifyScreenDump(buf, 'Test_relnr_colors_wrapped_5', {})
+
+  " clean up
+  call StopVimInTerminal(buf)
 endfunc
 
 func Test_relativenumber_callback()
@@ -313,14 +340,13 @@ func Test_relativenumber_callback()
 
       call timer_start(300, 'Func')
   END
-  call writefile(lines, 'Xrnu_timer')
+  call writefile(lines, 'Xrnu_timer', 'D')
 
   let buf = RunVimInTerminal('-S Xrnu_timer', #{rows: 8})
   call TermWait(buf, 310)
   call VerifyScreenDump(buf, 'Test_relativenumber_callback_1', {})
 
   call StopVimInTerminal(buf)
-  call delete('Xrnu_timer')
 endfunc
 
 " Test for displaying line numbers with 'rightleft'