rust: add handling of arbitrary predicate operators

cli/src/tests/query_test.rs

@@ -2,7 +2,8 @@ use super::helpers::allocations;
 use super::helpers::fixtures::get_language;
 use std::fmt::Write;
 use tree_sitter::{
-    Node, Parser, Query, QueryCapture, QueryCursor, QueryError, QueryMatch, QueryProperty,
+    Node, Parser, Query, QueryCapture, QueryCursor, QueryError, QueryMatch, QueryPredicate,
+    QueryPredicateArg, QueryProperty,
 };
 
 #[test]
@@ -948,7 +949,7 @@ fn test_query_captures_with_text_conditions() {
 }
 
 #[test]
-fn test_query_captures_with_set_properties() {
+fn test_query_captures_with_predicates() {
     allocations::record(|| {
         let language = get_language("javascript");
 
@@ -957,7 +958,8 @@ fn test_query_captures_with_set_properties() {
             r#"
             ((call_expression (identifier) @foo)
              (set! name something)
-             (set! cool))
+             (set! cool)
+             (something! @foo omg))
 
             ((property_identifier) @bar
              (is? cool)
@@ -972,6 +974,16 @@ fn test_query_captures_with_set_properties() {
                 QueryProperty::new("cool", None, None),
             ]
         );
+        assert_eq!(
+            query.general_predicates(0),
+            &[QueryPredicate {
+                operator: "something!".to_string().into_boxed_str(),
+                args: vec![
+                    QueryPredicateArg::Capture(0),
+                    QueryPredicateArg::String("omg".to_string().into_boxed_str()),
+                ],
+            },]
+        );
         assert_eq!(query.property_settings(1), &[]);
         assert_eq!(query.property_predicates(0), &[]);
         assert_eq!(
@@ -985,7 +997,7 @@ fn test_query_captures_with_set_properties() {
 }
 
 #[test]
-fn test_query_captures_with_set_quoted_properties() {
+fn test_query_captures_with_quoted_predicate_args() {
     allocations::record(|| {
         let language = get_language("javascript");
 

lib/binding_rust/lib.rs

@@ -100,6 +100,7 @@ pub struct Query {
     text_predicates: Vec<Box<[TextPredicate]>>,
     property_settings: Vec<Box<[QueryProperty]>>,
     property_predicates: Vec<Box<[(QueryProperty, bool)]>>,
+    general_predicates: Vec<Box<[QueryPredicate]>>,
 }
 
 /// A stateful object for executing a `Query` on a syntax `Tree`.
@@ -113,6 +114,19 @@ pub struct QueryProperty {
     pub capture_id: Option<usize>,
 }
 
+#[derive(Debug, PartialEq, Eq)]
+pub enum QueryPredicateArg {
+    Capture(u32),
+    String(Box<str>),
+}
+
+/// A key-value pair associated with a particular pattern in a `Query`.
+#[derive(Debug, PartialEq, Eq)]
+pub struct QueryPredicate {
+    pub operator: Box<str>,
+    pub args: Vec<QueryPredicateArg>,
+}
+
 /// A match of a `Query` to a particular set of `Node`s.
 pub struct QueryMatch<'a> {
     pub pattern_index: usize,
@@ -1199,6 +1213,7 @@ impl Query {
             text_predicates: Vec::with_capacity(pattern_count),
             property_predicates: Vec::with_capacity(pattern_count),
             property_settings: Vec::with_capacity(pattern_count),
+            general_predicates: Vec::with_capacity(pattern_count),
         };
 
         // Build a vector of strings to store the capture names.
@@ -1242,6 +1257,7 @@ impl Query {
             let mut text_predicates = Vec::new();
             let mut property_predicates = Vec::new();
             let mut property_settings = Vec::new();
+            let mut general_predicates = Vec::new();
             for p in predicate_steps.split(|s| s.type_ == type_done) {
                 if p.is_empty() {
                     continue;
@@ -1333,12 +1349,21 @@ impl Query {
                         operator_name == "is?",
                     )),
 
-                    _ => {
-                        return Err(QueryError::Predicate(format!(
-                            "Unknown query predicate function {}",
-                            operator_name,
-                        )))
-                    }
+                    _ => general_predicates.push(QueryPredicate {
+                        operator: operator_name.clone().into_boxed_str(),
+                        args: p[1..]
+                            .iter()
+                            .map(|a| {
+                                if a.type_ == type_capture {
+                                    QueryPredicateArg::Capture(a.value_id)
+                                } else {
+                                    QueryPredicateArg::String(
+                                        string_values[a.value_id as usize].clone().into_boxed_str(),
+                                    )
+                                }
+                            })
+                            .collect(),
+                    }),
                 }
             }
 
@@ -1351,6 +1376,9 @@ impl Query {
             result
                 .property_settings
                 .push(property_settings.into_boxed_slice());
+            result
+                .general_predicates
+                .push(general_predicates.into_boxed_slice());
         }
         Ok(result)
     }
@@ -1380,15 +1408,30 @@ impl Query {
     }
 
     /// Get the properties that are checked for the given pattern index.
+    ///
+    /// This includes predicates with the operators `is?` and `is-not?`.
     pub fn property_predicates(&self, index: usize) -> &[(QueryProperty, bool)] {
         &self.property_predicates[index]
     }
 
     /// Get the properties that are set for the given pattern index.
+    ///
+    /// This includes predicates with the operator `set!`.
     pub fn property_settings(&self, index: usize) -> &[QueryProperty] {
         &self.property_settings[index]
     }
 
+    /// Get the other user-defined predicates associated with the given index.
+    ///
+    /// This includes predicate with operators other than:
+    /// * `match?`
+    /// * `eq?` and `not-eq?
+    /// * `is?` and `is-not?`
+    /// * `set!`
+    pub fn general_predicates(&self, index: usize) -> &[QueryPredicate] {
+        &self.general_predicates[index]
+    }
+
     /// Disable a certain capture within a query.
     ///
     /// This prevents the capture from being returned in matches, and also avoids any