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.

1 files changed, 23 insertions, 10 deletions

diff --git a/runtime/lua/vim/treesitter/dev.lua b/runtime/lua/vim/treesitter/dev.lua
index ee4d981c16..14e21901df 100644
--- a/runtime/lua/vim/treesitter/dev.lua
+++ b/runtime/lua/vim/treesitter/dev.lua
@@ -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