feat(treesitter): show root nodes in :InspectTree (#26944)

Co-authored-by: altermo <>
Co-authored-by: Jongwook Choi <wookayin@gmail.com>

3 files changed, 48 insertions, 40 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 3dc85eddbb..c67713b657 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -225,6 +225,7 @@ The following new APIs and features were added.
   • Improved error messages for query parsing.
   • `:InspectTree` (|vim.treesitter.inspect_tree()|) shows node ranges in
     0-based indexing instead of 1-based indexing.
+  • `:InspectTree` (|vim.treesitter.inspect_tree()|) shows root nodes
 
 • |vim.ui.open()| opens URIs using the system default handler (macOS `open`,
   Windows `explorer`, Linux `xdg-open`, etc.)
diff --git a/runtime/lua/vim/treesitter/dev.lua b/runtime/lua/vim/treesitter/dev.lua
index e1f93a654b..551067533a 100644
--- a/runtime/lua/vim/treesitter/dev.lua
+++ b/runtime/lua/vim/treesitter/dev.lua
@@ -43,25 +43,26 @@ local TSTreeView = {}
 ---
 ---@param node TSNode Starting node to begin traversal |tsnode|
 ---@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, TSP.Injection> Mapping of node ids to root nodes
 ---                  of injected language trees (see explanation above)
 ---@param tree TSP.Node[] Output table containing a list of tables each representing a node in the tree
-local function traverse(node, depth, lang, injections, tree)
+local function traverse(node, depth, field, lang, injections, tree)
+  table.insert(tree, {
+    node = node,
+    depth = depth,
+    lang = lang,
+    field = field,
+  })
+
   local injection = injections[node:id()]
   if injection then
-    traverse(injection.root, depth, injection.lang, injections, tree)
+    traverse(injection.root, depth + 1, nil, injection.lang, injections, tree)
   end
 
-  for child, field in node:iter_children() do
-    table.insert(tree, {
-      node = child,
-      field = field,
-      depth = depth,
-      lang = lang,
-    })
-
-    traverse(child, depth + 1, lang, injections, tree)
+  for child, child_field in node:iter_children() do
+    traverse(child, depth + 1, child_field, lang, injections, tree)
   end
 
   return tree
@@ -106,7 +107,7 @@ function TSTreeView:new(bufnr, lang)
     end
   end)
 
-  local nodes = traverse(root, 0, parser:lang(), injections, {})
+  local nodes = traverse(root, 0, nil, parser:lang(), injections, {})
 
   local named = {} ---@type TSP.Node[]
   for _, v in ipairs(nodes) do
diff --git a/test/functional/treesitter/inspect_tree_spec.lua b/test/functional/treesitter/inspect_tree_spec.lua
index 0102838b82..a3d44ff906 100644
--- a/test/functional/treesitter/inspect_tree_spec.lua
+++ b/test/functional/treesitter/inspect_tree_spec.lua
@@ -26,9 +26,10 @@ describe('vim.treesitter.inspect_tree', function()
     ]])
 
     expect_tree [[
-      (function_call ; [0, 0] - [0, 7]
-        name: (identifier) ; [0, 0] - [0, 5]
-        arguments: (arguments)) ; [0, 5] - [0, 7]
+      (chunk ; [0, 0] - [2, 0]
+        (function_call ; [0, 0] - [0, 7]
+          name: (identifier) ; [0, 0] - [0, 5]
+          arguments: (arguments))) ; [0, 5] - [0, 7]
       ]]
   end)
 
@@ -44,11 +45,12 @@ describe('vim.treesitter.inspect_tree', function()
     feed('a')
 
     expect_tree [[
-      (function_call ; [0, 0] - [0, 7]
-        name: (identifier) ; [0, 0] - [0, 5]
-        arguments: (arguments ; [0, 5] - [0, 7]
-          "(" ; [0, 5] - [0, 6]
-          ")")) ; [0, 6] - [0, 7]
+      (chunk ; [0, 0] - [2, 0]
+        (function_call ; [0, 0] - [0, 7]
+          name: (identifier) ; [0, 0] - [0, 5]
+          arguments: (arguments ; [0, 5] - [0, 7]
+            "(" ; [0, 5] - [0, 6]
+            ")"))) ; [0, 6] - [0, 7]
       ]]
   end)
 
@@ -66,16 +68,18 @@ describe('vim.treesitter.inspect_tree', function()
     ]])
 
     expect_tree [[
-      (section ; [0, 0] - [4, 0]
-        (fenced_code_block ; [0, 0] - [3, 0]
-          (fenced_code_block_delimiter) ; [0, 0] - [0, 3]
-          (info_string ; [0, 3] - [0, 6]
-            (language)) ; [0, 3] - [0, 6]
-          (block_continuation) ; [1, 0] - [1, 0]
-          (code_fence_content ; [1, 0] - [2, 0]
-            (return_statement) ; [1, 0] - [1, 6]
-            (block_continuation)) ; [2, 0] - [2, 0]
-          (fenced_code_block_delimiter))) ; [2, 0] - [2, 3]
+      (document ; [0, 0] - [4, 0]
+        (section ; [0, 0] - [4, 0]
+          (fenced_code_block ; [0, 0] - [3, 0]
+            (fenced_code_block_delimiter) ; [0, 0] - [0, 3]
+            (info_string ; [0, 3] - [0, 6]
+              (language)) ; [0, 3] - [0, 6]
+            (block_continuation) ; [1, 0] - [1, 0]
+            (code_fence_content ; [1, 0] - [2, 0]
+              (chunk ; [1, 0] - [2, 0]
+                (return_statement)) ; [1, 0] - [1, 6]
+              (block_continuation)) ; [2, 0] - [2, 0]
+            (fenced_code_block_delimiter)))) ; [2, 0] - [2, 3]
       ]]
   end)
 
@@ -94,16 +98,18 @@ describe('vim.treesitter.inspect_tree', function()
     feed('I')
 
     expect_tree [[
-      (section ; [0, 0] - [4, 0] markdown
-        (fenced_code_block ; [0, 0] - [3, 0] markdown
-          (fenced_code_block_delimiter) ; [0, 0] - [0, 3] markdown
-          (info_string ; [0, 3] - [0, 6] markdown
-            (language)) ; [0, 3] - [0, 6] markdown
-          (block_continuation) ; [1, 0] - [1, 0] markdown
-          (code_fence_content ; [1, 0] - [2, 0] markdown
-            (return_statement) ; [1, 0] - [1, 6] lua
-            (block_continuation)) ; [2, 0] - [2, 0] markdown
-          (fenced_code_block_delimiter))) ; [2, 0] - [2, 3] markdown
+      (document ; [0, 0] - [4, 0] markdown
+        (section ; [0, 0] - [4, 0] markdown
+          (fenced_code_block ; [0, 0] - [3, 0] markdown
+            (fenced_code_block_delimiter) ; [0, 0] - [0, 3] markdown
+            (info_string ; [0, 3] - [0, 6] markdown
+              (language)) ; [0, 3] - [0, 6] markdown
+            (block_continuation) ; [1, 0] - [1, 0] markdown
+            (code_fence_content ; [1, 0] - [2, 0] markdown
+              (chunk ; [1, 0] - [2, 0] lua
+                (return_statement)) ; [1, 0] - [1, 6] lua
+              (block_continuation)) ; [2, 0] - [2, 0] markdown
+            (fenced_code_block_delimiter)))) ; [2, 0] - [2, 3] markdown
       ]]
   end)
 end)