feat(folds): support virtual text format for 'foldtext' (#25209)

Co-authored-by: Lewis Russell <lewis6991@gmail.com>

7 files changed, 214 insertions, 16 deletions

diff --git a/runtime/doc/fold.txt b/runtime/doc/fold.txt
index 24d7605e4a..8f7393f5e3 100644
--- a/runtime/doc/fold.txt
+++ b/runtime/doc/fold.txt
@@ -520,8 +520,10 @@ expression.  It can use these special Vim variables:
 			foldlevel.
 	v:foldlevel	the foldlevel of the fold
 
-In the result a TAB is replaced with a space and unprintable characters are
-made into printable characters.
+If the result is a |List|, it is parsed and drawn like "overlay" virtual text
+(see |nvim_buf_set_extmark()|), otherwise the result is converted to a string
+where a TAB is replaced with a space and unprintable characters are made into
+printable characters.
 
 The resulting line is truncated to fit in the window, it never wraps.
 When there is room after the text, it is filled with the character specified
diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 0162a99619..7c971097fb 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -89,6 +89,8 @@ The following new APIs and features were added.
 
 • Added inline virtual text support to |nvim_buf_set_extmark()|.
 
+• 'foldtext' now supports virtual text format. |fold-foldtext|
+
 • The terminal buffer now supports reflow (wrapped lines adapt when the buffer
   is resized horizontally). Note: Lines that are not visible and kept in
   |'scrollback'| are not reflown.
diff --git a/src/nvim/api/extmark.c b/src/nvim/api/extmark.c
index 05f62f6c7c..faab6e593c 100644
--- a/src/nvim/api/extmark.c
+++ b/src/nvim/api/extmark.c
@@ -1206,7 +1206,9 @@ VirtText parse_virt_text(Array chunks, Error *err, int *width)
     kv_push(virt_text, ((VirtTextChunk){ .text = text, .hl_id = hl_id }));
   }
 
-  *width = w;
+  if (width != NULL) {
+    *width = w;
+  }
   return virt_text;
 
 free_exit:
diff --git a/src/nvim/drawline.c b/src/nvim/drawline.c
index e1550e0ece..811cfc1eb2 100644
--- a/src/nvim/drawline.c
+++ b/src/nvim/drawline.c
@@ -1145,6 +1145,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool number_onl
   bool decor_need_recheck = false;      // call decor_recheck_draw_col() at next char
 
   char buf_fold[FOLD_TEXT_LEN];         // Hold value returned by get_foldtext
+  VirtText fold_vt = VIRTTEXT_EMPTY;
 
   // 'cursorlineopt' has "screenline" and cursor is in this line
   bool cul_screenline = false;
@@ -1916,7 +1917,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool number_onl
     if (draw_folded && wlv.n_extra == 0 && wlv.col == win_col_offset) {
       linenr_T lnume = lnum + foldinfo.fi_lines - 1;
       memset(buf_fold, ' ', FOLD_TEXT_LEN);
-      wlv.p_extra = get_foldtext(wp, lnum, lnume, foldinfo, buf_fold);
+      wlv.p_extra = get_foldtext(wp, lnum, lnume, foldinfo, buf_fold, &fold_vt);
       wlv.n_extra = (int)strlen(wlv.p_extra);
 
       if (wlv.p_extra != buf_fold) {
@@ -2881,6 +2882,9 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool number_onl
         }
       }
 
+      if (kv_size(fold_vt) > 0) {
+        draw_virt_text_item(buf, win_col_offset, fold_vt, kHlModeCombine, grid->cols, 0);
+      }
       draw_virt_text(wp, buf, win_col_offset, &wlv.col, grid->cols, wlv.row);
       grid_put_linebuf(grid, wlv.row, 0, wlv.col, grid->cols, wp->w_p_rl, wp, bg_attr, false);
       wlv.row++;
@@ -3207,6 +3211,7 @@ int win_line(win_T *wp, linenr_T lnum, int startrow, int endrow, bool number_onl
     }
   }     // for every character in the line
 
+  clear_virttext(&fold_vt);
   kv_destroy(virt_lines);
   xfree(wlv.p_extra_free);
   xfree(wlv.saved_p_extra_free);
