Merge pull request #19946 from bfredl/newnode

feat: upstream some nvim-treesitter functions

4 files changed, 157 insertions, 10 deletions

diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index 06409f9980..c80d8f4479 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -352,6 +352,15 @@ attribute: >
 ==============================================================================
 Lua module: vim.treesitter                               *lua-treesitter-core*
 
+get_node_range({node_or_range})                             *get_node_range()*
+    Get the node's range or unpack a range table
+
+    Parameters: ~
+        {node_or_range}  (table)
+
+    Return: ~
+        start_row, start_col, end_row, end_col
+
 get_parser({bufnr}, {lang}, {opts})                             *get_parser()*
     Gets the parser for this bufnr / ft combination.
 
@@ -374,6 +383,26 @@ get_string_parser({str}, {lang}, {opts})                 *get_string_parser()*
         {lang}  The language of this string
         {opts}  Options to pass to the created language tree
 
+is_ancestor({dest}, {source})                                  *is_ancestor()*
+    Determines whether a node is the ancestor of another
+
+    Parameters: ~
+        {dest}    (table) the possible ancestor
+        {source}  (table) the possible descendant node
+
+    Return: ~
+        (boolean) True if dest is an ancestor of source
+
+node_contains({node}, {range})                               *node_contains()*
+    Determines if a node contains a range
+
+    Parameters: ~
+        {node}   (table) The node
+        {range}  (table) The range
+
+    Return: ~
+        (boolean) True if the node contains the range
+
 
 ==============================================================================
 Lua module: vim.treesitter.language                      *treesitter-language*
@@ -427,12 +456,16 @@ add_predicate({name}, {handler}, {force})                    *add_predicate()*
         {handler}  the handler function to be used signature will be (match,
                    pattern, bufnr, predicate)
 
-get_node_text({node}, {source})                              *get_node_text()*
+get_node_text({node}, {source}, {opts})                      *get_node_text()*
     Gets the text corresponding to a given node
 
     Parameters: ~
-        {node}    the node
-        {source}  The buffer or string from which the node is extracted
+        {node}    (table) The node
+        {source}  (table) The buffer or string from which the node is
+                  extracted
+        {opts}    (table) Optional parameters.
+                  • concat: (boolean default true) Concatenate result in a
+                    string
 
 get_query({lang}, {query_name})                                  *get_query()*
     Returns the runtime query {query_name} for {lang}.
@@ -680,6 +713,17 @@ LanguageTree:language_for_range({self}, {range})
         {range}  A text range, see |LanguageTree:contains|
         {self}
 
+                                         *LanguageTree:named_node_for_range()*
+LanguageTree:named_node_for_range({self}, {range}, {opts})
+    Gets the smallest named node that contains {range}
+
+    Parameters: ~
+        {range}                   (table) A text range
+        {opts}                    (table) Options table
+        {opts.ignore_injections}  (boolean) (default true) Ignore injected
+                                  languages.
+        {self}
+
 LanguageTree:parse({self})                              *LanguageTree:parse()*
     Parses all defined regions using a treesitter parser for the language this
     tree represents. This will run the injection query for this language to
@@ -740,6 +784,17 @@ LanguageTree:source({self})                            *LanguageTree:source()*
     Parameters: ~
         {self}
 
+                                               *LanguageTree:tree_for_range()*
+LanguageTree:tree_for_range({self}, {range}, {opts})
+    Gets the tree that contains {range}
+
+    Parameters: ~
+        {range}                   (table) A text range
+        {opts}                    (table) Options table
+        {opts.ignore_injections}  (boolean) (default true) Ignore injected
+                                  languages.
+        {self}
+
 LanguageTree:trees({self})                              *LanguageTree:trees()*
     Returns all trees this language tree contains. Does not include child
     languages.
diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 70f2c425ed..82d41070ee 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -118,4 +118,53 @@ function M.get_string_parser(str, lang, opts)
   return LanguageTree.new(str, lang, opts)
 end
 
+--- Determines whether a node is the ancestor of another
+---
+---@param dest table the possible ancestor
+---@param source table the possible descendant node
+---
+---@returns (boolean) True if dest is an ancestor of source
+function M.is_ancestor(dest, source)
+  if not (dest and source) then
+    return false
+  end
+
+  local current = source
+  while current ~= nil do
+    if current == dest then
+      return true
+    end
+
+    current = current:parent()
+  end
+
+  return false
+end
+
+--- Get the node's range or unpack a range table
+---
+---@param node_or_range table
+---
+---@returns start_row, start_col, end_row, end_col
+function M.get_node_range(node_or_range)
+  if type(node_or_range) == 'table' then
+    return unpack(node_or_range)
+  else
+    return node_or_range:range()
+  end
+end
+
+---Determines if a node contains a range
+---@param node table The node
+---@param range table The range
+---
+---@returns (boolean) True if the node contains the range
+function M.node_contains(node, range)
+  local start_row, start_col, end_row, end_col = node:range()
+  local start_fits = start_row < range[1] or (start_row == range[1] and start_col <= range[2])
+  local end_fits = end_row > range[3] or (end_row == range[3] and end_col >= range[4])
+
+  return start_fits and end_fits
+end
+
 return M
diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index 4d3b0631a2..70317a9f94 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -299,7 +299,7 @@ function LanguageTree:included_regions()
 end
 
 ---@private
-local function get_node_range(node, id, metadata)
+local function get_range_from_metadata(node, id, metadata)
   if metadata[id] and metadata[id].range then
     return metadata[id].range
   end
@@ -362,7 +362,7 @@ function LanguageTree:_get_injections()
         elseif name == 'combined' then
           combined = true
         elseif name == 'content' and #ranges == 0 then
-          table.insert(ranges, get_node_range(node, id, metadata))
+          table.insert(ranges, get_range_from_metadata(node, id, metadata))
           -- Ignore any tags that start with "_"
           -- Allows for other tags to be used in matches
         elseif string.sub(name, 1, 1) ~= '_' then
@@ -371,7 +371,7 @@ function LanguageTree:_get_injections()
           end
 
           if #ranges == 0 then
-            table.insert(ranges, get_node_range(node, id, metadata))
+            table.insert(ranges, get_range_from_metadata(node, id, metadata))
           end
         end
       end
@@ -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/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 103e85abfd..697e2e7691 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -181,9 +181,14 @@ end
 
 --- Gets the text corresponding to a given node
 ---
----@param node the node
----@param source The buffer or string from which the node is extracted
-function M.get_node_text(node, source)
+---@param node table The node
+---@param source table The buffer or string from which the node is extracted
+---@param opts table Optional parameters.
+---          - concat: (boolean default true) Concatenate result in a string
+function M.get_node_text(node, source, opts)
+  opts = opts or {}
+  local concat = vim.F.if_nil(opts.concat, true)
+
   local start_row, start_col, start_byte = node:start()
   local end_row, end_col, end_byte = node:end_()
 
@@ -210,7 +215,7 @@ function M.get_node_text(node, source)
       end
     end
 
-    return table.concat(lines, '\n')
+    return concat and table.concat(lines, '\n') or lines
   elseif type(source) == 'string' then
     return source:sub(start_byte + 1, end_byte)
   end