fix(treesitter): return correct match table in iter_captures()

2 files changed, 33 insertions, 27 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
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 075fd0e99b..e68acac929 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -1,5 +1,6 @@
 local api = vim.api
 local language = require('vim.treesitter.language')
+local memoize = vim.func._memoize
 
 local M = {}
 
@@ -212,7 +213,7 @@ end
 ---@param query_name string Name of the query (e.g. "highlights")
 ---
 ---@return vim.treesitter.Query? : Parsed query. `nil` if no query files are found.
-M.get = vim.func._memoize('concat-2', function(lang, query_name)
+M.get = memoize('concat-2', function(lang, query_name)
   if explicit_queries[lang][query_name] then
     return explicit_queries[lang][query_name]
   end
@@ -245,7 +246,7 @@ end)
 ---@return vim.treesitter.Query : Parsed query
 ---
 ---@see [vim.treesitter.query.get()]
-M.parse = vim.func._memoize('concat-2', function(lang, query)
+M.parse = memoize('concat-2', function(lang, query)
   language.add(lang)
 
   local ts_query = vim._ts_parse_query(lang, query)
@@ -812,6 +813,12 @@ local function value_or_node_range(start, stop, node)
   return start, stop
 end
 
+--- @param match TSQueryMatch
+--- @return integer
+local function match_id_hash(_, match)
+  return (match:info())
+end
+
 --- Iterate over all captures from all matches inside {node}
 ---
 --- {source} is needed if the query contains predicates; then the caller
@@ -841,7 +848,7 @@ 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, table<integer,TSNode[]>?):
+---@return (fun(end_line: integer|nil): integer, TSNode, vim.treesitter.query.TSMetadata, TSQueryMatch):
 ---        capture id, capture node, metadata, match
 ---
 ---@note Captures are only returned if the query pattern of a specific capture contained predicates.
@@ -854,7 +861,8 @@ function Query:iter_captures(node, source, start, stop)
 
   local cursor = vim._create_ts_querycursor(node, self.query, start, stop, { match_limit = 256 })
 
-  local max_match_id = -1
+  local apply_directives = memoize(match_id_hash, self.apply_directives, true)
+  local match_preds = memoize(match_id_hash, self.match_preds, true)
 
   local function iter(end_line)
     local capture, captured_node, match = cursor:next_capture()
@@ -863,27 +871,18 @@ function Query:iter_captures(node, source, start, stop)
       return
     end
 
-    local captures --- @type table<integer,TSNode[]>?
-    local match_id, pattern_index = match:info()
-
-    local metadata = {}
-
-    local preds = self.info.patterns[pattern_index] or {}
-
-    if #preds > 0 and match_id > max_match_id then
-      captures = match:captures()
-      max_match_id = match_id
-      if not self:match_preds(match, source) then
-        cursor:remove_match(match_id)
-        if end_line and captured_node:range() > end_line then
-          return nil, captured_node, nil
-        end
-        return iter(end_line) -- tail call: try next match
+    if not match_preds(self, match, source) then
+      local match_id = match:info()
+      cursor:remove_match(match_id)
+      if end_line and captured_node:range() > end_line then
+        return nil, captured_node, nil, nil
       end
-
-      metadata = self:apply_directives(match, source)
+      return iter(end_line) -- tail call: try next match
     end
-    return capture, captured_node, metadata, captures
+
+    local metadata = apply_directives(self, match, source)
+
+    return capture, captured_node, metadata, match
   end
   return iter
 end
@@ -972,6 +971,7 @@ function Query:iter_matches(node, source, start, stop, opts)
       return pattern, old_match, metadata
     end
 
+    -- TODO(lewis6991): create a new function that returns {match, metadata}
     return pattern, captures, metadata
   end
   return iter