1 files changed, 10 insertions, 4 deletions
diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
index 1e6f128461..3f7e31212c 100644
--- a/runtime/lua/vim/treesitter/highlighter.lua
+++ b/runtime/lua/vim/treesitter/highlighter.lua
@@ -4,7 +4,7 @@ local Range = require('vim.treesitter._range')
 
 local ns = api.nvim_create_namespace('treesitter/highlighter')
 
----@alias vim.treesitter.highlighter.Iter fun(end_line: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, table<integer, TSNode[]>
+---@alias vim.treesitter.highlighter.Iter fun(end_line: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch
 
 ---@class (private) vim.treesitter.highlighter.Query
 ---@field private _query vim.treesitter.Query?
@@ -243,7 +243,7 @@ function TSHighlighter:get_query(lang)
   return self._queries[lang]
 end
 
---- @param match table<integer,TSNode[]>
+--- @param match TSQueryMatch
 --- @param bufnr integer
 --- @param capture integer
 --- @param metadata vim.treesitter.query.TSMetadata
@@ -256,13 +256,15 @@ local function get_url(match, bufnr, capture, metadata)
     return url
   end
 
-  if not match or not match[url] then
+  local captures = match:captures()
+
+  if not captures[url] then
     return
   end
 
   -- Assume there is only one matching node. If there is more than one, take the URL
   -- from the first.
-  local other_node = match[url][1]
+  local other_node = captures[url][1]
 
   return vim.treesitter.get_node_text(other_node, bufnr, {
     metadata = metadata[url],
@@ -296,6 +298,10 @@ local function on_line_impl(self, buf, line, is_spell_nav)
     end
 
     if state.iter == nil or state.next_row < line then
+      -- Mainly used to skip over folds
+
+      -- TODO(lewis6991): Creating a new iterator loses the cached predicate results for query
+      -- matches. Move this logic inside iter_captures() so we can maintain the cache.
       state.iter =
         state.highlighter_query:query():iter_captures(root_node, self.bufnr, line, root_end_row + 1)
     end