Merge pull request #13765 from jvgrootveld/ts-default-start-end-row-on-capture

treesitter: default start and end row when omitted

runtime/doc/treesitter.txt

@@ -195,7 +195,8 @@ query:iter_captures({node}, {bufnr}, {start_row}, {end_row})
 	text of the buffer. {start_row} and {end_row} can be used to limit
 	matches inside a row range (this is typically used with root node
 	as the node, i e to get syntax highlight matches in the current
-	viewport)
+	viewport). When omitted the start and end row values are used from
+        the given node.
 
 	The iterator returns three values, a numeric id identifying the capture,
 	the captured node, and metadata from any directives processing the match.

runtime/lua/vim/treesitter/query.lua

@@ -340,6 +340,19 @@ function Query:apply_directives(match, pattern, source, metadata)
   end
 end
 
+
+--- Returns the start and stop value if set else the node's range.
+-- When the node's range is used, the stop is incremented by 1
+-- to make the search inclusive.
+local function value_or_node_range(start, stop, node)
+  if start == nil and stop == nil then
+    local node_start, _, node_stop, _ = node:range()
+    return node_start, node_stop + 1 -- Make stop inclusive
+  end
+
+  return start, stop
+end
+
 --- Iterates of the captures of self on a given range.
 --
 -- @param node The node under witch the search will occur
@@ -353,6 +366,9 @@ function Query:iter_captures(node, source, start, stop)
   if type(source) == "number" and source == 0 then
     source = vim.api.nvim_get_current_buf()
   end
+
+  start, stop = value_or_node_range(start, stop, node)
+
   local raw_iter = node:_rawquery(self.query, true, start, stop)
   local function iter()
     local capture, captured_node, match = raw_iter()
@@ -385,6 +401,9 @@ function Query:iter_matches(node, source, start, stop)
   if type(source) == "number" and source == 0 then
     source = vim.api.nvim_get_current_buf()
   end
+
+  start, stop = value_or_node_range(start, stop, node)
+
   local raw_iter = node:_rawquery(self.query, false, start, stop)
   local function iter()
     local pattern, match = raw_iter()

test/functional/treesitter/parser_spec.lua

@@ -245,7 +245,7 @@ void ui_refresh(void)
       parser = vim.treesitter.get_parser(0, "c")
       tree = parser:parse()[1]
       res = {}
-      for pattern, match in cquery:iter_matches(tree:root(), 0, 0, 1) do
+      for pattern, match in cquery:iter_matches(tree:root(), 0) do
         -- can't transmit node over RPC. just check the name and range
         local mrepr = {}
         for cid,node in pairs(match) do
@@ -289,7 +289,7 @@ void ui_refresh(void)
     local query = query.parse_query("c", ...)
 
     local nodes = {}
-    for _, node in query:iter_captures(parser:parse()[1]:root(), 0, 0, 19) do
+    for _, node in query:iter_captures(parser:parse()[1]:root(), 0) do
       table.insert(nodes, {node:range()})
     end
 
@@ -365,7 +365,7 @@ void ui_refresh(void)
     query = vim.treesitter.parse_query("c", "(declaration) @decl")
 
     local nodes = {}
-    for _, node in query:iter_captures(parser:parse()[1]:root(), 0, 0, 19) do
+    for _, node in query:iter_captures(parser:parse()[1]:root(), 0) do
       table.insert(nodes, node)
     end
 
@@ -412,7 +412,7 @@ void ui_refresh(void)
     local nodes = {}
     local query = vim.treesitter.parse_query("c", '((identifier) @id (eq? @id "foo"))')
 
-    for _, node in query:iter_captures(parser:parse()[1]:root(), str, 0, 2) do
+    for _, node in query:iter_captures(parser:parse()[1]:root(), str) do
       table.insert(nodes, { node:range() })
     end
 
@@ -421,6 +421,29 @@ void ui_refresh(void)
     eq({ {0, 10, 0, 13} }, ret)
   end)
 
+  it("should use node range when omitted", function()
+    local txt = [[
+      int foo = 42;
+      int bar = 13;
+    ]]
+
+    local ret = exec_lua([[
+    local str = ...
+    local parser = vim.treesitter.get_string_parser(str, "c")
+
+    local nodes = {}
+    local query = vim.treesitter.parse_query("c", '((identifier) @foo)')
+    local first_child = parser:parse()[1]:root():child(1)
+
+    for _, node in query:iter_captures(first_child, str) do
+      table.insert(nodes, { node:range() })
+    end
+
+    return nodes]], txt)
+
+    eq({ {1, 10, 1, 13} }, ret)
+  end)
+
   describe("when creating a language tree", function()
     local function get_ranges()
       return exec_lua([[
@@ -539,7 +562,7 @@ int x = INT_MAX;
         query = vim.treesitter.parse_query("c", '((number_literal) @number (#set! "key" "value"))')
         parser = vim.treesitter.get_parser(0, "c")
 
-        for pattern, match, metadata in query:iter_matches(parser:parse()[1]:root(), 0, 0, 1) do
+        for pattern, match, metadata in query:iter_matches(parser:parse()[1]:root(), 0) do
           result = metadata.key
         end
 
@@ -562,7 +585,7 @@ int x = INT_MAX;
         query = vim.treesitter.parse_query("c", '((number_literal) @number (#set! @number "key" "value"))')
         parser = vim.treesitter.get_parser(0, "c")
 
-        for pattern, match, metadata in query:iter_matches(parser:parse()[1]:root(), 0, 0, 1) do
+        for pattern, match, metadata in query:iter_matches(parser:parse()[1]:root(), 0) do
           result = metadata[pattern].key
         end