fix: update add_predicate and add_directive calls for upstream (#6106)

Update custom predicates and directives to handle multiple nodes per
capture ID per changes upstream.

1 files changed, 35 insertions, 14 deletions

diff --git a/plugin/query_predicates.lua b/plugin/query_predicates.lua
index 7cd8fb618..06bf957f8 100644
--- a/plugin/query_predicates.lua
+++ b/plugin/query_predicates.lua
@@ -1,19 +1,44 @@
 local query = vim.treesitter.query
 
+local predicates = {
+  ---@param match TSQueryMatch
+  ---@param pred string[]
+  ---@param any boolean
+  ---@return boolean
+  ['kind-eq'] = function(match, pred, any)
+    local nodes = match[pred[2]]
+    if not nodes or #nodes == 0 then
+      return true
+    end
+
+    local types = { unpack(pred, 3) }
+    for _, node in ipairs(nodes) do
+      local res = vim.list_contains(types, node:type())
+      if any and res then
+        return true
+      elseif not any and not res then
+        return false
+      end
+    end
+    return not any
+  end,
+}
+
 -- register custom predicates (overwrite existing; needed for CI)
 
----@param match (TSNode|nil)[]
+---@param match TSQueryMatch
 ---@param pred string[]
 ---@return boolean|nil
 query.add_predicate('kind-eq?', function(match, _, _, pred)
-  local node = match[pred[2]]
-  if not node then
-    return true
-  end
+  return predicates['kind-eq'](match, pred, false)
+end, { force = true })
 
-  local types = { unpack(pred, 3) }
-  return vim.list_contains(types, node:type())
-end, true)
+---@param match TSQueryMatch
+---@param pred string[]
+---@return boolean|nil
+query.add_predicate('any-kind-eq?', function(match, _, _, pred)
+  return predicates['kind-eq'](match, pred, true)
+end, { force = true })
 
 -- register custom directives
 
@@ -24,7 +49,7 @@ local mimetype_aliases = {
   ['text/ecmascript'] = 'javascript',
 }
 
----@param match (TSNode|nil)[]
+---@param match TSQueryMatch
 ---@param _ string
 ---@param bufnr integer
 ---@param pred string[]
@@ -32,10 +57,6 @@ local mimetype_aliases = {
 query.add_directive('set-lang-from-mimetype!', function(match, _, bufnr, pred, metadata)
   local id = pred[2]
   local node = match[id]
-  if not node then
-    return
-  end
-
   local type_attr_value = vim.treesitter.get_node_text(node, bufnr, { metadata = metadata[id] })
   local configured = mimetype_aliases[type_attr_value]
   if configured then
@@ -44,4 +65,4 @@ query.add_directive('set-lang-from-mimetype!', function(match, _, bufnr, pred, m
     local parts = vim.split(type_attr_value, '/', {})
     metadata['injection.language'] = parts[#parts]
   end
-end, true)
+end, { force = true })