1 files changed, 339 insertions, 41 deletions

diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 6431162799..582922ecb6 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -3,10 +3,7 @@ local query = require('vim.treesitter.query')
 local language = require('vim.treesitter.language')
 local LanguageTree = require('vim.treesitter.languagetree')
 
--- TODO(bfredl): currently we retain parsers for the lifetime of the buffer.
--- Consider use weak references to release parser if all plugins are done with
--- it.
-local parsers = {}
+local parsers = setmetatable({}, { __mode = 'v' })
 
 local M = vim.tbl_extend('error', query, language)
 
@@ -28,13 +25,15 @@ setmetatable(M, {
   end,
 })
 
---- Creates a new parser.
+--- Creates a new parser
 ---
---- It is not recommended to use this, use vim.treesitter.get_parser() instead.
+--- It is not recommended to use this; use |get_parser()| instead.
 ---
----@param bufnr The buffer the parser will be tied to
----@param lang The language of the parser
----@param opts Options to pass to the created language tree
+---@param bufnr string Buffer the parser will be tied to (0 for current buffer)
+---@param lang string Language of the parser
+---@param opts (table|nil) Options to pass to the created language tree
+---
+---@return LanguageTree |LanguageTree| object to use for parsing
 function M._create_parser(bufnr, lang, opts)
   language.require_language(lang)
   if bufnr == 0 then
@@ -74,16 +73,15 @@ function M._create_parser(bufnr, lang, opts)
   return self
 end
 
---- Gets the parser for this bufnr / ft combination.
+--- Returns the parser for a specific buffer and filetype and attaches it to the buffer
 ---
---- If needed this will create the parser.
---- Unconditionally attach the provided callback
+--- If needed, this will create the parser.
 ---
----@param bufnr The buffer the parser should be tied to
----@param lang The filetype of this parser
----@param opts Options object to pass to the created language tree
+---@param bufnr (number|nil) Buffer the parser should be tied to (default: current buffer)
+---@param lang (string|nil) Filetype of this parser (default: buffer filetype)
+---@param opts (table|nil) Options to pass to the created language tree
 ---
----@returns The parser
+---@return LanguageTree |LanguageTree| object to use for parsing
 function M.get_parser(bufnr, lang, opts)
   opts = opts or {}
 
@@ -103,11 +101,13 @@ function M.get_parser(bufnr, lang, opts)
   return parsers[bufnr]
 end
 
---- Gets a string parser
+--- Returns a string parser
+---
+---@param str string Text to parse
+---@param lang string Language of this string
+---@param opts (table|nil) Options to pass to the created language tree
 ---
