feat(treesitter): table of contents for checkhealth, markdown (#32282)

Problem: It's difficult to navigate large structured text files (vim
help, checkhealth, Markdown).

Solution: Support `gO` for table of contents and `]]`/`[[` for moving
between headings for all these filetypes using treesitter queries.

Refactor: colorization of highlight groups is moved to the `help` ftplugin
while headings-related functionality is implemented in a private
`vim.treesitter` module for possible future use for other filetypes.

1 files changed, 45 insertions, 8 deletions

diff --git a/runtime/ftplugin/help.lua b/runtime/ftplugin/help.lua
index a6169a1d9d..f3bc28c552 100644
--- a/runtime/ftplugin/help.lua
+++ b/runtime/ftplugin/help.lua
@@ -1,15 +1,43 @@
 -- use treesitter over syntax (for highlighted code blocks)
 vim.treesitter.start()
 
+--- Apply current colorscheme to lists of default highlight groups
+---
+--- Note: {patterns} is assumed to be sorted by occurrence in the file.
+--- @param patterns {start:string,stop:string,match:string}[]
+local function colorize_hl_groups(patterns)
+  local ns = vim.api.nvim_create_namespace('nvim.vimhelp')
+  vim.api.nvim_buf_clear_namespace(0, ns, 0, -1)
+
+  local save_cursor = vim.fn.getcurpos()
+
+  for _, pat in pairs(patterns) do
+    local start_lnum = vim.fn.search(pat.start, 'c')
+    local end_lnum = vim.fn.search(pat.stop)
+    if start_lnum == 0 or end_lnum == 0 then
+      break
+    end
+
+    for lnum = start_lnum, end_lnum do
+      local word = vim.api.nvim_buf_get_lines(0, lnum - 1, lnum, true)[1]:match(pat.match)
+      if vim.fn.hlexists(word) ~= 0 then
+        vim.api.nvim_buf_set_extmark(0, ns, lnum - 1, 0, { end_col = #word, hl_group = word })
+      end
+    end
+  end
+
+  vim.fn.setpos('.', save_cursor)
+end
+
 -- Add custom highlights for list in `:h highlight-groups`.
 local bufname = vim.fs.normalize(vim.api.nvim_buf_get_name(0))
 if vim.endswith(bufname, '/doc/syntax.txt') then
-  require('vim.vimhelp').highlight_groups({
+  colorize_hl_groups({
     { start = [[\*group-name\*]], stop = '^======', match = '^(%w+)\t' },
     { start = [[\*highlight-groups\*]], stop = '^======', match = '^(%w+)\t' },
   })
 elseif vim.endswith(bufname, '/doc/treesitter.txt') then
-  require('vim.vimhelp').highlight_groups({
+  colorize_hl_groups({
     {
       start = [[\*treesitter-highlight-groups\*]],
       stop = [[\*treesitter-highlight-spell\*]],
@@ -17,24 +45,31 @@ elseif vim.endswith(bufname, '/doc/treesitter.txt') then
     },
   })
 elseif vim.endswith(bufname, '/doc/diagnostic.txt') then
-  require('vim.vimhelp').highlight_groups({
+  colorize_hl_groups({
     { start = [[\*diagnostic-highlights\*]], stop = '^======', match = '^(%w+)' },
   })
 elseif vim.endswith(bufname, '/doc/lsp.txt') then
-  require('vim.vimhelp').highlight_groups({
+  colorize_hl_groups({
     { start = [[\*lsp-highlight\*]], stop = '^------', match = '^(%w+)' },
     { start = [[\*lsp-semantic-highlight\*]], stop = '^======', match = '^@[%w%p]+' },
   })
 end
 
 vim.keymap.set('n', 'gO', function()
-  require('vim.vimhelp').show_toc()
-end, { buffer = 0, silent = true })
+  require('vim.treesitter._headings').show_toc()
+end, { buffer = 0, silent = true, desc = 'Show table of contents for current buffer' })
+
+vim.keymap.set('n', ']]', function()
+  require('vim.treesitter._headings').jump({ count = 1 })
+end, { buffer = 0, silent = false, desc = 'Jump to next section' })
+vim.keymap.set('n', '[[', function()
+  require('vim.treesitter._headings').jump({ count = -1 })
+end, { buffer = 0, silent = false, desc = 'Jump to previous section' })
 
 -- Add "runnables" for Lua/Vimscript code examples.
 ---@type table<integer, { lang: string, code: string }>
 local code_blocks = {}
-local tree = vim.treesitter.get_parser():parse()[1]
+local parser = assert(vim.treesitter.get_parser(0, 'vimdoc', { error = false }))
 local query = vim.treesitter.query.parse(
   'vimdoc',
   [[
@@ -46,10 +81,11 @@ local query = vim.treesitter.query.parse(
     (#set! @code lang @_lang))
 ]]
 )
+local root = parser:parse()[1]:root()
 local run_message_ns = vim.api.nvim_create_namespace('nvim.vimdoc.run_message')
 
 vim.api.nvim_buf_clear_namespace(0, run_message_ns, 0, -1)
-for _, match, metadata in query:iter_matches(tree:root(), 0, 0, -1) do
+for _, match, metadata in query:iter_matches(root, 0, 0, -1) do
   for id, nodes in pairs(match) do
     local name = query.captures[id]
     local node = nodes[1]
@@ -83,4 +119,5 @@ end, { buffer = true })
 
 vim.b.undo_ftplugin = (vim.b.undo_ftplugin or '')
   .. '\n exe "nunmap <buffer> gO" | exe "nunmap <buffer> g=="'
+  .. '\n exe "nunmap <buffer> ]]" | exe "nunmap <buffer> [["'
 vim.b.undo_ftplugin = vim.b.undo_ftplugin .. ' | call v:lua.vim.treesitter.stop()'