diff options
Diffstat (limited to 'runtime/lua/vim/treesitter/playground.lua')
| -rw-r--r-- | runtime/lua/vim/treesitter/playground.lua | 186 |
1 files changed, 186 insertions, 0 deletions
diff --git a/runtime/lua/vim/treesitter/playground.lua b/runtime/lua/vim/treesitter/playground.lua new file mode 100644 index 0000000000..bb073290c6 --- /dev/null +++ b/runtime/lua/vim/treesitter/playground.lua @@ -0,0 +1,186 @@ +local api = vim.api + +local M = {} + +---@class Playground +---@field ns number API namespace +---@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 +---@param lang string|nil Language of source buffer +--- +---@return Playground|nil +---@return string|nil Error message, if any +--- +---@private +function M.new(self, bufnr, lang) + local ok, parser = pcall(vim.treesitter.get_parser, bufnr or 0, lang) + 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 |