fix(treesitter): InspectTree only show the largest injection #37906

Problem:
:InspectTree don't show luadoc injection lua file. Since luadoc share
the same "root" with comment in their common primary (lua) tree.
Current logic simply show the largest (comment injection) and ignore all
smaller one (luadoc injection).

Solution:
Handle different lang injections separately. Then sort them by
byte_length to ensure the draw tree consistent.
This commit is contained in:
phanium
2026-02-25 05:22:30 +08:00
committed by GitHub
parent 39c4e0f336
commit 16aab4cb48
2 changed files with 62 additions and 10 deletions

View File

@@ -45,7 +45,7 @@ local TSTreeView = {}
---@param depth integer Current recursion depth
---@param field string|nil The field of the current node
---@param lang string Language of the tree currently being traversed
---@param injections table<string, vim.treesitter.dev.Injection> Mapping of node ids to root nodes
---@param injections table<string, vim.treesitter.dev.Injection[]> Mapping of node ids to root nodes
--- of injected language trees (see explanation above)
---@param tree vim.treesitter.dev.Node[] Output table containing a list of tables each representing a node in the tree
local function traverse(node, depth, field, lang, injections, tree)
@@ -56,8 +56,7 @@ local function traverse(node, depth, field, lang, injections, tree)
field = field,
})
local injection = injections[node:id()]
if injection then
for _, injection in ipairs(injections[node:id()] or {}) do
traverse(injection.root, depth + 1, nil, injection.lang, injections, tree)
end
@@ -94,7 +93,7 @@ function TSTreeView:new(bufnr, lang)
-- 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(true)[1]:root()
local injections = {} ---@type table<string, vim.treesitter.dev.Injection>
local injections = {} ---@type table<string, table<string, TSNode>>
parser:for_each_tree(function(parent_tree, parent_ltree)
local parent = parent_tree:root()
@@ -106,18 +105,32 @@ function TSTreeView:new(bufnr, lang)
if Range.contains(parent_range, r_range) then
local node = assert(parent:named_descendant_for_range(r:range()))
local id = node:id()
if not injections[id] or r:byte_length() > injections[id].root:byte_length() then
injections[id] = {
lang = child:lang(),
root = r,
}
local ilang = child:lang()
injections[id] = injections[id] or {}
local injection = injections[id][ilang]
if not injection or r:byte_length() > injection:byte_length() then
injections[id][ilang] = r
end
end
end
end
end)
local nodes = traverse(root, 0, nil, parser:lang(), injections, {})
local sorted_injections = {} ---@type table<string, vim.treesitter.dev.Injection[]>
for id, lang_injections in pairs(injections) do
local langs = vim.tbl_keys(lang_injections)
---@param a string
---@param b string
table.sort(langs, function(a, b)
return lang_injections[a]:byte_length() > lang_injections[b]:byte_length()
end)
---@param ilang string
sorted_injections[id] = vim.tbl_map(function(ilang)
return { lang = ilang, root = lang_injections[ilang] }
end, langs)
end
local nodes = traverse(root, 0, nil, parser:lang(), sorted_injections, {})
local named = {} ---@type vim.treesitter.dev.Node[]
for _, v in ipairs(nodes) do