feat(treesitter): add foldtext with treesitter highlighting (#25391)

runtime/doc/news.txt

@@ -145,6 +145,8 @@ The following new APIs and features were added.
   • Added `vim.treesitter.query.edit()`, for live editing of treesitter
     queries.
   • Improved error messages for query parsing.
+  • Added |vim.treesitter.foldtext()| to apply treesitter highlighting to
+    foldtext.
 
 • |vim.ui.open()| opens URIs using the system default handler (macOS `open`,
   Windows `explorer`, Linux `xdg-open`, etc.)

runtime/doc/treesitter.txt

@@ -560,6 +560,16 @@ foldexpr({lnum})                                   *vim.treesitter.foldexpr()*
     Return: ~
         (string)
 
+foldtext()                                         *vim.treesitter.foldtext()*
+    Returns the highlighted content of the first line of the fold or falls
+    back to |foldtext()| if no treesitter parser is found. Can be set directly
+    to 'foldtext': >lua
+        vim.wo.foldtext = 'v:lua.vim.treesitter.foldtext()'
+<
+
+    Return: ~
+        `{ [1]: string, [2]: string[] }[]` | string
+
                                      *vim.treesitter.get_captures_at_cursor()*
 get_captures_at_cursor({winnr})
     Returns a list of highlight capture names under the cursor

runtime/lua/vim/treesitter.lua

@@ -508,4 +508,16 @@ function M.foldexpr(lnum)
   return require('vim.treesitter._fold').foldexpr(lnum)
 end
 
+--- Returns the highlighted content of the first line of the fold or falls back to |foldtext()|
+--- if no treesitter parser is found. Can be set directly to 'foldtext':
+---
+--- ```lua
+--- vim.wo.foldtext = 'v:lua.vim.treesitter.foldtext()'
+--- ```
+---
+---@return { [1]: string, [2]: string[] }[] | string
+function M.foldtext()
+  return require('vim.treesitter._fold').foldtext()
+end
+
 return M

runtime/lua/vim/treesitter/_fold.lua

@@ -361,4 +361,96 @@ function M.foldexpr(lnum)
   return foldinfos[bufnr].levels[lnum] or '0'
 end
 
+---@package
+---@return { [1]: string, [2]: string[] }[]|string
+function M.foldtext()
+  local foldstart = vim.v.foldstart
+  local bufnr = api.nvim_get_current_buf()
+
+  ---@type boolean, LanguageTree
+  local ok, parser = pcall(ts.get_parser, bufnr)
+  if not ok then
+    return vim.fn.foldtext()
+  end
+
+  local query = ts.query.get(parser:lang(), 'highlights')
+  if not query then
+    return vim.fn.foldtext()
+  end
+
+  local tree = parser:parse({ foldstart - 1, foldstart })[1]
+
+  local line = api.nvim_buf_get_lines(bufnr, foldstart - 1, foldstart, false)[1]
+  if not line then
+    return vim.fn.foldtext()
+  end
+
+  ---@type { [1]: string, [2]: string[], range: { [1]: integer, [2]: integer } }[] | { [1]: string, [2]: string[] }[]
+  local result = {}
+
+  local line_pos = 0
+
+  for id, node, metadata in query:iter_captures(tree:root(), 0, foldstart - 1, foldstart) do
+    local name = query.captures[id]
+    local start_row, start_col, end_row, end_col = node:range()
+
+    local priority = tonumber(metadata.priority or vim.highlight.priorities.treesitter)
+
+    if start_row == foldstart - 1 and end_row == foldstart - 1 then
+      -- check for characters ignored by treesitter
+      if start_col > line_pos then
+        table.insert(result, {
+          line:sub(line_pos + 1, start_col),
+          { { 'Folded', priority } },
+          range = { line_pos, start_col },
+        })
+      end
+      line_pos = end_col
+
+      local text = line:sub(start_col + 1, end_col)
+      table.insert(result, { text, { { '@' .. name, priority } }, range = { start_col, end_col } })
+    end
+  end
+
+  local i = 1
+  while i <= #result do
+    -- find first capture that is not in current range and apply highlights on the way
+    local j = i + 1
+    while
+      j <= #result
+      and result[j].range[1] >= result[i].range[1]
+      and result[j].range[2] <= result[i].range[2]
+    do
+      for k, v in ipairs(result[i][2]) do
+        if not vim.tbl_contains(result[j][2], v) then
+          table.insert(result[j][2], k, v)
+        end
+      end
+      j = j + 1
+    end
+
+    -- remove the parent capture if it is split into children
+    if j > i + 1 then
+      table.remove(result, i)
+    else
+      -- highlights need to be sorted by priority, on equal prio, the deeper nested capture (earlier
+      -- in list) should be considered higher prio
+      if #result[i][2] > 1 then
+        table.sort(result[i][2], function(a, b)
+          return a[2] < b[2]
+        end)
+      end
+
+      result[i][2] = vim.tbl_map(function(tbl)
+        return tbl[1]
+      end, result[i][2])
+      result[i] = { result[i][1], result[i][2] }
+
+      i = i + 1
+    end
+  end
+
+  return result
+end
+
 return M

test/functional/treesitter/fold_spec.lua

@@ -359,3 +359,175 @@ void ui_refresh(void)
   end)
 
 end)
