From cc2531428e78d11fd182787eba177e890ea5fdac Mon Sep 17 00:00:00 2001
From: Michael Davis <mcarsondavis@gmail.com>
Date: Wed, 1 Feb 2023 15:50:12 -0600
Subject: [PATCH] Share sort_key function between query iterators

This is brings the fix from d5f17d3 to the QueryIter layers. A trait
for getting the cursor and sort-key from each layer helps cut down on
code duplicated between these iterators.
---
 helix-core/src/syntax.rs | 51 ++++++++++++++++++++++++++++++++++------
 1 file changed, 44 insertions(+), 7 deletions(-)

diff --git a/helix-core/src/syntax.rs b/helix-core/src/syntax.rs
index e80f50ff7..74950c740 100644
--- a/helix-core/src/syntax.rs
+++ b/helix-core/src/syntax.rs
@@ -1425,7 +1425,7 @@ impl Syntax {
     where
         F: Fn(&'a HighlightConfiguration) -> &'a Query,
     {
-        let layers = self
+        let mut layers: Vec<_> = self
             .layers
             .iter()
             .filter_map(|(_, layer)| {
@@ -1457,12 +1457,14 @@ impl Syntax {
 
                 Some(QueryIterLayer {
                     cursor,
-                    captures,
+                    captures: RefCell::new(captures),
                     layer,
                 })
             })
             .collect();
 
+        layers.sort_unstable_by_key(|layer| layer.sort_key());
+
         QueryIter { layers }
     }
 
@@ -2282,11 +2284,21 @@ impl HighlightConfiguration {
     }
 }
 
-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(&self) -> Option<(usize, bool, isize)> {
+trait IterLayer {
+    type SortKey: PartialOrd;
+
+    fn sort_key(&self) -> Option<Self::SortKey>;
+
+    fn cursor(self) -> QueryCursor;
+}
+
+impl<'a> IterLayer for HighlightIterLayer<'a> {
+    type SortKey = (usize, bool, isize);
+
+    fn sort_key(&self) -> Option<Self::SortKey> {
+        // 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.
         let depth = -(self.depth as isize);
         let next_start = self
             .captures
@@ -2307,6 +2319,30 @@ impl<'a> HighlightIterLayer<'a> {
             _ => None,
         }
     }
+
+    fn cursor(self) -> QueryCursor {
+        self.cursor
+    }
+}
+
+impl<'a> IterLayer for QueryIterLayer<'a> {
+    type SortKey = (usize, isize);
+
+    fn sort_key(&self) -> Option<Self::SortKey> {
+        // Sort the layers so that the first layer in the Vec has the next
+        // capture ordered by start byte and depth (descending).
+        let depth = -(self.layer.depth as isize);
+        let mut captures = self.captures.borrow_mut();
+        let (match_, capture_index) = captures.peek()?;
+        let start = match_.captures[*capture_index].node.start_byte();
+
+        Some((start, depth))
+    }
+
+    fn cursor(self) -> QueryCursor {
+        self.cursor
+    }
+}
 }
 
 #[derive(Clone)]
@@ -2969,6 +3005,7 @@ impl<'a> Iterator for QueryIter<'a> {
         let inner = layer.layer;
         layer
             .captures
+            .borrow_mut()
             .next()
             .map(|(match_, index)| (inner, match_, index))
     }