diff --git a/src/nvim/eval.c b/src/nvim/eval.c
index ec90803884..a279b6d051 100644
--- a/src/nvim/eval.c
+++ b/src/nvim/eval.c
@@ -12,7 +12,9 @@
 #include <string.h>
 
 #include "auto/config.h"
+#include "nvim/api/private/converter.h"
 #include "nvim/api/private/defs.h"
+#include "nvim/api/private/helpers.h"
 #include "nvim/ascii.h"
 #include "nvim/buffer.h"
 #include "nvim/buffer_defs.h"
@@ -994,7 +996,7 @@ char *eval_to_string(char *arg, bool convert)
 /// textlock.
 ///
 /// @param use_sandbox  when true, use the sandbox.
-char *eval_to_string_safe(char *arg, int use_sandbox)
+char *eval_to_string_safe(char *arg, const bool use_sandbox)
 {
   char *retval;
   funccal_entry_T funccal_entry;
@@ -1267,11 +1269,10 @@ void *call_func_retlist(const char *func, int argc, typval_T *argv)
 
 /// Evaluate 'foldexpr'.  Returns the foldlevel, and any character preceding
 /// it in "*cp".  Doesn't give error messages.
-int eval_foldexpr(char *arg, int *cp)
+int eval_foldexpr(win_T *wp, int *cp)
 {
-  typval_T tv;
-  varnumber_T retval;
-  int use_sandbox = was_set_insecurely(curwin, "foldexpr", OPT_LOCAL);
+  const bool use_sandbox = was_set_insecurely(wp, "foldexpr", OPT_LOCAL);
+  char *arg = wp->w_p_fde;
 
   emsg_off++;
   if (use_sandbox) {
@@ -1279,6 +1280,9 @@ int eval_foldexpr(char *arg, int *cp)
   }
   textlock++;
   *cp = NUL;
+
+  typval_T tv;
+  varnumber_T retval;
   if (eval0(arg, &tv, NULL, &EVALARG_EVALUATE) == FAIL) {
     retval = 0;
   } else {
@@ -1298,6 +1302,7 @@ int eval_foldexpr(char *arg, int *cp)
     }
     tv_clear(&tv);
   }
+
   emsg_off--;
   if (use_sandbox) {
     sandbox--;
@@ -1308,6 +1313,42 @@ int eval_foldexpr(char *arg, int *cp)
   return (int)retval;
 }
 
+/// Evaluate 'foldtext', returning an Array or a String (NULL_STRING on failure).
+Object eval_foldtext(win_T *wp)
+{
+  const bool use_sandbox = was_set_insecurely(wp, "foldtext", OPT_LOCAL);
+  char *arg = wp->w_p_fdt;
+  funccal_entry_T funccal_entry;
+
+  save_funccal(&funccal_entry);
+  if (use_sandbox) {
+    sandbox++;
+  }
+  textlock++;
+
+  typval_T tv;
+  Object retval;
+  if (eval0(arg, &tv, NULL, &EVALARG_EVALUATE) == FAIL) {
+    retval = STRING_OBJ(NULL_STRING);
+  } else {
+    if (tv.v_type == VAR_LIST) {
+      retval = vim_to_object(&tv);
+    } else {
+      retval = STRING_OBJ(cstr_to_string(tv_get_string(&tv)));
+    }
+    tv_clear(&tv);
+  }
+  clear_evalarg(&EVALARG_EVALUATE, NULL);
+
+  if (use_sandbox) {
+    sandbox--;
+  }
+  textlock--;
+  restore_funccal();
+
+  return retval;
+}
+
 /// Get an lvalue
 ///
 /// Lvalue may be
diff --git a/src/nvim/fold.c b/src/nvim/fold.c
index a6cb0b568c..1d5ba49301 100644
--- a/src/nvim/fold.c
+++ b/src/nvim/fold.c
@@ -12,12 +12,14 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "nvim/api/extmark.h"
 #include "nvim/ascii.h"
 #include "nvim/buffer_defs.h"
 #include "nvim/buffer_updates.h"
 #include "nvim/change.h"
 #include "nvim/charset.h"
 #include "nvim/cursor.h"
+#include "nvim/decoration.h"
 #include "nvim/diff.h"
 #include "nvim/drawscreen.h"
 #include "nvim/eval.h"
@@ -1702,8 +1704,9 @@ static void foldDelMarker(buf_T *buf, linenr_T lnum, char *marker, size_t marker
 /// @return the text for a closed fold
 ///
 /// Otherwise the result is in allocated memory.
-char *get_foldtext(win_T *wp, linenr_T lnum, linenr_T lnume, foldinfo_T foldinfo, char *buf)
-  FUNC_ATTR_NONNULL_ARG(1)
+char *get_foldtext(win_T *wp, linenr_T lnum, linenr_T lnume, foldinfo_T foldinfo, char *buf,
+                   VirtText *vt)
+  FUNC_ATTR_NONNULL_ALL
 {
   char *text = NULL;
   // an error occurred when evaluating 'fdt' setting
@@ -1750,8 +1753,22 @@ char *get_foldtext(win_T *wp, linenr_T lnum, linenr_T lnume, foldinfo_T foldinfo
       current_sctx = wp->w_p_script_ctx[WV_FDT].script_ctx;
 
       emsg_off++;  // handle exceptions, but don't display errors
-      text = eval_to_string_safe(wp->w_p_fdt,
-                                 was_set_insecurely(wp, "foldtext", OPT_LOCAL));
+
+      Object obj = eval_foldtext(wp);
+      if (obj.type == kObjectTypeArray) {
+        Error err = ERROR_INIT;
+        *vt = parse_virt_text(obj.data.array, &err, NULL);
+        if (!ERROR_SET(&err)) {
+          *buf = NUL;
+          text = buf;
+        }
+        api_clear_error(&err);
+      } else if (obj.type == kObjectTypeString) {
+        text = obj.data.string.data;
+        obj = NIL;
+      }
+      api_free_object(obj);
+
       emsg_off--;
 
       if (text == NULL || did_emsg) {
@@ -2929,7 +2946,7 @@ static void foldlevelExpr(fline_T *flp)
   const bool save_keytyped = KeyTyped;
 
   int c;
-  const int n = eval_foldexpr(flp->wp->w_p_fde, &c);
+  const int n = eval_foldexpr(flp->wp, &c);
   KeyTyped = save_keytyped;
 
   switch (c) {
@@ -3320,10 +3337,20 @@ void f_foldtextresult(typval_T *argvars, typval_T *rettv, EvalFuncData fptr)
 
   foldinfo_T info = fold_info(curwin, lnum);
   if (info.fi_lines > 0) {
-    char *text = get_foldtext(curwin, lnum, lnum + info.fi_lines - 1, info, buf);
+    VirtText vt = VIRTTEXT_EMPTY;
+    char *text = get_foldtext(curwin, lnum, lnum + info.fi_lines - 1, info, buf, &vt);
     if (text == buf) {
       text = xstrdup(text);
     }
+    if (kv_size(vt) > 0) {
+      assert(*text == NUL);
+      for (size_t i = 0; i < kv_size(vt); i++) {
+        char *new_text = concat_str(text, kv_A(vt, i).text);
+        xfree(text);
+        text = new_text;
+      }
+    }
+    clear_virttext(&vt);
     rettv->vval.v_string = text;
   }
 
diff --git a/test/functional/ui/fold_spec.lua b/test/functional/ui/fold_spec.lua
index f00fba331e..46de6b114e 100644
--- a/test/functional/ui/fold_spec.lua
+++ b/test/functional/ui/fold_spec.lua
@@ -40,12 +40,16 @@ describe("folded lines", function()
         [8] = {foreground = Screen.colors.Brown },
         [9] = {bold = true, foreground = Screen.colors.Brown},
         [10] = {background = Screen.colors.LightGrey, underline = true},
-        [11] = {bold=true},
+        [11] = {bold = true},
         [12] = {foreground = Screen.colors.Red},
         [13] = {foreground = Screen.colors.Red, background = Screen.colors.LightGrey},
         [14] = {background = Screen.colors.Red},
         [15] = {foreground = Screen.colors.DarkBlue, background = Screen.colors.Red},
         [16] = {background = Screen.colors.LightGrey},
+        [17] = {background = Screen.colors.Yellow, foreground = Screen.colors.Red},
+        [18] = {background = Screen.colors.LightGrey, bold = true, foreground = Screen.colors.Blue},
+        [19] = {background = Screen.colors.Yellow, foreground = Screen.colors.DarkBlue},
+        [20] = {background = Screen.colors.Red, bold = true, foreground = Screen.colors.Blue},
       })
     end)
 
@@ -2816,6 +2820,121 @@ describe("folded lines", function()
         ]])
       end
     end)
+
+    it('support foldtext with virtual text format', function()
+      screen:try_resize(30, 7)
+      insert(content1)
+      command("hi! CursorLine guibg=NONE guifg=Red gui=NONE")
+      meths.set_option_value('cursorline', true, {})
+      meths.set_option_value('foldcolumn', '4', {})
+      meths.set_option_value('foldtext',
+        '[[v:folddashes], ["\t", "Search"], [getline(v:foldstart), "NonText"]]', {})
+
+      command('3,4fold')
+      command('5,6fold')
+      command('2,6fold')
+      if multigrid then
+        screen:expect{grid=[[
+        ## grid 1
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [3:------------------------------]|
+        ## grid 2
+          {7:    }This is a                 |
+          {7:+   }{13:^-}{17:       }{18:valid English}{13:·····}|
+          {1:~                             }|
+          {1:~                             }|
+          {1:~                             }|
+          {1:~                             }|
+        ## grid 3
+                                        |
+        ]]}
+      else
+        screen:expect([[
+          {7:    }This is a                 |
+          {7:+   }{13:^-}{17:       }{18:valid English}{13:·····}|
+          {1:~                             }|
+          {1:~                             }|
+          {1:~                             }|
+          {1:~                             }|
+                                        |
+        ]])
+      end
+      eq('-\tvalid English', funcs.foldtextresult(2))
+
+      feed('zo')
+      if multigrid then
+        screen:expect{grid=[[
+        ## grid 1
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [3:------------------------------]|
+        ## grid 2
+          {7:    }This is a                 |
+          {7:-   }valid English             |
+          {7:│+  }{5:--}{19:      }{18:sentence composed }|
+          {7:│+  }{13:^--}{17:      }{18:in his cave.}{13:······}|
+          {1:~                             }|
+          {1:~                             }|
+        ## grid 3
+                                        |
+        ]]}
+      else
+        screen:expect([[
+          {7:    }This is a                 |
+          {7:-   }valid English             |
+          {7:│+  }{5:--}{19:      }{18:sentence composed }|
+          {7:│+  }{13:^--}{17:      }{18:in his cave.}{13:······}|
+          {1:~                             }|
+          {1:~                             }|
+                                        |
+        ]])
+      end
+      eq('--\tsentence composed by', funcs.foldtextresult(3))
+      eq('--\tin his cave.', funcs.foldtextresult(5))
+
+      command('hi! Visual guibg=Red')
+      feed('V2k')
+      if multigrid then
+        screen:expect{grid=[[
+        ## grid 1
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [2:------------------------------]|
+          [3:------------------------------]|
+        ## grid 2
+          {7:    }This is a                 |
+          {7:-   }^v{14:alid English}             |
+          {7:│+  }{15:--}{19:      }{20:sentence composed }|
+          {7:│+  }{15:--}{19:      }{20:in his cave.}{15:······}|
+          {1:~                             }|
+          {1:~                             }|
+        ## grid 3
+          {11:-- VISUAL LINE --}             |
+        ]]}
+      else
+        screen:expect([[
+          {7:    }This is a                 |
+          {7:-   }^v{14:alid English}             |
+          {7:│+  }{15:--}{19:      }{20:sentence composed }|
+          {7:│+  }{15:--}{19:      }{20:in his cave.}{15:······}|
+          {1:~                             }|
+          {1:~                             }|
+          {11:-- VISUAL LINE --}             |
+        ]])
+      end
+    end)
   end
 
   describe("with ext_multigrid", function()