+
+describe('treesitter foldtext', function()
+  local test_text = [[
+void qsort(void *base, size_t nel, size_t width, int (*compar)(const void *, const void *))
+{
+  int width = INT_MAX, height = INT_MAX;
+  bool ext_widgets[kUIExtCount];
+  for (UIExtension i = 0; (int)i < kUIExtCount; i++) {
+    ext_widgets[i] = true;
+  }
+
+  bool inclusive = ui_override();
+  for (size_t i = 0; i < ui_count; i++) {
+    UI *ui = uis[i];
+    width = MIN(ui->width, width);
+    height = MIN(ui->height, height);
+    foo = BAR(ui->bazaar, bazaar);
+    for (UIExtension j = 0; (int)j < kUIExtCount; j++) {
+      ext_widgets[j] &= (ui->ui_ext[j] || inclusive);
+    }
+  }
+}]]
+
+  it('displays highlighted content', function()
+    local screen = Screen.new(60, 21)
+    screen:attach()
+
+    command([[set foldmethod=manual foldtext=v:lua.vim.treesitter.foldtext() updatetime=50]])
+    insert(test_text)
+    exec_lua([[vim.treesitter.get_parser(0, "c")]])
+
+    feed('ggVGzf')
+
+    screen:expect({
+      grid = [[
+{1:^void}{2: }{3:qsort}{4:(}{1:void}{2: }{5:*}{3:base}{4:,}{2: }{1:size_t}{2: }{3:nel}{4:,}{2: }{1:size_t}{2: }{3:width}{4:,}{2: }{1:int}{2: }{4:(}{5:*}{3:compa}|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+{6:~                                                           }|
+                                                            |
+]],
+      attr_ids = {
+        [1] = {
+          foreground = Screen.colors.SeaGreen4,
+          background = Screen.colors.LightGrey,
+          bold = true,
+        },
+        [2] = { background = Screen.colors.LightGrey, foreground = Screen.colors.Blue4 },
+        [3] = { background = Screen.colors.LightGrey, foreground = Screen.colors.DarkCyan },
+        [4] = { background = Screen.colors.LightGrey, foreground = Screen.colors.SlateBlue },
+        [5] = {
+          foreground = Screen.colors.Brown,
+          background = Screen.colors.LightGrey,
+          bold = true,
+        },
+        [6] = { foreground = Screen.colors.Blue, bold = true },
+      },
+    })
+  end)
+
+  it('handles deep nested captures', function()
+    local screen = Screen.new(60, 21)
+    screen:attach()
+
+    command([[set foldmethod=manual foldtext=v:lua.vim.treesitter.foldtext() updatetime=50]])
+    insert([[
+function FoldInfo.new()
+  return setmetatable({
+    start_counts = {},
+    stop_counts = {},
+    levels0 = {},
+    levels = {},
+  }, FoldInfo)
+end
+    ]])
+    exec_lua([[vim.treesitter.get_parser(0, "lua")]])
+
+    feed('ggjVGkzf')
+
+    screen:expect({
+      grid = [[
+function FoldInfo.new()                                     |
+{1:^  }{2:return}{1: }{3:setmetatable({}{1:·····································}|
+                                                            |
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+{4:~                                                           }|
+                                                            |
+]],
+      attr_ids = {
+        [1] = { foreground = Screen.colors.Blue4, background = Screen.colors.LightGray },
+        [2] = {
+          foreground = Screen.colors.Brown,
+          bold = true,
+          background = Screen.colors.LightGray,
+        },
+        [3] = { foreground = Screen.colors.SlateBlue, background = Screen.colors.LightGray },
+        [4] = { bold = true, foreground = Screen.colors.Blue },
+      },
+    })
+  end)
+
+  it('falls back to default', function()
+    local screen = Screen.new(60, 21)
+    screen:attach()
+
+    command([[set foldmethod=manual foldtext=v:lua.vim.treesitter.foldtext()]])
+    insert(test_text)
+
+    feed('ggVGzf')
+
+    screen:expect({
+      grid = [[
+{1:^+-- 19 lines: void qsort(void *base, size_t nel, size_t widt}|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+{2:~                                                           }|
+                                                            |
+]],
+      attr_ids = {
+        [1] = { foreground = Screen.colors.Blue4, background = Screen.colors.LightGray },
+        [2] = { bold = true, foreground = Screen.colors.Blue },
+      },
+    })
+  end)
+end)