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.

7 files changed, 221 insertions, 79 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 5c7354e3b5..05d14600f8 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -265,6 +265,9 @@ EDITOR
   to a literal "~" directory.
 • |hl-ComplMatchIns| shows matched text of the currently inserted completion.
 • |hl-PmenuMatch| and |hl-PmenuMatchSel| show matched text in completion popup.
+• |gO| now works in `help`, `checkhealth`, and `markdown` buffers.
+• Jump between sections in `help` and `checkhealth` buffers with `[[` and
+  `]]`.
 
 EVENTS
 
diff --git a/runtime/ftplugin/checkhealth.lua b/runtime/ftplugin/checkhealth.lua
new file mode 100644
index 0000000000..232846fa47
--- /dev/null
+++ b/runtime/ftplugin/checkhealth.lua
@@ -0,0 +1,14 @@
+vim.keymap.set('n', 'gO', function()
+  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, level = 1 })
+end, { buffer = 0, silent = false, desc = 'Jump to next section' })
+vim.keymap.set('n', '[[', function()
+  require('vim.treesitter._headings').jump({ count = -1, level = 1 })
+end, { buffer = 0, silent = false, desc = 'Jump to previous section' })
+
+vim.b.undo_ftplugin = (vim.b.undo_ftplugin or '')
+  .. '\n exe "nunmap <buffer> gO"'
+  .. '\n exe "nunmap <buffer> ]]" | exe "nunmap <buffer> [["'
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()'
diff --git a/runtime/ftplugin/markdown.lua b/runtime/ftplugin/markdown.lua
new file mode 100644
index 0000000000..9319ca7757
--- /dev/null
+++ b/runtime/ftplugin/markdown.lua
@@ -0,0 +1,14 @@
+vim.keymap.set('n', 'gO', function()
+  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' })
+
+vim.b.undo_ftplugin = (vim.b.undo_ftplugin or '')
+  .. '\n exe "nunmap <buffer> gO"'
+  .. '\n exe "nunmap <buffer> ]]" | exe "nunmap <buffer> [["'
diff --git a/runtime/lua/vim/treesitter/_headings.lua b/runtime/lua/vim/treesitter/_headings.lua
new file mode 100644
index 0000000000..4e8833c93a
--- /dev/null
+++ b/runtime/lua/vim/treesitter/_headings.lua
@@ -0,0 +1,144 @@
+local ts = vim.treesitter
+local api = vim.api
+
+--- Treesitter-based navigation functions for headings
+local M = {}
+
+-- TODO(clason): use runtimepath queries (for other languages)
+local heading_queries = {
+  vimdoc = [[
+    (h1 (heading) @h1)
+    (h2 (heading) @h2)
+    (h3 (heading) @h3)
+    (column_heading (heading) @h4)
+  ]],
+  markdown = [[
+    (setext_heading
+      heading_content: (_) @h1
+      (setext_h1_underline))
+    (setext_heading
+      heading_content: (_) @h2
+      (setext_h2_underline))
+    (atx_heading
+      (atx_h1_marker)
+      heading_content: (_) @h1)
+    (atx_heading
+      (atx_h2_marker)
+      heading_content: (_) @h2)
+    (atx_heading
+      (atx_h3_marker)
+      heading_content: (_) @h3)
+    (atx_heading
+      (atx_h4_marker)
+      heading_content: (_) @h4)
+    (atx_heading
+      (atx_h5_marker)
+      heading_content: (_) @h5)
+    (atx_heading
+      (atx_h6_marker)
+      heading_content: (_) @h6)
+  ]],
+}
+
+local function hash_tick(bufnr)
+  return tostring(vim.b[bufnr].changedtick)
+end
+
+---@class TS.Heading
+---@field bufnr integer
+---@field lnum integer
+---@field text string
+---@field level integer
+
+--- Extract headings from buffer
+--- @param bufnr integer buffer to extract headings from
+--- @return TS.Heading[]
+local get_headings = vim.func._memoize(hash_tick, function(bufnr)
+  local lang = ts.language.get_lang(vim.bo[bufnr].filetype)
+  if not lang then
+    return {}
+  end
+  local parser = assert(ts.get_parser(bufnr, lang, { error = false }))
+  local query = ts.query.parse(lang, heading_queries[lang])
+  local root = parser:parse()[1]:root()
+  local headings = {}
+  for id, node, _, _ in query:iter_captures(root, bufnr) do
+    local text = ts.get_node_text(node, bufnr)
+    local row, col = node:start()
+    --- why can't you just be normal?!
+    local skip ---@type boolean|integer
+    if lang == 'vimdoc' then
+      -- only column_headings at col 1 are headings, otherwise it's code examples
+      skip = (id == 4 and col > 0)
+        -- ignore tabular material
+        or (id == 4 and (text:find('\t') or text:find('  ')))
+        -- ignore tag-only headings
+        or (node:child_count() == 1 and node:child(0):type() == 'tag')
+    end
+    if not skip then
+      table.insert(headings, {
+        bufnr = bufnr,
+        lnum = row + 1,
+        text = text,
+        level = id,
+      })
+    end
+  end
+  return headings
+end)
+
+--- Show a table of contents for the help buffer in a loclist
+function M.show_toc()
+  local bufnr = api.nvim_get_current_buf()
+  local headings = get_headings(bufnr)
+  if #headings == 0 then
+    return
+  end
+  -- add indentation for nicer list formatting
+  for _, heading in pairs(headings) do
+    if heading.level > 2 then
+      heading.text = '  ' .. heading.text
+    end
+    if heading.level > 4 then
+      heading.text = '  ' .. heading.text
+    end
+  end
+  vim.fn.setloclist(0, headings, ' ')
+  vim.fn.setloclist(0, {}, 'a', { title = 'Help TOC' })
+  vim.cmd.lopen()
+end
+
+--- Jump to section
+--- @param opts table jump options
+---  - count integer direction to jump (>0 forward, <0 backward)
+---  - level integer only consider headings up to level
+--- todo(clason): support count
+function M.jump(opts)
+  local bufnr = api.nvim_get_current_buf()
+  local headings = get_headings(bufnr)
+  if #headings == 0 then
+    return
+  end
+
+  local winid = api.nvim_get_current_win()
+  local curpos = vim.fn.getcurpos(winid)[2] --[[@as integer]]
+  local maxlevel = opts.level or 6
+
+  if opts.count > 0 then
+    for _, heading in ipairs(headings) do
+      if heading.lnum > curpos and heading.level <= maxlevel then
+        api.nvim_win_set_cursor(winid, { heading.lnum, 0 })
+        return
+      end
+    end
+  elseif opts.count < 0 then
+    for i = #headings, 1, -1 do
+      if headings[i].lnum < curpos and headings[i].level <= maxlevel then
+        api.nvim_win_set_cursor(winid, { headings[i].lnum, 0 })
+        return
+      end
+    end
+  end
+end
+
+return M
diff --git a/runtime/lua/vim/treesitter/language.lua b/runtime/lua/vim/treesitter/language.lua
index 16d19bfc5a..38d309a102 100644
--- a/runtime/lua/vim/treesitter/language.lua
+++ b/runtime/lua/vim/treesitter/language.lua
@@ -5,6 +5,7 @@ local M = {}
 ---@type table<string,string>
 local ft_to_lang = {
   help = 'vimdoc',
+  checkhealth = 'vimdoc',
 }
 
 --- Returns the filetypes for which a parser named {lang} is used.
diff --git a/runtime/lua/vim/vimhelp.lua b/runtime/lua/vim/vimhelp.lua
deleted file mode 100644
index a494d311b1..0000000000
--- a/runtime/lua/vim/vimhelp.lua
+++ /dev/null
@@ -1,71 +0,0 @@
--- Extra functionality for displaying Vim help.
-
-local M = {}
-
---- 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}[]
-function M.highlight_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
-
---- Show a table of contents for the help buffer in a loclist
-function M.show_toc()
-  local bufnr = vim.api.nvim_get_current_buf()
-  local parser = assert(vim.treesitter.get_parser(bufnr, 'vimdoc', { error = false }))
-  local query = vim.treesitter.query.parse(
-    parser:lang(),
-    [[
-    (h1 (heading) @h1)
-    (h2 (heading) @h2)
-    (h3 (heading) @h3)
-    (column_heading (heading) @h4)
-  ]]
-  )
-  local root = parser:parse()[1]:root()
-  local headings = {}
-  for id, node, _, _ in query:iter_captures(root, bufnr) do
-    local text = vim.treesitter.get_node_text(node, bufnr)
-    local capture = query.captures[id]
-    local row, col = node:start()
-    -- only column_headings at col 1 are headings, otherwise it's code examples
-    local is_code = (capture == 'h4' and col > 0)
-    -- ignore tabular material
-    local is_table = (capture == 'h4' and (text:find('\t') or text:find('  ')))
-    -- ignore tag-only headings
-    local is_tag = node:child_count() == 1 and node:child(0):type() == 'tag'
-    if not (is_code or is_table or is_tag) then
-      table.insert(headings, {
-        bufnr = bufnr,
-        lnum = row + 1,
-        text = (capture == 'h3' or capture == 'h4') and '  ' .. text or text,
-      })
-    end
-  end
-  vim.fn.setloclist(0, headings, ' ')
-  vim.fn.setloclist(0, {}, 'a', { title = 'Help TOC' })
-  vim.cmd.lopen()
-end
-
-return M