WIP

Co-authored-by: Kirill <kirill@zed.dev>

crates/language/src/buffer.rs

@@ -66,6 +66,7 @@ pub use text::{
     Transaction, TransactionId, Unclipped,
 };
 use theme::{ActiveTheme as _, SyntaxTheme};
+
 #[cfg(any(test, feature = "test-support"))]
 use util::RandomCharIter;
 use util::{debug_panic, maybe, RangeExt};
@@ -475,12 +476,236 @@ struct IndentSuggestion {
 }
 
 struct BufferChunkHighlights<'a> {
+    highlighters: Vec<Box<dyn Highlighter<'a>>>,
+}
+
+trait Highlighter<'a>: 'a {
+    fn narrow(&mut self, range: Range<usize>);
+    // Return starting offset of the next capture.
+    fn seek_next_capture_offset(&mut self, current_offset: usize);
+    fn current_capture(&self) -> Option<(Range<usize>, HighlightId)>;
+    fn advance(&mut self);
+    fn set_byte_range(&mut self, range: Range<usize>);
+}
+
+struct TreeSitterHighlights<'a> {
     captures: SyntaxMapCaptures<'a>,
     next_capture: Option<SyntaxMapCapture<'a>>,
-    stack: Vec<(usize, HighlightId)>,
+    stack: Vec<(Range<usize>, HighlightId)>,
     highlight_maps: Vec<HighlightMap>,
 }
 
