1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
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
|