From 9ce1623837a817c3f4f5deff9c8ba862578b6009 Mon Sep 17 00:00:00 2001 From: Till Bungert Date: Sun, 1 Oct 2023 21:10:51 +0200 Subject: [PATCH] feat(treesitter): add foldtext with treesitter highlighting (#25391) --- runtime/doc/news.txt | 2 + runtime/doc/treesitter.txt | 10 ++ runtime/lua/vim/treesitter.lua | 12 ++ runtime/lua/vim/treesitter/_fold.lua | 92 ++++++++++++ test/functional/treesitter/fold_spec.lua | 172 +++++++++++++++++++++++ 5 files changed, 288 insertions(+) diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt index f82cb7b7e0..db0c7b4407 100644 --- a/runtime/doc/news.txt +++ b/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.) diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt index 34971c7acf..e19fda8fd1 100644 --- a/runtime/doc/treesitter.txt +++ b/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 diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua index 0e34cbcbcc..f863942d3b 100644 --- a/runtime/lua/vim/treesitter.lua +++ b/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 diff --git a/runtime/lua/vim/treesitter/_fold.lua b/runtime/lua/vim/treesitter/_fold.lua index 8bc08c9c2e..c6a4b48d4f 100644 --- a/runtime/lua/vim/treesitter/_fold.lua +++ b/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 diff --git a/test/functional/treesitter/fold_spec.lua b/test/functional/treesitter/fold_spec.lua index 9ed86e87f1..8cf9a91bbd 100644 --- a/test/functional/treesitter/fold_spec.lua +++ b/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)