+impl<'a> TreeSitterHighlights<'a> {
+    fn new(captures: SyntaxMapCaptures<'a>, highlight_maps: Vec<HighlightMap>) -> Self {
+        TreeSitterHighlights {
+            captures,
+            next_capture: None,
+            stack: Default::default(),
+            highlight_maps,
+        }
+    }
+}
+
+impl<'a> Highlighter<'a> for TreeSitterHighlights<'a> {
+    /// Preserve the existing highlights only if they fall within a provided range.
+    fn narrow(&mut self, range: Range<usize>) {
+        self.stack
+            .retain(|(offset_range, _)| offset_range.end > range.start);
+        if let Some(capture) = &self.next_capture {
+            let next_capture_start = capture.node.start_byte();
+            if range.start >= next_capture_start {
+                let next_capture_end = capture.node.end_byte();
+                if range.start < next_capture_end {
+                    self.stack.push((
+                        next_capture_start..next_capture_end,
+                        self.highlight_maps[capture.grammar_index].get(capture.index),
+                    ));
+                }
+                self.next_capture.take();
+            }
+        }
+    }
+
+    fn seek_next_capture_offset(&mut self, current_offset: usize) {
+        while let Some((parent_range, _)) = self.stack.last() {
+            if parent_range.end <= current_offset {
+                self.stack.pop();
+            } else {
+                break;
+            }
+        }
+
+        if self.next_capture.is_none() {
+            self.next_capture = self.captures.next();
+        }
+
+        while let Some(capture) = self.next_capture.as_ref() {
+            let start_byte = capture.node.start_byte();
+            if current_offset < start_byte {
+                break;
+            } else {
+                let highlight_id = self.highlight_maps[capture.grammar_index].get(capture.index);
+                self.stack
+                    .push((start_byte..capture.node.end_byte(), highlight_id));
+                self.next_capture = self.captures.next();
+            }
+        }
+    }
+    fn set_byte_range(&mut self, range: Range<usize>) {
+        self.captures.set_byte_range(range);
+    }
+    fn current_capture(&self) -> Option<(Range<usize>, HighlightId)> {
+        self.stack.last().cloned()
+    }
+    fn advance(&mut self) {
+        self.stack.pop();
+    }
+}
+
+type OrdRange = (usize, usize);
+struct RainbowBracketsHighlighter<'a> {
+    matches: SyntaxMapMatches<'a>,
+
+    colors: BTreeMap<OrdRange, HighlightId>,
+    /// Index of next color to use.
+    palette_index: usize,
+    palette: Vec<HighlightId>,
+}
+
+impl<'a> RainbowBracketsHighlighter<'a> {
+    fn new(_palette: HighlightMap, matches: SyntaxMapMatches<'a>) -> Option<Self> {
+        // if palette.is_empty() {
+        //     return None;
+        // }
+        // let palette = HighlightMap::new()
+        Some(RainbowBracketsHighlighter {
+            matches,
+            colors: BTreeMap::new(),
+            palette_index: 0,
+            palette: vec![HighlightId(0), HighlightId(1)],
+        })
+    }
+    fn next_highlight_id(palette: &[HighlightId], palette_index: &mut usize) -> HighlightId {
+        let next_highlight_id = palette.get(*palette_index);
+        *palette_index = (*palette_index + 1) % palette.len();
+        *(next_highlight_id.unwrap())
+    }
+}
+impl<'a> Highlighter<'a> for RainbowBracketsHighlighter<'a> {
+    fn narrow(&mut self, range: Range<usize>) {
+        self.colors
+            .retain(|(start, end), _| range.contains(start) || range.contains(end));
+        // if let Some(capture) = self.matches.peek() {
+        //     if range.start >= capture.node.start_byte() {
+        //         let next_capture_end = capture.node.end_byte();
+        //         if range.start < next_capture_end {
+        //             self.stack.push((
+        //                 next_capture_end,
+        //                 self.highlight_maps[capture.grammar_index].get(capture.index),
+        //             ));
+        //         }
+        //         self.next_capture.take();
+        //     }
+        // }
+    }
+
+    fn seek_next_capture_offset(&mut self, current_offset: usize) {
+        // First let's fetch captures up until the opening brace is past the current_offset
+        // The ending brace might fall within visible range.
+
+        if self
+            .colors
+            .iter()
+            .next()
+            .map_or(true, |((start_byte, _), _)| {
+                dbg!(*start_byte) < dbg!(current_offset)
+            })
+        {
+            while let Some(matches) = self.matches.peek() {
+                let [start, end] = matches.captures else {
+                    break;
+                };
+                let start_byte;
+                let highlight_id = Self::next_highlight_id(&self.palette, &mut self.palette_index);
+                {
+                    let Range { start, end } = start.node.byte_range();
+                    start_byte = start;
+                    self.colors.insert((start, end), highlight_id);
+                }
+                let Range { start, end } = end.node.byte_range();
+                self.colors.insert((start, end), highlight_id);
+                if !dbg!(self.matches.advance()) {
+                    break;
+                }
+                if start_byte > current_offset {
+                    break;
+                }
+            }
+        }
+
+        // Now let's throw away anything that starts before current_offset
+        self.colors
+            .retain(|(start, end), _| *start > current_offset || *end > current_offset);
+    }
+    fn set_byte_range(&mut self, range: Range<usize>) {
+        self.matches.set_byte_range(range);
+    }
+
+    fn current_capture(&self) -> Option<(Range<usize>, HighlightId)> {
+        dbg!(&self.colors);
+        self.colors
+            .iter()
+            .next()
+            .map(|((start_offset, end_offset), highlight_id)| {
+                (*start_offset..*end_offset, *highlight_id)
+            })
+    }
+    fn advance(&mut self) {
+        self.colors.pop_first();
+    }
+}
+
+impl<'a> BufferChunkHighlights<'a> {
+    fn new(highlighters: Vec<Box<dyn Highlighter<'a>>>) -> Self {
+        BufferChunkHighlights { highlighters }
+    }
+    /// Preserve the existing highlights only if they fall within a provided range.
+    fn narrow(&mut self, range: Range<usize>) {
+        self.highlighters.iter_mut().for_each(|highlighter| {
+            highlighter.narrow(range.clone());
+        });
+    }
+    /// Preserve the existing highlights only if they fall within a provided range.
+    fn set_byte_range(&mut self, range: Range<usize>) {
+        self.highlighters.iter_mut().for_each(|highlighter| {
+            highlighter.set_byte_range(range.clone());
+        });
+    }
+
+    fn seek_next_capture_offset(&mut self, current_offset: usize) {
+        for highlighter in &mut self.highlighters {
+            highlighter.seek_next_capture_offset(current_offset);
+        }
+    }
+    fn current_capture(&mut self) -> Option<(Range<usize>, HighlightId)> {
+        self.highlighters
+            .iter_mut()
+            .min_by_key(|highlighter| {
+                highlighter
+                    .current_capture()
+                    .map_or((usize::MAX, usize::MAX), |(range, _)| {
+                        (range.start, range.end)
+                    })
+            })
+            .and_then(|highlighter| {
+                let ret = highlighter.current_capture();
+                highlighter.advance();
+                dbg!(&ret);
+                ret
+            })
+    }
+}
+
 /// An iterator that yields chunks of a buffer's text, along with their
 /// syntax highlights and diagnostic status.
 pub struct BufferChunks<'a> {
