2 files changed, 27 insertions, 10 deletions
diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index 9c620c422c..c864fe5878 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -425,23 +425,21 @@ function LanguageTree:register_cbs(cbs)
   end
 end
 
-local function region_contains(region, range)
-  for _, node in ipairs(region) do
-    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])
+local function tree_contains(tree, range)
+  local start_row, start_col, end_row, end_col = tree:root():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])
 
-    if start_fits and end_fits then
-      return true
-    end
+  if start_fits and end_fits then
+    return true
   end
 
   return false
 end
 
 function LanguageTree:contains(range)
-  for _, region in pairs(self._regions) do
-    if region_contains(region, range) then
+  for _, tree in pairs(self._trees) do
+    if tree_contains(tree, range) then
       return true
     end
   end
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 5a27d740a2..e49f54681d 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/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()