----@param str The string to parse
----@param lang The language of this string
----@param opts Options to pass to the created language tree
+---@return LanguageTree |LanguageTree| object to use for parsing
 function M.get_string_parser(str, lang, opts)
   vim.validate({
     str = { str, 'string' },
@@ -120,10 +120,10 @@ end
 
 --- Determines whether a node is the ancestor of another
 ---
----@param dest table the possible ancestor
----@param source table the possible descendant node
+---@param dest userdata Possible ancestor |tsnode|
+---@param source userdata Possible descendant |tsnode|
 ---
----@returns (boolean) True if dest is an ancestor of source
+---@return boolean True if {dest} is an ancestor of {source}
 function M.is_ancestor(dest, source)
   if not (dest and source) then
     return false
@@ -141,11 +141,11 @@ function M.is_ancestor(dest, source)
   return false
 end
 
---- Get the node's range or unpack a range table
+--- Returns the node's range or an unpacked range table
 ---
----@param node_or_range table
+---@param node_or_range (userdata|table) |tsnode| or table of positions
 ---
----@returns start_row, start_col, end_row, end_col
+---@return table `{ start_row, start_col, end_row, end_col }`
 function M.get_node_range(node_or_range)
   if type(node_or_range) == 'table' then
     return unpack(node_or_range)
@@ -154,11 +154,13 @@ function M.get_node_range(node_or_range)
   end
 end
 
----Determines whether (line, col) position is in node range
+--- Determines whether (line, col) position is in node range
+---
+---@param node userdata |tsnode| defining the range
+---@param line number Line (0-based)
+---@param col number Column (0-based)
 ---
----@param node Node defining the range
----@param line A line (0-based)
----@param col A column (0-based)
+---@return boolean True if the position is in node range
 function M.is_in_node_range(node, line, col)
   local start_line, start_col, end_line, end_col = M.get_node_range(node)
   if line >= start_line and line <= end_line then
@@ -176,11 +178,12 @@ function M.is_in_node_range(node, line, col)
   end
 end
 
----Determines if a node contains a range
----@param node table The node
----@param range table The range
+--- Determines if a node contains a range
 ---
----@returns (boolean) True if the node contains the range
+---@param node userdata |tsnode|
+---@param range table
+---
+---@return boolean True if the {node} contains the {range}
 function M.node_contains(node, range)
   local start_row, start_col, end_row, end_col = node:range()
   local start_fits = start_row < range[1] or (start_row == range[1] and start_col <= range[2])
@@ -189,13 +192,17 @@ function M.node_contains(node, range)
   return start_fits and end_fits
 end
 
----Gets a list of captures for a given cursor position
----@param bufnr number The buffer number
----@param row number The position row
----@param col number The position column
+--- Returns a list of highlight captures at the given position
+---
+--- Each capture is represented by a table containing the capture name as a string as
+--- well as a table of metadata (`priority`, `conceal`, ...; empty if none are defined).
+---
+---@param bufnr number Buffer number (0 for current buffer)
+---@param row number Position row
+---@param col number Position column
 ---
----@returns (table) A table of captures
-function M.get_captures_at_position(bufnr, row, col)
+---@return table[] List of captures `{ capture = "capture name", metadata = { ... } }`
+function M.get_captures_at_pos(bufnr, row, col)
   if bufnr == 0 then
     bufnr = a.nvim_get_current_buf()
   end
@@ -233,7 +240,7 @@ function M.get_captures_at_position(bufnr, row, col)
       if M.is_in_node_range(node, row, col) then
         local c = q._query.captures[capture] -- name of the capture in the query
         if c ~= nil then
-          table.insert(matches, { capture = c, priority = metadata.priority })
+          table.insert(matches, { capture = c, metadata = metadata, lang = tree:lang() })
         end
       end
     end
@@ -241,4 +248,295 @@ function M.get_captures_at_position(bufnr, row, col)
   return matches
 end
 
+--- Returns a list of highlight capture names under the cursor
+---
+---@param winnr (number|nil) Window handle or 0 for current window (default)
+---
+---@return string[] List of capture names
+function M.get_captures_at_cursor(winnr)
+  winnr = winnr or 0
+  local bufnr = a.nvim_win_get_buf(winnr)
+  local cursor = a.nvim_win_get_cursor(winnr)
+
+  local data = M.get_captures_at_pos(bufnr, cursor[1] - 1, cursor[2])
+
+  local captures = {}
+
+  for _, capture in ipairs(data) do
+    table.insert(captures, capture.capture)
+  end
+
+  return captures
+end
+
+--- Returns the smallest named node at the given position
+---
+---@param bufnr number Buffer number (0 for current buffer)
+---@param row number Position row
+---@param col number Position column
+---@param opts table Optional keyword arguments:
+---             - lang string|nil Parser language
+---             - ignore_injections boolean Ignore injected languages (default true)
+---
+---@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()
+  end
+  local ts_range = { row, col, row, col }
+
+  local root_lang_tree = M.get_parser(bufnr, opts.lang)
+  if not root_lang_tree then
+    return
+  end
+
+  return root_lang_tree:named_node_for_range(ts_range, opts)
+end
+
+--- Returns the smallest named node under the cursor
+---
+---@param winnr (number|nil) Window handle or 0 for current window (default)
+---
+---@return string Name of node under the cursor
+function M.get_node_at_cursor(winnr)
+  winnr = winnr or 0
+  local bufnr = a.nvim_win_get_buf(winnr)
+  local cursor = a.nvim_win_get_cursor(winnr)
+
+  return M.get_node_at_pos(bufnr, cursor[1] - 1, cursor[2], { ignore_injections = false }):type()
+end
+
+--- Starts treesitter highlighting for a buffer
+---
+--- Can be used in an ftplugin or FileType autocommand.
+---
+--- Note: By default, disables regex syntax highlighting, which may be required for some plugins.
+--- In this case, add ``vim.bo.syntax = 'on'`` after the call to `start`.
+---
+--- Example:
+--- <pre>lua
+--- vim.api.nvim_create_autocmd( 'FileType', { pattern = 'tex',
+---     callback = function(args)
+---         vim.treesitter.start(args.buf, 'latex')
+---         vim.bo[args.buf].syntax = 'on'  -- only if additional legacy syntax is needed
+---     end
+--- })
+--- </pre>
+---
+---@param bufnr (number|nil) Buffer to be highlighted (default: current buffer)
+---@param lang (string|nil) Language of the parser (default: buffer filetype)
+function M.start(bufnr, lang)
+  bufnr = bufnr or a.nvim_get_current_buf()
+  local parser = M.get_parser(bufnr, lang)
+  M.highlighter.new(parser)
+end
+
+--- Stops treesitter highlighting for a buffer
+---
+---@param bufnr (number|nil) Buffer to stop highlighting (default: current buffer)
+function M.stop(bufnr)
+  bufnr = bufnr or a.nvim_get_current_buf()
+
+  if M.highlighter.active[bufnr] then
+    M.highlighter.active[bufnr]:destroy()
+  end
+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:
+---                      - lang (string|nil): The language of the source buffer. If omitted, the
+---                        filetype of the source buffer is used.
+---                      - 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 },
+  })
+
+  opts = opts or {}
+
+  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.lang))
+
+  -- 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], {
+        lang = opts.lang,
+        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, opts.lang))
+      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