Refactor to be able to pass a query

crates/highlight/src/c_lib.rs

@@ -297,7 +297,7 @@ impl TSHighlighter {
                     })
                 })
             },
-            &|_| true,
+            &|_, _| true,
         );
 
         if let Ok(highlights) = highlights {

crates/highlight/src/highlight.rs

@@ -165,7 +165,7 @@ struct LocalScope<'a> {
 struct HighlightIter<'a, F, G>
 where
     F: FnMut(&str) -> Option<&'a HighlightConfiguration> + 'a,
-    G: Fn(&(tree_sitter::QueryMatch<'a, 'a>, usize)) -> bool,
+    G: Fn(&QueryMatch, &Query) -> bool + 'a,
 {
     source: &'a [u8],
     language_name: &'a str,
@@ -174,19 +174,16 @@ where
     injection_callback: F,
     capture_filter: &'a G,
     cancellation_flag: Option<&'a AtomicUsize>,
-    layers: Vec<HighlightIterLayer<'a, G>>,
+    layers: Vec<HighlightIterLayer<'a>>,
     iter_count: usize,
     next_event: Option<HighlightEvent>,
     last_highlight_range: Option<(usize, usize, usize)>,
 }
 
-struct HighlightIterLayer<'a, G>
-where
-    G: Fn(&(tree_sitter::QueryMatch<'a, 'a>, usize)) -> bool,
-{
+struct HighlightIterLayer<'a> {
     _tree: Tree,
     cursor: QueryCursor,
-    captures: iter::Peekable<iter::Filter<_QueryCaptures<'a, 'a, &'a [u8], &'a [u8]>, &'a G>>,
+    captures: iter::Peekable<Box<dyn Iterator<Item = (QueryMatch<'a, 'a>, usize)> + 'a>>,
     config: &'a HighlightConfiguration,
     highlight_end_stack: Vec<usize>,
     scope_stack: Vec<LocalScope<'a>>,
@@ -294,7 +291,7 @@ impl Highlighter {
         source: &'a [u8],
         cancellation_flag: Option<&'a AtomicUsize>,
         mut injection_callback: impl FnMut(&str) -> Option<&'a HighlightConfiguration> + 'a,
-        query_filter: &'a impl Fn(&(QueryMatch, usize)) -> bool,
+        query_filter: &'a impl Fn(&QueryMatch, &Query) -> bool,
     ) -> Result<impl Iterator<Item = Result<HighlightEvent, Error>> + 'a, Error> {
         let layers = HighlightIterLayer::new(
             source,
@@ -512,17 +509,17 @@ impl HighlightConfiguration {
     }
 }
 
-impl<'a, G> HighlightIterLayer<'a, G>
-where
-    G: Fn(&(tree_sitter::QueryMatch<'a, 'a>, usize)) -> bool,
-{
+impl<'a> HighlightIterLayer<'a> {
     /// Create a new 'layer' of highlighting for this document.
     ///
     /// In the event that the new layer contains "combined injections" (injections where multiple
     /// disjoint ranges are parsed as one syntax tree), these will be eagerly processed and
     /// added to the returned vector.
     #[allow(clippy::too_many_arguments)]
-    fn new<F: FnMut(&str) -> Option<&'a HighlightConfiguration> + 'a>(
+    fn new<
+        F: FnMut(&str) -> Option<&'a HighlightConfiguration> + 'a,
+        G: Fn(&QueryMatch, &Query) -> bool,
+    >(
         source: &'a [u8],
         parent_name: Option<&str>,
         highlighter: &mut Highlighter,
@@ -615,13 +612,21 @@ where
                 let cursor_ref = unsafe {
                     mem::transmute::<&mut QueryCursor, &'static mut QueryCursor>(&mut cursor)
                 };
+
                 let captures = unsafe {
-                    std::mem::transmute::<QueryCaptures<_, _>, _QueryCaptures<_, _>>(
-                        cursor_ref.captures(&config.query, tree_ref.root_node(), source),
-                    )
-                }
-                .filter(query_filter)
-                .peekable();
+                    std::mem::transmute::<
+                        QueryCaptures<_, _>,
+                        _QueryCaptures<'a, 'a, &'a [u8], &'a [u8]>,
+                    >(cursor_ref.captures(
+                        &config.query,
+                        tree_ref.root_node(),
+                        source,
+                    ))
+                };
+
+                let captures: Box<dyn Iterator<Item = _>> = Box::new(
+                    captures.filter(|(result, _): &(_, _)| query_filter(result, &config.query)),
+                );
 
                 result.push(HighlightIterLayer {
                     highlight_end_stack: Vec::new(),
@@ -633,7 +638,7 @@ where
                     cursor,
                     depth,
                     _tree: tree,
-                    captures,
+                    captures: captures.peekable(),
                     config,
                     ranges,
                 });
@@ -774,7 +779,7 @@ where
 impl<'a, F, G> HighlightIter<'a, F, G>
 where
     F: FnMut(&str) -> Option<&'a HighlightConfiguration> + 'a,
-    G: Fn(&(tree_sitter::QueryMatch<'a, 'a>, usize)) -> bool,
+    G: Fn(&QueryMatch, &Query) -> bool,
 {
     fn emit_event(
         &mut self,
@@ -819,7 +824,7 @@ where
         }
     }
 
-    fn insert_layer(&mut self, mut layer: HighlightIterLayer<'a, G>) {
+    fn insert_layer(&mut self, mut layer: HighlightIterLayer<'a>) {
         if let Some(sort_key) = layer.sort_key() {
             let mut i = 1;
             while i < self.layers.len() {
@@ -841,7 +846,7 @@ where
 impl<'a, F, G> Iterator for HighlightIter<'a, F, G>
 where
     F: FnMut(&str) -> Option<&'a HighlightConfiguration> + 'a,
-    G: Fn(&(tree_sitter::QueryMatch<'a, 'a>, usize)) -> bool,
+    G: Fn(&QueryMatch, &Query) -> bool,
 {
     type Item = Result<HighlightEvent, Error>;
 
@@ -926,7 +931,7 @@ where
                 // to the highlighted document.
                 if let (Some(language_name), Some(content_node)) = (language_name, content_node) {
                     if let Some(config) = (self.injection_callback)(language_name) {
-                        let ranges = HighlightIterLayer::<G>::intersect_ranges(
+                        let ranges = HighlightIterLayer::intersect_ranges(
                             &self.layers[0].ranges,
                             &[content_node],
                             include_children,