1 files changed, 9 insertions, 5 deletions
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index b0b0fecd38..66ab0d52f0 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -913,8 +913,8 @@ end
 ---@param start? integer Starting line for the search. Defaults to `node:start()`.
 ---@param stop? integer Stopping line for the search (end-exclusive). Defaults to `node:end_()`.
 ---
----@return (fun(end_line: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch):
----        capture id, capture node, metadata, match
+---@return (fun(end_line: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch, TSTree):
+---        capture id, capture node, metadata, match, tree
 ---
 ---@note Captures are only returned if the query pattern of a specific capture contained predicates.
 function Query:iter_captures(node, source, start, stop)
@@ -924,6 +924,8 @@ function Query:iter_captures(node, source, start, stop)
 
   start, stop = value_or_node_range(start, stop, node)
 
+  -- Copy the tree to ensure it is valid during the entire lifetime of the iterator
+  local tree = node:tree():copy()
   local cursor = vim._create_ts_querycursor(node, self.query, start, stop, { match_limit = 256 })
 
   -- For faster checks that a match is not in the cache.
@@ -970,7 +972,7 @@ function Query:iter_captures(node, source, start, stop)
       match_cache[match_id] = metadata
     end
 
-    return capture, captured_node, metadata, match
+    return capture, captured_node, metadata, match, tree
   end
   return iter
 end
@@ -1011,7 +1013,7 @@ end
 ---   (last) node instead of the full list of matching nodes. This option is only for backward
 ---   compatibility and will be removed in a future release.
 ---
----@return (fun(): integer, table<integer, TSNode[]>, vim.treesitter.query.TSMetadata): pattern id, match, metadata
+---@return (fun(): integer, table<integer, TSNode[]>, vim.treesitter.query.TSMetadata, TSTree): pattern id, match, metadata, tree
 function Query:iter_matches(node, source, start, stop, opts)
   opts = opts or {}
   opts.match_limit = opts.match_limit or 256
@@ -1022,6 +1024,8 @@ function Query:iter_matches(node, source, start, stop, opts)
 
   start, stop = value_or_node_range(start, stop, node)
 
+  -- Copy the tree to ensure it is valid during the entire lifetime of the iterator
+  local tree = node:tree():copy()
   local cursor = vim._create_ts_querycursor(node, self.query, start, stop, opts)
 
   local function iter()
@@ -1059,7 +1063,7 @@ function Query:iter_matches(node, source, start, stop, opts)
     end
 
     -- TODO(lewis6991): create a new function that returns {match, metadata}
-    return pattern_i, captures, metadata
+    return pattern_i, captures, metadata, tree
   end
   return iter
 end