@@ -734,11 +959,13 @@ impl EditPreview {
             .iter()
             .map(|grammar| grammar.highlight_map())
             .collect();
-
+        let matches = syntax_snapshot.matches(range.clone(), snapshot, |grammar| {
+            grammar.brackets_config.as_ref().map(|config| &config.query)
+        });
         BufferChunks::new(
             snapshot.as_rope(),
             range,
-            Some((captures, highlight_maps)),
+            Some((captures, highlight_maps, matches)),
             false,
             None,
         )
@@ -2954,8 +3181,11 @@ impl BufferSnapshot {
         None
     }
 
-    fn get_highlights(&self, range: Range<usize>) -> (SyntaxMapCaptures, Vec<HighlightMap>) {
-        let captures = self.syntax.captures(range, &self.text, |grammar| {
+    fn get_highlights(
+        &self,
+        range: Range<usize>,
+    ) -> (SyntaxMapCaptures, Vec<HighlightMap>, SyntaxMapMatches) {
+        let captures = self.syntax.captures(range.clone(), &self.text, |grammar| {
             grammar.highlights_query.as_ref()
         });
         let highlight_maps = captures
@@ -2963,7 +3193,10 @@ impl BufferSnapshot {
             .iter()
             .map(|grammar| grammar.highlight_map())
             .collect();
-        (captures, highlight_maps)
+        let matches = self.syntax.matches(range, &self.text, |grammar| {
+            grammar.brackets_config.as_ref().map(|config| &config.query)
+        });
+        (captures, highlight_maps, matches)
     }
 
     /// Iterates over chunks of text in the given range of the buffer. Text is chunked
@@ -4067,18 +4300,27 @@ impl<'a> BufferChunks<'a> {
     pub(crate) fn new(
         text: &'a Rope,
         range: Range<usize>,
-        syntax: Option<(SyntaxMapCaptures<'a>, Vec<HighlightMap>)>,
+        syntax: Option<(
+            SyntaxMapCaptures<'a>,
+            Vec<HighlightMap>,
+            SyntaxMapMatches<'a>,
+        )>,
         diagnostics: bool,
         buffer_snapshot: Option<&'a BufferSnapshot>,
     ) -> Self {
         let mut highlights = None;
-        if let Some((captures, highlight_maps)) = syntax {
-            highlights = Some(BufferChunkHighlights {
-                captures,
-                next_capture: None,
-                stack: Default::default(),
-                highlight_maps,
-            })
+        if let Some((captures, highlight_maps, matches)) = syntax {
+            let highlighter = highlight_maps.iter().next().cloned();
+            // let mut highlighters: Vec<Box<dyn Highlighter>> = vec![Box::new(
+            //     TreeSitterHighlights::new(captures, highlight_maps),
+            // )];
+            let mut highlighters: Vec<Box<dyn Highlighter>> = vec![];
+            highlighters.extend(highlighter.and_then(|highlighter| {
+                RainbowBracketsHighlighter::new(highlighter, matches)
+                    .map(|brackets| Box::new(brackets) as Box<dyn Highlighter>)
+            }));
+
+            highlights = Some(BufferChunkHighlights::new(highlighters));
         }
 
         let diagnostic_endpoints = diagnostics.then(|| Vec::new().into_iter().peekable());
@@ -4107,36 +4349,22 @@ impl<'a> BufferChunks<'a> {
         if let Some(highlights) = self.highlights.as_mut() {
             if old_range.start <= self.range.start && old_range.end >= self.range.end {
                 // Reuse existing highlights stack, as the new range is a subrange of the old one.
-                highlights
-                    .stack
-                    .retain(|(end_offset, _)| *end_offset > range.start);
-                if let Some(capture) = &highlights.next_capture {
-                    if range.start >= capture.node.start_byte() {
-                        let next_capture_end = capture.node.end_byte();
-                        if range.start < next_capture_end {
-                            highlights.stack.push((
-                                next_capture_end,
-                                highlights.highlight_maps[capture.grammar_index].get(capture.index),
-                            ));
-                        }
-                        highlights.next_capture.take();
-                    }
-                }
+                highlights.narrow(range);
             } else if let Some(snapshot) = self.buffer_snapshot {
-                let (captures, highlight_maps) = snapshot.get_highlights(self.range.clone());
-                *highlights = BufferChunkHighlights {
-                    captures,
-                    next_capture: None,
-                    stack: Default::default(),
-                    highlight_maps,
-                };
+                let (captures, highlight_maps, matches) =
+                    snapshot.get_highlights(self.range.clone());
+                let highlighter = highlight_maps.iter().next().cloned().unwrap();
+                *highlights = BufferChunkHighlights::new(vec![
+                    Box::new(TreeSitterHighlights::new(captures, highlight_maps)),
+                    Box::new(RainbowBracketsHighlighter::new(highlighter, matches).unwrap()),
+                ]);
             } else {
                 // We cannot obtain new highlights for a language-aware buffer iterator, as we don't have a buffer snapshot.
                 // Seeking such BufferChunks is not supported.
                 debug_assert!(false, "Attempted to seek on a language-aware buffer iterator without associated buffer snapshot");
             }
 
-            highlights.captures.set_byte_range(self.range.clone());
+            highlights.set_byte_range(self.range.clone());
             self.initialize_diagnostic_endpoints();
         }
     }
@@ -4225,34 +4453,15 @@ impl<'a> Iterator for BufferChunks<'a> {
     type Item = Chunk<'a>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let mut next_capture_start = usize::MAX;
         let mut next_diagnostic_endpoint = usize::MAX;
-
+        let mut next_capture_start = usize::MAX;
+        let mut current_capture = None;
         if let Some(highlights) = self.highlights.as_mut() {
-            while let Some((parent_capture_end, _)) = highlights.stack.last() {
-                if *parent_capture_end <= self.range.start {
-                    highlights.stack.pop();
-                } else {
-                    break;
-                }
-            }
-
-            if highlights.next_capture.is_none() {
-                highlights.next_capture = highlights.captures.next();
-            }
-
-            while let Some(capture) = highlights.next_capture.as_ref() {
-                if self.range.start < capture.node.start_byte() {
-                    next_capture_start = capture.node.start_byte();
-                    break;
-                } else {
-                    let highlight_id =
-                        highlights.highlight_maps[capture.grammar_index].get(capture.index);
-                    highlights
-                        .stack
-                        .push((capture.node.end_byte(), highlight_id));
-                    highlights.next_capture = highlights.captures.next();
-                }
+            highlights.seek_next_capture_offset(self.range.start);
+            current_capture = highlights.current_capture();
+            dbg!(&current_capture, &self.range);
+            if let Some(next_capture) = &current_capture {
+                next_capture_start = next_capture.0.start;
             }
         }
 
@@ -4272,13 +4481,14 @@ impl<'a> Iterator for BufferChunks<'a> {
 
         if let Some(chunk) = self.chunks.peek() {
             let chunk_start = self.range.start;
+
             let mut chunk_end = (self.chunks.offset() + chunk.len())
                 .min(next_capture_start)
                 .min(next_diagnostic_endpoint);
             let mut highlight_id = None;
-            if let Some(highlights) = self.highlights.as_ref() {
-                if let Some((parent_capture_end, parent_highlight_id)) = highlights.stack.last() {
-                    chunk_end = chunk_end.min(*parent_capture_end);
+            if let Some(highlights) = self.highlights.as_mut() {
+                if let Some((parent_range, parent_highlight_id)) = &current_capture {
+                    chunk_end = chunk_end.min(parent_range.end);
                     highlight_id = Some(*parent_highlight_id);
                 }
             }

crates/language/src/buffer_tests.rs

@@ -3118,6 +3118,31 @@ fn test_trailing_whitespace_ranges(mut rng: StdRng) {
     );
 }
 
+#[gpui::test]
+async fn test_rainbow_brackets(cx: &mut TestAppContext) {
+    cx.update(|cx| {
+        init_settings(cx, |_| {});
+        theme::init(theme::LoadThemes::JustBase, cx);
+    });
+
+    static BUFFER_TEXT: &str = r#"
+    use crate::{bar};
+    "#;
+
+    let rust_lang = Arc::new(rust_lang().with_highlights_query("").unwrap());
+    rust_lang.set_theme(&SyntaxTheme {
+        highlights: vec![("foo".to_string(), HighlightStyle::default())],
+    });
+    let buffer = cx.new(|cx| Buffer::local(BUFFER_TEXT, cx).with_language(rust_lang, cx));
+    buffer.update(cx, |buffer, cx| {
+        let snapshot = buffer.snapshot();
+        let chunks = snapshot
+            .chunks(0..BUFFER_TEXT.len(), true)
+            .collect::<Vec<_>>();
+        dbg!(&chunks);
+    });
+}
+
 fn ruby_lang() -> Language {
     Language::new(
         LanguageConfig {

crates/language/src/highlight_map.rs

@@ -48,6 +48,13 @@ impl HighlightMap {
             .copied()
             .unwrap_or(DEFAULT_SYNTAX_HIGHLIGHT_ID)
     }
+
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
 }
 
 impl HighlightId {

crates/language/src/language.rs

@@ -1485,11 +1485,19 @@ impl Language {
                 SyntaxSnapshot::single_tree_captures(range.clone(), text, &tree, self, |grammar| {
                     grammar.highlights_query.as_ref()
                 });
+            let matches =
+                SyntaxSnapshot::single_tree_matches(range.clone(), text, &tree, self, |grammar| {
+                    grammar.brackets_config.as_ref().map(|config| &config.query)
+                });
             let highlight_maps = vec![grammar.highlight_map()];
             let mut offset = 0;
-            for chunk in
-                BufferChunks::new(text, range, Some((captures, highlight_maps)), false, None)
-            {
+            for chunk in BufferChunks::new(
+                text,
+                range,
+                Some((captures, highlight_maps, matches)),
+                false,
+                None,
+            ) {
                 let end_offset = offset + chunk.text.len();
                 if let Some(highlight_id) = chunk.syntax_highlight_id {
                     if !highlight_id.is_default() {

crates/language/src/syntax_map.rs

@@ -789,6 +789,28 @@ impl SyntaxSnapshot {
         )
     }
 
+    pub fn single_tree_matches<'a>(
+        range: Range<usize>,
+        text: &'a Rope,
+        tree: &'a Tree,
+        language: &'a Arc<Language>,
+        query: fn(&Grammar) -> Option<&Query>,
+    ) -> SyntaxMapMatches<'a> {
+        SyntaxMapMatches::new(
+            range.clone(),
+            text,
+            [SyntaxLayer {
+                language,
+                tree,
+                depth: 0,
+                offset: (0, tree_sitter::Point::new(0, 0)),
+            }]
+            .into_iter(),
+            query,
+            TreeSitterOptions::default(),
+        )
+    }
+
     pub fn captures<'a>(
         &'a self,
         range: Range<usize>,
@@ -1118,9 +1140,7 @@ impl<'a> SyntaxMapMatches<'a> {
     }
 
     pub fn advance(&mut self) -> bool {
-        let layer = if let Some(layer) = self.layers.first_mut() {
-            layer
-        } else {
+        let Some(layer) = self.layers.first_mut() else {
             return false;
         };
 
@@ -1139,6 +1159,23 @@ impl<'a> SyntaxMapMatches<'a> {
 
         true
     }
+    pub fn set_byte_range(&mut self, range: Range<usize>) {
+        // for layer in &mut self.layers {
+        //     layer.matches.set_byte_range(range.clone());
+        //     if let Some(capture) = &layer.next {
+        //         if capture.node.end_byte() > range.start {
+        //             continue;
+        //         }
+        //     }
+        //     layer.advance();
+        // }
+        // self.layers.sort_unstable_by_key(|layer| layer.sort_key());
+        // self.active_layer_count = self
+        //     .layers
+        //     .iter()
+        //     .position(|layer| layer.next_capture.is_none())
+        //     .unwrap_or(self.layers.len());
+    }
 }
 
 impl<'a> SyntaxMapCapturesLayer<'a> {