treesitter: add predicate "any-of?" (#14344)

For the case of Clojure and other Lisp syntax highlighting, it is
necessary to create huge regexps consisting of hundreds of symbols with
the pipe (|) character. To make things more difficult, these Lisp
symbols sometimes consists of special characters that are themselves
part of special regexp characters like '*'. In addition to being
difficult to maintain, it's performance is suboptimal.

This patch introduces a new predicate to perform 'source' matching in
amortized constant time. This is accomplished by compiling a hash table
on the first use.

3 files changed, 84 insertions, 2 deletions

diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index 1f4b5d3097..39522898f9 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -212,6 +212,11 @@ Here is a list of built-in predicates :
 			((identifier) @foo (#contains? @foo "foo"))
 			((identifier) @foo-bar (#contains @foo-bar "foo" "bar"))
 <
+	`any-of?`					*ts-predicate-any-of?*
+		Will check if the text is the same as any of the following
+		arguments : >
+			((identifier) @foo (#any-of? @foo "foo" "bar"))
+<
 							*lua-treesitter-not-predicate*
 Each predicate has a `not-` prefixed predicate that is just the negation of
 the predicate.
diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index c0140f9186..b81eb18945 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -260,7 +260,25 @@ local predicate_handlers = {
     end
 
     return false
-  end
+  end,
+
+  ["any-of?"] = function(match, _, source, predicate)
+    local node = match[predicate[2]]
+    local node_text = M.get_node_text(node, source)
+
+    -- Since 'predicate' will not be used by callers of this function, use it
+    -- to store a string set built from the list of words to check against.
+    local string_set = predicate["string_set"]
+    if not string_set then
+      string_set = {}
+      for i=3,#predicate do
+        string_set[predicate[i]] = true
+      end
+      predicate["string_set"] = string_set
+    end
+
+    return string_set[node_text]
+  end,
 }
 
 -- As we provide lua-match? also expose vim-match?
diff --git a/test/functional/treesitter/parser_spec.lua b/test/functional/treesitter/parser_spec.lua
index 1017913709..d2f9148e8f 100644
--- a/test/functional/treesitter/parser_spec.lua
+++ b/test/functional/treesitter/parser_spec.lua
@@ -270,6 +270,65 @@ void ui_refresh(void)
     }, res)
   end)
 
+  it('supports builtin query predicate any-of?', function()
+    if pending_c_parser(pending) then return end
+
+    insert([[
+      #include <stdio.h>
+
+      int main(void) {
+        int i;
+        for(i=1; i<=100; i++) {
+          if(((i%3)||(i%5))== 0)
+            printf("number= %d FizzBuzz\n", i);
+          else if((i%3)==0)
+            printf("number= %d Fizz\n", i);
+          else if((i%5)==0)
+            printf("number= %d Buzz\n", i);
+          else
+            printf("number= %d\n",i);
+        }
+        return 0;
+      }
+    ]])
+    exec_lua([[
+      function get_query_result(query_text)
+        cquery = vim.treesitter.parse_query("c", query_text)
+        parser = vim.treesitter.get_parser(0, "c")
+        tree = parser:parse()[1]
+        res = {}
+        for cid, node in cquery:iter_captures(tree:root(), 0) do
+          -- can't transmit node over RPC. just check the name, range, and text
+          local text = vim.treesitter.get_node_text(node, 0)
+          local range = {node:range()}
+          table.insert(res, {cquery.captures[cid], node:type(), range, text})
+        end
+        return res
+      end
+    ]])
+
+    local res0 = exec_lua([[return get_query_result(...)]],
+      [[((primitive_type) @c-keyword (#any-of? @c-keyword "int" "float"))]])
+    eq({
+      { "c-keyword", "primitive_type", { 2, 2, 2, 5 }, "int" },
+      { "c-keyword", "primitive_type", { 3, 4, 3, 7 }, "int" },
+    }, res0)
+
+    local res1 = exec_lua([[return get_query_result(...)]],
+      [[
+        ((string_literal) @fizzbuzz-strings (#any-of? @fizzbuzz-strings
+          "\"number= %d FizzBuzz\\n\""
+          "\"number= %d Fizz\\n\""
+          "\"number= %d Buzz\\n\""
+        ))
+      ]])
+    eq({
+      { "fizzbuzz-strings", "string_literal", { 6, 15, 6, 38 }, "\"number= %d FizzBuzz\\n\""},
+      { "fizzbuzz-strings", "string_literal", { 8, 15, 8, 34 }, "\"number= %d Fizz\\n\""},
+      { "fizzbuzz-strings", "string_literal", { 10, 15, 10, 34 }, "\"number= %d Buzz\\n\""},
+    }, res1)
+  end)
+
   it('allow loading query with escaped quotes and capture them with `lua-match?` and `vim-match?`', function()
     if pending_c_parser(pending) then return end
 
@@ -343,7 +402,7 @@ void ui_refresh(void)
     return list
     ]]
 
-    eq({ 'contains?', 'eq?', 'is-main?', 'lua-match?', 'match?', 'vim-match?' }, res_list)
+    eq({ 'any-of?', 'contains?', 'eq?', 'is-main?', 'lua-match?', 'match?', 'vim-match?' }, res_list)
   end)