Allow to match arbitrary predicates

Cargo.lock

@@ -55,6 +55,12 @@ version = "1.0.15"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c"
 
+[[package]]
+name = "lua-patterns"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2f3da90ef01f4466adefe593c05d534fe270d04a2b067608cf8b230123e5a281"
+
 [[package]]
 name = "memchr"
 version = "2.7.6"
@@ -218,6 +224,7 @@ dependencies = [
 name = "tree-sitter-highlight"
 version = "0.26.3"
 dependencies = [
+ "lua-patterns",
  "regex",
  "streaming-iterator",
  "thiserror",

crates/highlight/Cargo.toml

@@ -29,3 +29,4 @@ thiserror.workspace = true
 streaming-iterator.workspace = true
 
 tree-sitter = "0.26"
+lua-patterns = "0.4.0"

crates/highlight/src/highlight.rs

@@ -16,11 +16,13 @@ use std::{
 };
 
 pub use c_lib as c;
+use lua_patterns::LuaPattern;
 use streaming_iterator::StreamingIterator;
 use thiserror::Error;
 use tree_sitter::{
     ffi, Language, LossyUtf8, Node, ParseOptions, Parser, Point, Query, QueryCapture,
-    QueryCaptures, QueryCursor, QueryError, QueryMatch, Range, TextProvider, Tree,
+    QueryCaptures, QueryCursor, QueryError, QueryMatch, QueryPredicateArg, Range, TextProvider,
+    Tree,
 };
 
 const CANCELLATION_CHECK_INTERVAL: usize = 100;
@@ -225,6 +227,37 @@ impl<'tree> _QueryMatch<'_, 'tree> {
     }
 }
 
+struct NodeText<'a, T> {
+    buffer: &'a mut Vec<u8>,
+    first_chunk: Option<T>,
+}
+impl<'a, T: AsRef<[u8]>> NodeText<'a, T> {
+    const fn new(buffer: &'a mut Vec<u8>) -> Self {
+        Self {
+            buffer,
+            first_chunk: None,
+        }
+    }
+
+    fn get_text(&mut self, chunks: &mut impl Iterator<Item = T>) -> &[u8] {
+        self.first_chunk = chunks.next();
+        if let Some(next_chunk) = chunks.next() {
+            self.buffer.clear();
+            self.buffer
+                .extend_from_slice(self.first_chunk.as_ref().unwrap().as_ref());
+            self.buffer.extend_from_slice(next_chunk.as_ref());
+            for chunk in chunks {
+                self.buffer.extend_from_slice(chunk.as_ref());
+            }
+            self.buffer.as_slice()
+        } else if let Some(ref first_chunk) = self.first_chunk {
+            first_chunk.as_ref()
+        } else {
+            &[]
+        }
+    }
+}
+
 impl<'query, 'tree: 'query, T: TextProvider<I>, I: AsRef<[u8]>> Iterator
     for _QueryCaptures<'query, 'tree, T, I>
 {
@@ -244,14 +277,51 @@ impl<'query, 'tree: 'query, T: TextProvider<I>, I: AsRef<[u8]>> Iterator
                         &m.assume_init(),
                         self.ptr,
                     ));
-                    if result.satisfies_text_predicates(
+
+                    let mut matches = true;
+
+                    // Check that we don’t have any is/is-not (Not handled)
+                    let props = self.query.property_predicates(result.pattern_index);
+                    assert_eq!(props, [], "Unhandled is/is-not?");
+
+                    let mut node_text = NodeText::new(&mut self.buffer1);
+
+                    'pred: for predicate in self.query.general_predicates(result.pattern_index) {
+                        match &*predicate.operator {
+                            "lua-match?" => {
+                                let [QueryPredicateArg::Capture(capture), QueryPredicateArg::String(pattern)] =
+                                    &*predicate.args
+                                else {
+                                    panic!("Unexpected arguments: {:?}", predicate.args);
+                                };
+
+                                let mut matcher = LuaPattern::new(pattern);
+                                for node in result.nodes_for_capture_index(*capture) {
+                                    let mut text = self.text_provider.text(node);
+                                    let text = node_text.get_text(&mut text);
+                                    if !matcher.matches_bytes(text) {
+                                        matches = false;
+                                        break 'pred;
+                                    }
+                                }
+                            }
+                            _ => panic!("Unhandled operator: {}", predicate.operator),
+                        }
+                    }
+
+                    if !result.satisfies_text_predicates(
                         self.query,
                         &mut self.buffer1,
                         &mut self.buffer2,
                         &mut self.text_provider,
                     ) {
+                        matches = false;
+                    }
+
+                    if matches {
                         return Some((result, capture_index as usize));
                     }
+
                     result.remove();
                 } else {
                     return None;