treesitter: default start and end row when omitted

Add support for default start and end row when omitted in the
query:iter_captures and query:iter_matches functions.

When the start and end row values are omitted, the values of the given
node is used. The end row value is incremented by 1 to include the node end
row in the match.

Updated tests and docs accordingly.

3 files changed, 43 insertions, 7 deletions

diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index 911e7b8b47..1696d3b9ba 100644
--- a/runtime/doc/treesitter.txt
+++ b/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.
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 5a27d740a2..9d1f265e54 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -353,6 +353,12 @@ function Query:iter_captures(node, source, start, stop)
   if type(source) == "number" and source == 0 then
     source = vim.api.nvim_get_current_buf()
   end
+
+  if start == nil and stop == nil then
+    start, _, stop, _ = node:range()
+    stop = stop + 1 -- Make stop inclusive
+  end
+
   local raw_iter = node:_rawquery(self.query, true, start, stop)
   local function iter()
     local capture, captured_node, match = raw_iter()
@@ -385,6 +391,12 @@ function Query:iter_matches(node, source, start, stop)
   if type(source) == "number" and source == 0 then
     source = vim.api.nvim_get_current_buf()
   end
+
+  if start == nil and stop == nil then
+    start, _, stop, _ = node:range()
+    stop = stop + 1 -- Make stop inclusive
+  end
+
   local raw_iter = node:_rawquery(self.query, false, start, stop)
   local function iter()
     local pattern, match = raw_iter()
diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua
index 520574c08a..f99362fbdf 100644
--- a/test/functional/treesitter/parser_spec.lua
+++ b/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