fix(extmarks): problems with folded virtual lines (#21930)

Problem:    When a folded line has virtual lines attached, the following
            problems occur:
              - The virtual lines are drawn empty.
              - The 'foldtext' line is drawn empty.
              - The cursor is drawn incorrectly.
Solution:   Check whether virtual lines belong to a folded line.

Fix #17027
Fix #19557
Fix #21837

Co-authored-by: zeertzjq <zeertzjq@outlook.com>

6 files changed, 139 insertions, 11 deletions

diff --git a/src/nvim/decoration.c b/src/nvim/decoration.c
index 9274138cb6..63c55ec602 100644
--- a/src/nvim/decoration.c
+++ b/src/nvim/decoration.c
@@ -7,6 +7,7 @@
 #include "nvim/decoration.h"
 #include "nvim/drawscreen.h"
 #include "nvim/extmark.h"
+#include "nvim/fold.h"
 #include "nvim/highlight.h"
 #include "nvim/highlight_group.h"
 #include "nvim/memory.h"
@@ -550,7 +551,8 @@ void decor_add_ephemeral(int start_row, int start_col, int end_row, int end_col,
   decor_add(&decor_state, start_row, start_col, end_row, end_col, decor, true, ns_id, mark_id);
 }
 
-int decor_virt_lines(win_T *wp, linenr_T lnum, VirtLines *lines)
+/// @param has_fold  whether line "lnum" has a fold, or kNone when not calculated yet
+int decor_virt_lines(win_T *wp, linenr_T lnum, VirtLines *lines, TriState has_fold)
 {
   buf_T *buf = wp->w_buffer;
   if (!buf->b_virt_line_blocks) {
@@ -564,6 +566,10 @@ int decor_virt_lines(win_T *wp, linenr_T lnum, VirtLines *lines)
   int end_row = (int)lnum;
   MarkTreeIter itr[1] = { 0 };
   marktree_itr_get(buf->b_marktree, row, 0,  itr);
+  bool below_fold = lnum > 1 && hasFoldingWin(wp, lnum - 1, NULL, NULL, true, NULL);
+  if (has_fold == kNone) {
+    has_fold = hasFoldingWin(wp, lnum, NULL, NULL, true, NULL);
+  }
   while (true) {
     mtkey_t mark = marktree_itr_current(itr);
     if (mark.pos.row < 0 || mark.pos.row >= end_row) {
@@ -572,8 +578,9 @@ int decor_virt_lines(win_T *wp, linenr_T lnum, VirtLines *lines)
       goto next_mark;
     }
     bool above = mark.pos.row > (lnum - 2);
+    bool has_fold_cur = above ? has_fold : below_fold;
     Decoration *decor = mark.decor_full;
-    if (decor && decor->virt_lines_above == above) {
+    if (!has_fold_cur && decor && decor->virt_lines_above == above) {
       virt_lines += (int)kv_size(decor->virt_lines);
       if (lines) {
         kv_splice(*lines, decor->virt_lines);
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index e7efb14559..e24d86b353 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -940,7 +940,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool nochange,
     area_highlighting = true;
   }
   VirtLines virt_lines = KV_INITIAL_VALUE;
-  int n_virt_lines = decor_virt_lines(wp, lnum, &virt_lines);
+  int n_virt_lines = decor_virt_lines(wp, lnum, &virt_lines, has_fold);
   filler_lines += n_virt_lines;
   if (lnum == wp->w_topline) {
     filler_lines = wp->w_topfill;
@@ -1507,7 +1507,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool nochange,
         && has_fold
         && col == win_col_offset
         && n_extra == 0
-        && row == startrow) {
+        && row == startrow + filler_lines) {
       char_attr = win_hl_attr(wp, HLF_FL);
 
       linenr_T lnume = lnum + foldinfo.fi_lines - 1;
@@ -1528,7 +1528,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool nochange,
         && has_fold
         && col < grid->cols
         && n_extra == 0
-        && row == startrow) {
+        && row == startrow + filler_lines) {
       // fill rest of line with 'fold'
       c_extra = wp->w_p_fcs_chars.fold;
       c_final = NUL;
@@ -1540,7 +1540,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool nochange,
         && has_fold
         && col >= grid->cols
         && n_extra != 0
-        && row == startrow) {
+        && row == startrow + filler_lines) {
       // Truncate the folding.
       n_extra = 0;
     }
@@ -2753,7 +2753,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool nochange,
     // At end of screen line and there is more to come: Display the line
     // so far.  If there is no more to display it is caught above.
     if ((wp->w_p_rl ? (col < 0) : (col >= grid->cols))
-        && foldinfo.fi_lines == 0
+        && (!has_fold || virt_line_offset >= 0)
         && (draw_state != WL_LINE
             || *ptr != NUL
             || filler_todo > 0
diff --git a/src/nvim/drawscreen.c b/src/nvim/drawscreen.c
index 568dbe99c0..3f81fd1daa 100644
--- a/src/nvim/drawscreen.c
+++ b/src/nvim/drawscreen.c
@@ -1906,7 +1906,7 @@ static void win_update(win_T *wp, DecorProviders *providers)
       if (j > 0 && !wp->w_botfill && row < wp->w_grid.rows) {
         // Display filler text below last line. win_line() will check
         // for ml_line_count+1 and only draw filler lines
-        foldinfo_T info = FOLDINFO_INIT;
+        foldinfo_T info = { 0 };
         row = win_line(wp, wp->w_botline, row, wp->w_grid.rows,
                        false, false, info, &line_providers, &provider_err);
       }
diff --git a/src/nvim/fold.h b/src/nvim/fold.h
index 9c3a34ab31..ac1e8c9419 100644
--- a/src/nvim/fold.h
+++ b/src/nvim/fold.h
@@ -20,8 +20,6 @@ typedef struct foldinfo {
   linenr_T fi_lines;
 } foldinfo_T;
 
-#define FOLDINFO_INIT { 0, 0, 0, 0 }
-
 EXTERN int disable_fold_update INIT(= 0);
 
 #ifdef INCLUDE_GENERATED_DECLARATIONS
diff --git a/src/nvim/plines.c b/src/nvim/plines.c
index 6393c8dea5..5469d94800 100644
--- a/src/nvim/plines.c
+++ b/src/nvim/plines.c
@@ -51,7 +51,7 @@ int plines_win(win_T *wp, linenr_T lnum, bool winheight)
 /// @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, NULL);
+  int virt_lines = decor_virt_lines(wp, lnum, NULL, kNone);
 
   // be quick when there are no filler lines
   if (diffopt_filler()) {
diff --git a/test/functional/ui/fold_spec.lua b/test/functional/ui/fold_spec.lua
index bfa7167100..46a478c1ea 100644
--- a/test/functional/ui/fold_spec.lua
+++ b/test/functional/ui/fold_spec.lua
@@ -8,6 +8,7 @@ local expect = helpers.expect
 local funcs = helpers.funcs
 local meths = helpers.meths
 local exec = helpers.exec
+local exec_lua = helpers.exec_lua
 local assert_alive = helpers.assert_alive
 
 
@@ -1852,6 +1853,128 @@ describe("folded lines", function()
         ]])
       end
     end)
+
+    it('fold attached virtual lines are drawn correctly #21837', function()
+      funcs.setline(1, 'line 1')
+      funcs.setline(2, 'line 2')
+      funcs.setline(3, 'line 3')
+      funcs.setline(4, 'line 4')
+      feed("zfj")
+      exec_lua([[
+        local ns = vim.api.nvim_create_namespace("ns")
+        vim.api.nvim_buf_set_extmark(0, ns, 0, 0, { virt_lines_above = true, virt_lines = {{{"virt_line above line 1", ""}}} })
+        vim.api.nvim_buf_set_extmark(0, ns, 1, 0, { virt_lines = {{{"virt_line below line 2", ""}}} })
+        vim.api.nvim_buf_set_extmark(0, ns, 2, 0, { virt_lines_above = true, virt_lines = {{{"virt_line above line 3", ""}}} })
+        vim.api.nvim_buf_set_extmark(0, ns, 3, 0, { virt_lines = {{{"virt_line below line 4", ""}}} })
+      ]])
+      if multigrid then
+        screen:expect([[
+        ## grid 1
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [3:---------------------------------------------]|
+        ## grid 2
+          {5:^+--  2 lines: line 1·························}|
+          virt_line above line 3                       |
+          line 3                                       |
+          line 4                                       |
+          virt_line below line 4                       |
+          {1:~                                            }|
+          {1:~                                            }|
+        ## grid 3
+                                                       |
+        ]])
+      else
+        screen:expect([[
+          {5:^+--  2 lines: line 1·························}|
+          virt_line above line 3                       |
+          line 3                                       |
+          line 4                                       |
+          virt_line below line 4                       |
+          {1:~                                            }|
+          {1:~                                            }|
+                                                       |
+        ]])
+      end
+
+      feed('jzfj')
+      if multigrid then
+        screen:expect([[
+        ## grid 1
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [3:---------------------------------------------]|
+        ## grid 2
+          {5:+--  2 lines: line 1·························}|
+          {5:^+--  2 lines: line 3·························}|
+          {1:~                                            }|
+          {1:~                                            }|
+          {1:~                                            }|
+          {1:~                                            }|
+          {1:~                                            }|
+        ## grid 3
+                                                       |
+        ]])
+      else
+        screen:expect([[
+          {5:+--  2 lines: line 1·························}|
+          {5:^+--  2 lines: line 3·························}|
+          {1:~                                            }|
+          {1:~                                            }|
+          {1:~                                            }|
+          {1:~                                            }|
+          {1:~                                            }|
+                                                       |
+        ]])
+      end
+
+      feed('kzo<C-Y>')
+      funcs.setline(5, 'line 5')
+      if multigrid then
+        screen:expect([[
+        ## grid 1
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [2:---------------------------------------------]|
+          [3:---------------------------------------------]|
+        ## grid 2
+          virt_line above line 1                       |
+          ^line 1                                       |
+          line 2                                       |
+          virt_line below line 2                       |
+          {5:+--  2 lines: line 3·························}|
+          line 5                                       |
+          {1:~                                            }|
+        ## grid 3
+                                                       |
+        ]])
+      else
+        screen:expect([[
+          virt_line above line 1                       |
+          ^line 1                                       |
+          line 2                                       |
+          virt_line below line 2                       |
+          {5:+--  2 lines: line 3·························}|
+          line 5                                       |
+          {1:~                                            }|
+                                                       |
+        ]])
+      end
+    end)
   end
 
   describe("with ext_multigrid", function()