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.

2 files changed, 24 insertions, 1 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?