feat(treesitter): add ability to retreive a tree/node given a range

2 files changed, 67 insertions, 0 deletions

diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index 4d3b0631a2..f87a66ddab 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -549,6 +549,44 @@ function LanguageTree:contains(range)
   return false
 end
 
+--- Gets the tree that contains {range}
+---
+---@param range table A text range
+---@param opts table Options table
+---@param opts.ignore_injections boolean (default true) Ignore injected languages.
+function LanguageTree:tree_for_range(range, opts)
+  opts = opts or {}
+  local ignore = vim.F.if_nil(opts.ignore_injections, true)
+
+  if not ignore then
+    for _, child in pairs(self._children) do
+      for _, tree in pairs(child:trees()) do
+        if tree_contains(tree, range) then
+          return tree
+        end
+      end
+    end
+  end
+
+  for _, tree in pairs(self._trees) do
+    if tree_contains(tree, range) then
+      return tree
+    end
+  end
+
+  return nil
+end
+
+--- Gets the smallest named node that contains {range}
+---
+---@param range table A text range
+---@param opts table Options table
+---@param opts.ignore_injections boolean (default true) Ignore injected languages.
+function LanguageTree:named_node_for_range(range, opts)
+  local tree = self:tree_for_range(range, opts)
+  return tree:root():named_descendant_for_range(unpack(range))
+end
+
 --- Gets the appropriate language that contains {range}
 ---
 ---@param range A text range, see |LanguageTree:contains|
diff --git a/test/functional/treesitter/language_spec.lua b/test/functional/treesitter/language_spec.lua
index 8aa8524a26..b66a56264f 100644
--- a/test/functional/treesitter/language_spec.lua
+++ b/test/functional/treesitter/language_spec.lua
@@ -7,6 +7,7 @@ local exec_lua = helpers.exec_lua
 local pcall_err = helpers.pcall_err
 local matches = helpers.matches
 local pending_c_parser = helpers.pending_c_parser
+local insert = helpers.insert
 
 before_each(clear)
 
@@ -84,5 +85,33 @@ describe('treesitter language API', function()
     eq("Error executing lua: .../language.lua:0: no parser for 'borklang' language, see :help treesitter-parsers",
        pcall_err(exec_lua, "new_parser = vim.treesitter.get_parser(0)"))
   end)
+
+  it('retrieve the tree given a range', function ()
+    insert([[
+      int main() {
+        int x = 3;
+      }]])
+
+    exec_lua([[
+      langtree = vim.treesitter.get_parser(0, "c")
+      tree = langtree:tree_for_range({1, 3, 1, 3})
+    ]])
+
+    eq('<node translation_unit>', exec_lua('return tostring(tree:root())'))
+  end)
+
+  it('retrieve the node given a range', function ()
+    insert([[
+      int main() {
+        int x = 3;
+      }]])
+
+    exec_lua([[
+      langtree = vim.treesitter.get_parser(0, "c")
+      node = langtree:named_node_for_range({1, 3, 1, 3})
+    ]])
+
+    eq('<node primitive_type>', exec_lua('return tostring(node)'))
+  end)
 end)