Fix initial highlight layer sort order (#5196)

The purpose of this change is to remove the mutable self borrow on
`HighlightIterLayer::sort_key` so that we can sort layers with the
correct ordering using the `Vec::sort` function family.
`HighlightIterLayer::sort_key` needs `&mut self` since it calls
`Peekable::peek` which needs `&mut self`. `Vec::sort` functions
only give immutable borrows of the elements to ensure the
correctness of the sort.

We could instead approach this by creating an eager Peekable and using
that instead of `std::iter::Peekable` to wrap `QueryCaptures`:

```rust
struct EagerPeekable<I: Iterator> {
    iter: I,
    peeked: Option<I::Item>,
}

impl<I: Iterator> EagerPeekable<I> {
    fn new(mut iter: I) -> Self {
        let peeked = iter.next();
        Self { iter, peeked }
    }

    fn peek(&self) -> Option<&I::Item> {
        self.peeked.as_ref()
    }
}

impl<I: Iterator> Iterator for EagerPeekable<I> {
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        std::mem::replace(&mut self.peeked, self.iter.next())
    }
}
```

This would be a cleaner approach (notice how `EagerPeekable::peek`
takes `&self` rather than `&mut self`), however this doesn't work in
practice because the Items emitted by the `tree_sitter::QueryCaptures`
Iterator must be consumed before the next Item is returned.
`Iterator::next` on `tree_sitter::QueryCaptures` modifies the
`QueryMatch` returned by the last call of `next`. This behavior is
not currently reflected in the lifetimes/structure of `QueryCaptures`.

This fixes an issue with layers being out of order when using combined
injections since the old code only checked the first range in the
layer. Layers being out of order could cause missing highlights for
combined-injections content.

1 files changed, 14 insertions, 20 deletions

diff --git a/helix-core/src/syntax.rs b/helix-core/src/syntax.rs
index 41ab23e1..ca4da3dc 100644
--- a/helix-core/src/syntax.rs
+++ b/helix-core/src/syntax.rs
@@ -1092,21 +1092,14 @@ impl Syntax {
                     }],
                     cursor,
                     _tree: None,
-                    captures,
+                    captures: RefCell::new(captures),
                     config: layer.config.as_ref(), // TODO: just reuse `layer`
                     depth: layer.depth,            // TODO: just reuse `layer`
-                    ranges: &layer.ranges,         // TODO: temp
                 })
             })
             .collect::<Vec<_>>();
 
-        // HAXX: arrange layers by byte range, with deeper layers positioned first
-        layers.sort_by_key(|layer| {
-            (
-                layer.ranges.first().cloned(),
-                std::cmp::Reverse(layer.depth),
-            )
-        });
+        layers.sort_unstable_by_key(|layer| layer.sort_key());
 
         let mut result = HighlightIter {
             source,
@@ -1424,12 +1417,11 @@ impl<'a> TextProvider<'a> for RopeProvider<'a> {
 struct HighlightIterLayer<'a> {
     _tree: Option<Tree>,
     cursor: QueryCursor,
-    captures: iter::Peekable<QueryCaptures<'a, 'a, RopeProvider<'a>>>,
+    captures: RefCell<iter::Peekable<QueryCaptures<'a, 'a, RopeProvider<'a>>>>,
     config: &'a HighlightConfiguration,
     highlight_end_stack: Vec<usize>,
     scope_stack: Vec<LocalScope<'a>>,
     depth: u32,
-    ranges: &'a [Range],
 }
 
 impl<'a> fmt::Debug for HighlightIterLayer<'a> {
@@ -1610,10 +1602,11 @@ impl<'a> HighlightIterLayer<'a> {
     // First, sort scope boundaries by their byte offset in the document. At a
     // given position, emit scope endings before scope beginnings. Finally, emit
     // scope boundaries from deeper layers first.
-    fn sort_key(&mut self) -> Option<(usize, bool, isize)> {
+    fn sort_key(&self) -> Option<(usize, bool, isize)> {
         let depth = -(self.depth as isize);
         let next_start = self
             .captures
+            .borrow_mut()
             .peek()
             .map(|(m, i)| m.captures[*i].node.start_byte());
         let next_end = self.highlight_end_stack.last().cloned();
@@ -1838,7 +1831,8 @@ impl<'a> Iterator for HighlightIter<'a> {
             // Get the next capture from whichever layer has the earliest highlight boundary.
             let range;
             let layer = &mut self.layers[0];
-            if let Some((next_match, capture_index)) = layer.captures.peek() {
+            let captures = layer.captures.get_mut();
+            if let Some((next_match, capture_index)) = captures.peek() {
                 let next_capture = next_match.captures[*capture_index];
                 range = next_capture.node.byte_range();
 
@@ -1861,7 +1855,7 @@ impl<'a> Iterator for HighlightIter<'a> {
                 return self.emit_event(self.source.len_bytes(), None);
             };
 
-            let (mut match_, capture_index) = layer.captures.next().unwrap();
+            let (mut match_, capture_index) = captures.next().unwrap();
             let mut capture = match_.captures[capture_index];
 
             // Remove from the local scope stack any local scopes that have already ended.
@@ -1937,11 +1931,11 @@ impl<'a> Iterator for HighlightIter<'a> {
                 }
 
                 // Continue processing any additional matches for the same node.
-                if let Some((next_match, next_capture_index)) = layer.captures.peek() {
+                if let Some((next_match, next_capture_index)) = captures.peek() {
                     let next_capture = next_match.captures[*next_capture_index];
                     if next_capture.node == capture.node {
                         capture = next_capture;
-                        match_ = layer.captures.next().unwrap().0;
+                        match_ = captures.next().unwrap().0;
                         continue;
                     }
                 }
@@ -1964,11 +1958,11 @@ impl<'a> Iterator for HighlightIter<'a> {
             // highlighting patterns that are disabled for local variables.
             if definition_highlight.is_some() || reference_highlight.is_some() {
                 while layer.config.non_local_variable_patterns[match_.pattern_index] {
-                    if let Some((next_match, next_capture_index)) = layer.captures.peek() {
+                    if let Some((next_match, next_capture_index)) = captures.peek() {
                         let next_capture = next_match.captures[*next_capture_index];
                         if next_capture.node == capture.node {
                             capture = next_capture;
-                            match_ = layer.captures.next().unwrap().0;
+                            match_ = captures.next().unwrap().0;
                             continue;
                         }
                     }
@@ -1983,10 +1977,10 @@ impl<'a> Iterator for HighlightIter<'a> {
             // for a given node are ordered by pattern index, so these subsequent
             // captures are guaranteed to be for highlighting, not injections or
             // local variables.
-            while let Some((next_match, next_capture_index)) = layer.captures.peek() {
+            while let Some((next_match, next_capture_index)) = captures.peek() {
                 let next_capture = next_match.captures[*next_capture_index];
                 if next_capture.node == capture.node {
-                    layer.captures.next();
+                    captures.next();
                 } else {
                     break;
                 }