From 458b5de0fc99d25403e5e440d905a45924301210 Mon Sep 17 00:00:00 2001
From: Amaan Qureshi <amaanq12@gmail.com>
Date: Thu, 20 Jul 2023 04:41:47 -0400
Subject: [PATCH] feat: add `any-of` predicate

---
 cli/src/tests/query_test.rs | 17 ++++++++++++++
 lib/binding_rust/lib.rs     | 44 +++++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/cli/src/tests/query_test.rs b/cli/src/tests/query_test.rs
index c3bf54a2..c0994d31 100644
--- a/cli/src/tests/query_test.rs
+++ b/cli/src/tests/query_test.rs
@@ -2842,6 +2842,14 @@ fn test_query_captures_with_text_conditions() {
             ((identifier) @function.builtin
              (#eq? @function.builtin "require"))
 
+            ((identifier) @variable.builtin
+              (#any-of? @variable.builtin
+                        "arguments"
+                        "module"
+                        "console"
+                        "window"
+                        "document"))
+
             ((identifier) @variable
              (#not-match? @variable "^(lambda|load)$"))
             "#,
@@ -2855,6 +2863,9 @@ fn test_query_captures_with_text_conditions() {
           lambda
           const ab = require('./ab');
           new Cd(EF);
+          document;
+          module;
+          console;
         ";
 
         let mut parser = Parser::new();
@@ -2876,6 +2887,12 @@ fn test_query_captures_with_text_conditions() {
                 ("constant", "EF"),
                 ("constructor", "EF"),
                 ("variable", "EF"),
+                ("variable.builtin", "document"),
+                ("variable", "document"),
+                ("variable.builtin", "module"),
+                ("variable", "module"),
+                ("variable.builtin", "console"),
+                ("variable", "console"),
             ],
         );
     });
diff --git a/lib/binding_rust/lib.rs b/lib/binding_rust/lib.rs
index b259aacf..932fc452 100644
--- a/lib/binding_rust/lib.rs
+++ b/lib/binding_rust/lib.rs
@@ -251,6 +251,7 @@ enum TextPredicate {
     CaptureEqString(u32, String, bool),
     CaptureEqCapture(u32, u32, bool),
     CaptureMatchString(u32, regex::bytes::Regex, bool),
+    CaptureAnyString(u32, Vec<String>, bool),
 }
 
 // TODO: Remove this struct at at some point. If `core::str::lossy::Utf8Lossy`
@@ -1811,6 +1812,38 @@ impl Query {
                         operator_name == "is?",
                     )),
 
+                    "any-of?" | "not-any-of?" => {
+                        if p.len() < 2 {
+                            return Err(predicate_error(row, format!(
+                                "Wrong number of arguments to #any-of? predicate. Expected at least 1, got {}.",
+                                p.len() - 1
+                            )));
+                        }
+                        if p[1].type_ != type_capture {
+                            return Err(predicate_error(row, format!(
+                                "First argument to #any-of? predicate must be a capture name. Got literal \"{}\".",
+                                string_values[p[1].value_id as usize],
+                            )));
+                        }
+
+                        let is_positive = operator_name == "any-of?";
+                        let mut values = Vec::new();
+                        for arg in &p[2..] {
+                            if arg.type_ == type_capture {
+                                return Err(predicate_error(row, format!(
+                                    "Arguments to #any-of? predicate must be literals. Got capture @{}.",
+                                    result.capture_names[arg.value_id as usize],
+                                )));
+                            }
+                            values.push(string_values[arg.value_id as usize].clone());
+                        }
+                        text_predicates.push(TextPredicate::CaptureAnyString(
+                            p[1].value_id,
+                            values,
+                            is_positive,
+                        ));
+                    }
+
                     _ => general_predicates.push(QueryPredicate {
                         operator: operator_name.clone().into_boxed_str(),
                         args: p[1..]
@@ -2265,6 +2298,17 @@ impl<'tree> QueryMatch<'_, 'tree> {
                         None => true,
                     }
                 }
+                TextPredicate::CaptureAnyString(i, v, is_positive) => {
+                    let node = self.nodes_for_capture_index(*i).next();
+                    match node {
+                        Some(node) => {
+                            let mut text = text_provider.text(node);
+                            let text = node_text1.get_text(&mut text);
+                            v.iter().any(|s| text == s.as_bytes()) == *is_positive
+                        }
+                        None => true,
+                    }
+                }
             })
     }
 }