feat(treesitter): add vim.treesitter.show_tree() (#21322)

Add a "show_tree" function to view a textual representation of the
nodes in a language tree in a window. Moving the cursor in the
window highlights the corresponding text in the source buffer, and
moving the cursor in the source buffer highlights the corresponding
nodes in the window.

3 files changed, 381 insertions, 2 deletions

diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 0603ddb421..5c9210cc2d 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -277,7 +277,7 @@ end
 ---@param opts table Optional keyword arguments:
 ---             - ignore_injections boolean Ignore injected languages (default true)
 ---
----@return userdata |tsnode| under the cursor
+---@return userdata|nil |tsnode| under the cursor
 function M.get_node_at_pos(bufnr, row, col, opts)
   if bufnr == 0 then
     bufnr = a.nvim_get_current_buf()
@@ -347,4 +347,197 @@ function M.stop(bufnr)
   vim.bo[bufnr].syntax = 'on'
 end
 
+--- Open a window that displays a textual representation of the nodes in the language tree.
+---
+--- While in the window, press "a" to toggle display of anonymous nodes, "I" to toggle the
+--- display of the source language of each node, and press <Enter> to jump to the node under the
+--- cursor in the source buffer.
+---
+---@param opts table|nil Optional options table with the following possible keys:
+---                      - bufnr (number|nil): Buffer to draw the tree into. If omitted, a new
+---                        buffer is created.
+---                      - winid (number|nil): Window id to display the tree buffer in. If omitted,
+---                        a new window is created with {command}.
+---                      - command (string|nil): Vimscript command to create the window. Default
+---                        value is "topleft 60vnew". Only used when {winid} is nil.
+---                      - title (string|fun(bufnr:number):string|nil): Title of the window. If a
+---                        function, it accepts the buffer number of the source buffer as its only
+---                        argument and should return a string.
+function M.show_tree(opts)
+  vim.validate({
+    opts = { opts, 't', true },
+  })
+
+  local Playground = require('vim.treesitter.playground')
+  local buf = a.nvim_get_current_buf()
+  local win = a.nvim_get_current_win()
+  local pg = assert(Playground:new(buf))
+
+  opts = opts or {}
+
+  -- Close any existing playground window
+  if vim.b[buf].playground then
+    local w = vim.b[buf].playground
+    if a.nvim_win_is_valid(w) then
+      a.nvim_win_close(w, true)
+    end
+  end
+
+  local w = opts.winid
+  if not w then
+    vim.cmd(opts.command or 'topleft 60vnew')
+    w = a.nvim_get_current_win()
+  end
+
+  local b = opts.bufnr
+  if b then
+    a.nvim_win_set_buf(w, b)
+  else
+    b = a.nvim_win_get_buf(w)
+  end
+
+  vim.b[buf].playground = w
+
+  vim.wo[w].scrolloff = 5
+  vim.wo[w].wrap = false
+  vim.bo[b].buflisted = false
+  vim.bo[b].buftype = 'nofile'
+  vim.bo[b].bufhidden = 'wipe'
+
+  local title = opts.title
+  if not title then
+    local bufname = a.nvim_buf_get_name(buf)
+    title = string.format('Syntax tree for %s', vim.fn.fnamemodify(bufname, ':.'))
+  elseif type(title) == 'function' then
+    title = title(buf)
+  end
+
+  assert(type(title) == 'string', 'Window title must be a string')
+  a.nvim_buf_set_name(b, title)
+
+  pg:draw(b)
+
+  vim.fn.matchadd('Comment', '\\[[0-9:-]\\+\\]')
+  vim.fn.matchadd('String', '".*"')
+
+  a.nvim_buf_clear_namespace(buf, pg.ns, 0, -1)
+  a.nvim_buf_set_keymap(b, 'n', '<CR>', '', {
+    desc = 'Jump to the node under the cursor in the source buffer',
+    callback = function()
+      local row = a.nvim_win_get_cursor(w)[1]
+      local pos = pg:get(row)
+      a.nvim_set_current_win(win)
+      a.nvim_win_set_cursor(win, { pos.lnum + 1, pos.col })
+    end,
+  })
+  a.nvim_buf_set_keymap(b, 'n', 'a', '', {
+    desc = 'Toggle anonymous nodes',
+    callback = function()
+      pg.opts.anon = not pg.opts.anon
+      pg:draw(b)
+    end,
+  })
+  a.nvim_buf_set_keymap(b, 'n', 'I', '', {
+    desc = 'Toggle language display',
+    callback = function()
+      pg.opts.lang = not pg.opts.lang
+      pg:draw(b)
+    end,
+  })
+
+  local group = a.nvim_create_augroup('treesitter/playground', {})
+
+  a.nvim_create_autocmd('CursorMoved', {
+    group = group,
+    buffer = b,
+    callback = function()
+      a.nvim_buf_clear_namespace(buf, pg.ns, 0, -1)
+      local row = a.nvim_win_get_cursor(w)[1]
+      local pos = pg:get(row)
+      a.nvim_buf_set_extmark(buf, pg.ns, pos.lnum, pos.col, {
+        end_row = pos.end_lnum,
+        end_col = math.max(0, pos.end_col),
+        hl_group = 'Visual',
+      })
+    end,
+  })
+
+  a.nvim_create_autocmd('CursorMoved', {
+    group = group,
+    buffer = buf,
+    callback = function()
+      if not a.nvim_buf_is_loaded(b) then
+        return true
+      end
+
+      a.nvim_buf_clear_namespace(b, pg.ns, 0, -1)
+
+      local cursor = a.nvim_win_get_cursor(win)
+      local cursor_node =
+        M.get_node_at_pos(buf, cursor[1] - 1, cursor[2], { ignore_injections = false })
+      if not cursor_node then
+        return
+      end
+
+      local cursor_node_id = cursor_node:id()
+      for i, v in pg:iter() do
+        if v.id == cursor_node_id then
+          local start = v.depth
+          local end_col = start + #v.text
+          a.nvim_buf_set_extmark(b, pg.ns, i - 1, start, {
+            end_col = end_col,
+            hl_group = 'Visual',
+          })
+          a.nvim_win_set_cursor(w, { i, 0 })
+          break
+        end
+      end
+    end,
+  })
+
+  a.nvim_create_autocmd({ 'TextChanged', 'InsertLeave' }, {
+    group = group,
+    buffer = buf,
+    callback = function()
+      if not a.nvim_buf_is_loaded(b) then
+        return true
+      end
+
+      pg = assert(Playground:new(buf))
+      pg:draw(b)
+    end,
+  })
+
+  a.nvim_create_autocmd('BufLeave', {
+    group = group,
+    buffer = b,
+    callback = function()
+      a.nvim_buf_clear_namespace(buf, pg.ns, 0, -1)
+    end,
+  })
+
+  a.nvim_create_autocmd('BufLeave', {
+    group = group,
+    buffer = buf,
+    callback = function()
+      if not a.nvim_buf_is_loaded(b) then
+        return true
+      end
+
+      a.nvim_buf_clear_namespace(b, pg.ns, 0, -1)
+    end,
+  })
+
+  a.nvim_create_autocmd('BufHidden', {
+    group = group,
+    buffer = buf,
+    once = true,
+    callback = function()
+      if a.nvim_win_is_valid(w) then
+        a.nvim_win_close(w, true)
+      end
+    end,
+  })
+end
+
 return M
diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index e9d70c4204..a1e96f8ef2 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -608,7 +608,9 @@ end
 ---@return userdata|nil Found |tsnode|
 function LanguageTree:named_node_for_range(range, opts)
   local tree = self:tree_for_range(range, opts)
-  return tree:root():named_descendant_for_range(unpack(range))
+  if tree then
+    return tree:root():named_descendant_for_range(unpack(range))
+  end
 end
 
 --- Gets the appropriate language that contains {range}.
diff --git a/runtime/lua/vim/treesitter/playground.lua b/runtime/lua/vim/treesitter/playground.lua
new file mode 100644
index 0000000000..325d303df5
--- /dev/null
+++ b/runtime/lua/vim/treesitter/playground.lua
@@ -0,0 +1,184 @@
+local api = vim.api
+
+local M = {}
+
+---@class Playground
+---@field opts table Options table with the following keys:
+---                  - anon (boolean): If true, display anonymous nodes
+---                  - lang (boolean): If true, display the language alongside each node
+---
+---@class Node
+---@field id number Node id
+---@field text string Node text
+---@field named boolean True if this is a named (non-anonymous) node
+---@field depth number Depth of the node within the tree
+---@field lnum number Beginning line number of this node in the source buffer
+---@field col number Beginning column number of this node in the source buffer
+---@field end_lnum number Final line number of this node in the source buffer
+---@field end_col number Final column number of this node in the source buffer
+---@field lang string Source language of this node
+
+--- Traverse all child nodes starting at {node}.
+---
+--- This is a recursive function. The {depth} parameter indicates the current recursion level.
+--- {lang} is a string indicating the language of the tree currently being traversed. Each traversed
+--- node is added to {tree}. When recursion completes, {tree} is an array of all nodes in the order
+--- they were visited.
+---
+--- {injections} is a table mapping node ids from the primary tree to language tree injections. Each
+--- injected language has a series of trees nested within the primary language's tree, and the root
+--- node of each of these trees is contained within a node in the primary tree. The {injections}
+--- table maps nodes in the primary tree to root nodes of injected trees.
+---
+---@param node userdata Starting node to begin traversal |tsnode|
+---@param depth number Current recursion depth
+---@param lang string Language of the tree currently being traversed
+---@param injections table Mapping of node ids to root nodes of injected language trees (see
+---                        explanation above)
+---@param tree Node[] Output table containing a list of tables each representing a node in the tree
+---@private
+local function traverse(node, depth, lang, injections, tree)
+  local injection = injections[node:id()]
+  if injection then
+    traverse(injection.root, depth, injection.lang, injections, tree)
+  end
+
+  for child, field in node:iter_children() do
+    local type = child:type()
+    local lnum, col, end_lnum, end_col = child:range()
+    local named = child:named()
+    local text
+    if named then
+      if field then
+        text = string.format('%s: (%s)', field, type)
+      else
+        text = string.format('(%s)', type)
+      end
+    else
+      text = string.format('"%s"', type:gsub('\n', '\\n'))
+    end
+
+    table.insert(tree, {
+      id = child:id(),
+      text = text,
+      named = named,
+      depth = depth,
+      lnum = lnum,
+      col = col,
+      end_lnum = end_lnum,
+      end_col = end_col,
+      lang = lang,
+    })
+
+    traverse(child, depth + 1, lang, injections, tree)
+  end
+
+  return tree
+end
+
+--- Create a new Playground object.
+---
+---@param bufnr number Source buffer number
+---
+---@return Playground|nil
+---@return string|nil Error message, if any
+---
+---@private
+function M.new(self, bufnr)
+  local ok, parser = pcall(vim.treesitter.get_parser, bufnr or 0)
+  if not ok then
+    return nil, 'No parser available for the given buffer'
+  end
+
+  -- For each child tree (injected language), find the root of the tree and locate the node within
+  -- the primary tree that contains that root. Add a mapping from the node in the primary tree to
+  -- the root in the child tree to the {injections} table.
+  local root = parser:parse()[1]:root()
+  local injections = {}
+  parser:for_each_child(function(child, lang)
+    child:for_each_tree(function(tree)
+      local r = tree:root()
+      local node = root:named_descendant_for_range(r:range())
+      if node then
+        injections[node:id()] = {
+          lang = lang,
+          root = r,
+        }
+      end
+    end)
+  end)
+
+  local nodes = traverse(root, 0, parser:lang(), injections, {})
+
+  local named = {}
+  for _, v in ipairs(nodes) do
+    if v.named then
+      named[#named + 1] = v
+    end
+  end
+
+  local t = {
+    ns = api.nvim_create_namespace(''),
+    nodes = nodes,
+    named = named,
+    opts = {
+      anon = false,
+      lang = false,
+    },
+  }
+
+  setmetatable(t, self)
+  self.__index = self
+  return t
+end
+
+--- Write the contents of this Playground into {bufnr}.
+---
+---@param bufnr number Buffer number to write into.
+---@private
+function M.draw(self, bufnr)
+  vim.bo[bufnr].modifiable = true
+  local lines = {}
+  for _, item in self:iter() do
+    lines[#lines + 1] = table.concat({
+      string.rep(' ', item.depth),
+      item.text,
+      item.lnum == item.end_lnum
+          and string.format(' [%d:%d-%d]', item.lnum + 1, item.col + 1, item.end_col)
+        or string.format(
+          ' [%d:%d-%d:%d]',
+          item.lnum + 1,
+          item.col + 1,
+          item.end_lnum + 1,
+          item.end_col
+        ),
+      self.opts.lang and string.format(' %s', item.lang) or '',
+    })
+  end
+  api.nvim_buf_set_lines(bufnr, 0, -1, false, lines)
+  vim.bo[bufnr].modifiable = false
+end
+
+--- Get node {i} from this Playground object.
+---
+--- The node number is dependent on whether or not anonymous nodes are displayed.
+---
+---@param i number Node number to get
+---@return Node
+---@private
+function M.get(self, i)
+  local t = self.opts.anon and self.nodes or self.named
+  return t[i]
+end
+
+--- Iterate over all of the nodes in this Playground object.
+---
+---@return function Iterator over all nodes in this Playground
+---@return table
+---@return number
+---@private
+function M.iter(self)
+  return ipairs(self.opts.anon and self.nodes or self.named)
+end
+
+return M