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

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

4 files changed, 88 insertions, 13 deletions

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;
   }