Support bracket colorization (rainbow brackets) (#43172)

Deals with https://github.com/zed-industries/zed/issues/5259

Highlights brackets with different colors based on their depth.
Uses existing tree-sitter queries from brackets.scm to find brackets,
uses theme's accents to color them.


https://github.com/user-attachments/assets/cc5f3aba-22fa-446d-9af7-ba6e772029da

1. Adds `colorize_brackets` language setting that allows, per language
or globally for all languages, to configure whether Zed should color the
brackets for a particular language.

Disabled for all languages by default.

2. Any given language can opt-out a certain bracket pair by amending the
brackets.scm like `("\"" @open "\"" @close) ` -> `(("\"" @open "\""
@close) (#set! rainbow.exclude))`

3. Brackets are using colors from theme accents, which can be overridden
as

```jsonc
"theme_overrides": {
  "One Dark": {
    "accents": ["#ff69b4", "#7fff00", "#ff1493", "#00ffff", "#ff8c00", "#9400d3"]
  }
},
```

Release Notes:

- Added bracket colorization (rainbow brackets) support. Use
`colorize_brackets` language setting to enable.

---------

Co-authored-by: MrSubidubi <dev@bahn.sh>
Co-authored-by: Lukas Wirth <lukas@zed.dev>
Co-authored-by: MrSubidubi <finn@zed.dev>
Co-authored-by: Lukas Wirth <me@lukaswirth.dev>
Co-authored-by: Smit Barmase <heysmitbarmase@gmail.com>

crates/language/src/buffer.rs

@@ -1,9 +1,12 @@
+pub mod row_chunk;
+
 use crate::{
     DebuggerTextObject, LanguageScope, Outline, OutlineConfig, RunnableCapture, RunnableTag,
     TextObject, TreeSitterOptions,
     diagnostic_set::{DiagnosticEntry, DiagnosticEntryRef, DiagnosticGroup},
     language_settings::{LanguageSettings, language_settings},
     outline::OutlineItem,
+    row_chunk::RowChunks,
     syntax_map::{
         SyntaxLayer, SyntaxMap, SyntaxMapCapture, SyntaxMapCaptures, SyntaxMapMatch,
         SyntaxMapMatches, SyntaxSnapshot, ToTreeSitterPoint,
@@ -18,9 +21,9 @@ pub use crate::{
     proto,
 };
 use anyhow::{Context as _, Result};
-use clock::Lamport;
 pub use clock::ReplicaId;
-use collections::HashMap;
+use clock::{Global, Lamport};
+use collections::{HashMap, HashSet};
 use fs::MTime;
 use futures::channel::oneshot;
 use gpui::{
@@ -28,8 +31,9 @@ use gpui::{
     Task, TaskLabel, TextStyle,
 };
 
+use itertools::Itertools;
 use lsp::{LanguageServerId, NumberOrString};
-use parking_lot::Mutex;
+use parking_lot::{Mutex, RawMutex, lock_api::MutexGuard};
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use settings::WorktreeId;
@@ -45,7 +49,7 @@ use std::{
     iter::{self, Iterator, Peekable},
     mem,
     num::NonZeroU32,
-    ops::{Deref, Range},
+    ops::{Deref, Not, Range},
     path::PathBuf,
     rc,
     sync::{Arc, LazyLock},
@@ -126,6 +130,29 @@ pub struct Buffer {
     has_unsaved_edits: Cell<(clock::Global, bool)>,
     change_bits: Vec<rc::Weak<Cell<bool>>>,
     _subscriptions: Vec<gpui::Subscription>,
+    tree_sitter_data: Arc<Mutex<TreeSitterData>>,
+}
+
+#[derive(Debug, Clone)]
+pub struct TreeSitterData {
+    chunks: RowChunks,
+    brackets_by_chunks: Vec<Option<Vec<BracketMatch<usize>>>>,
+}
+
+const MAX_ROWS_IN_A_CHUNK: u32 = 50;
+
+impl TreeSitterData {
+    fn clear(&mut self) {
+        self.brackets_by_chunks = vec![None; self.chunks.len()];
+    }
+
+    fn new(snapshot: text::BufferSnapshot) -> Self {
+        let chunks = RowChunks::new(snapshot, MAX_ROWS_IN_A_CHUNK);
+        Self {
+            brackets_by_chunks: vec![None; chunks.len()],
+            chunks,
+        }
+    }
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
@@ -149,6 +176,7 @@ pub struct BufferSnapshot {
     remote_selections: TreeMap<ReplicaId, SelectionSet>,
     language: Option<Arc<Language>>,
     non_text_state_update_count: usize,
+    tree_sitter_data: Arc<Mutex<TreeSitterData>>,
 }
 
 /// The kind and amount of indentation in a particular line. For now,
@@ -819,11 +847,18 @@ impl EditPreview {
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
-pub struct BracketMatch {
-    pub open_range: Range<usize>,
-    pub close_range: Range<usize>,
+pub struct BracketMatch<T> {
+    pub open_range: Range<T>,
+    pub close_range: Range<T>,
     pub newline_only: bool,
-    pub depth: usize,
+    pub syntax_layer_depth: usize,
+    pub color_index: Option<usize>,
+}
+
+impl<T> BracketMatch<T> {
+    pub fn bracket_ranges(self) -> (Range<T>, Range<T>) {
+        (self.open_range, self.close_range)
+    }
 }
 
 impl Buffer {
@@ -974,8 +1009,10 @@ impl Buffer {
         let saved_mtime = file.as_ref().and_then(|file| file.disk_state().mtime());
         let snapshot = buffer.snapshot();
         let syntax_map = Mutex::new(SyntaxMap::new(&snapshot));
+        let tree_sitter_data = TreeSitterData::new(snapshot);
         Self {
             saved_mtime,
+            tree_sitter_data: Arc::new(Mutex::new(tree_sitter_data)),
             saved_version: buffer.version(),
             preview_version: buffer.version(),
             reload_task: None,
@@ -1025,12 +1062,14 @@ impl Buffer {
                 let language_registry = language_registry.clone();
                 syntax.reparse(&text, language_registry, language);
             }
+            let tree_sitter_data = TreeSitterData::new(text.clone());
             BufferSnapshot {
                 text,
                 syntax,
                 file: None,
                 diagnostics: Default::default(),
                 remote_selections: Default::default(),
+                tree_sitter_data: Arc::new(Mutex::new(tree_sitter_data)),
                 language,
                 non_text_state_update_count: 0,
             }
@@ -1048,9 +1087,11 @@ impl Buffer {
         )
         .snapshot();
         let syntax = SyntaxMap::new(&text).snapshot();
+        let tree_sitter_data = TreeSitterData::new(text.clone());
         BufferSnapshot {
             text,
             syntax,
+            tree_sitter_data: Arc::new(Mutex::new(tree_sitter_data)),
             file: None,
             diagnostics: Default::default(),
             remote_selections: Default::default(),
@@ -1075,9 +1116,11 @@ impl Buffer {
         if let Some(language) = language.clone() {
             syntax.reparse(&text, language_registry, language);
         }
+        let tree_sitter_data = TreeSitterData::new(text.clone());
         BufferSnapshot {
             text,
             syntax,
+            tree_sitter_data: Arc::new(Mutex::new(tree_sitter_data)),
             file: None,
             diagnostics: Default::default(),
             remote_selections: Default::default(),
@@ -1097,6 +1140,7 @@ impl Buffer {
         BufferSnapshot {
             text,
             syntax,
+            tree_sitter_data: self.tree_sitter_data.clone(),
             file: self.file.clone(),
             remote_selections: self.remote_selections.clone(),
             diagnostics: self.diagnostics.clone(),
@@ -1611,6 +1655,7 @@ impl Buffer {
         self.syntax_map.lock().did_parse(syntax_snapshot);
         self.request_autoindent(cx);
         self.parse_status.0.send(ParseStatus::Idle).unwrap();
+        self.tree_sitter_data.lock().clear();
         cx.emit(BufferEvent::Reparsed);
         cx.notify();
     }
@@ -4120,61 +4165,166 @@ impl BufferSnapshot {
         self.syntax.matches(range, self, query)
     }
 
+    /// Finds all [`RowChunks`] applicable to the given range, then returns all bracket pairs that intersect with those chunks.
+    /// Hence, may return more bracket pairs than the range contains.
+    ///
+    /// Will omit known chunks.
+    /// The resulting bracket match collections are not ordered.
+    pub fn fetch_bracket_ranges(
+        &self,
+        range: Range<usize>,
+        known_chunks: Option<(&Global, &HashSet<Range<BufferRow>>)>,
+    ) -> HashMap<Range<BufferRow>, Vec<BracketMatch<usize>>> {
+        let mut tree_sitter_data = self.latest_tree_sitter_data().clone();
+
+        let known_chunks = match known_chunks {
+            Some((known_version, known_chunks)) => {
+                if !tree_sitter_data
+                    .chunks
+                    .version()
+                    .changed_since(known_version)
+                {
+                    known_chunks.clone()
+                } else {
+                    HashSet::default()
+                }
+            }
+            None => HashSet::default(),
+        };
+
+        let mut new_bracket_matches = HashMap::default();
+        let mut all_bracket_matches = HashMap::default();
+
+        for chunk in tree_sitter_data
+            .chunks
+            .applicable_chunks(&[self.anchor_before(range.start)..self.anchor_after(range.end)])
+        {
+            if known_chunks.contains(&chunk.row_range()) {
+                continue;
+            }
+            let Some(chunk_range) = tree_sitter_data.chunks.chunk_range(chunk) else {
+                continue;
+            };
+            let chunk_range = chunk_range.to_offset(&tree_sitter_data.chunks.snapshot);
+
+            let bracket_matches = match tree_sitter_data.brackets_by_chunks[chunk.id].take() {
+                Some(cached_brackets) => cached_brackets,
+                None => {
+                    let mut bracket_pairs_ends = Vec::new();
+                    let mut matches =
+                        self.syntax
+                            .matches(chunk_range.clone(), &self.text, |grammar| {
+                                grammar.brackets_config.as_ref().map(|c| &c.query)
+                            });
+                    let configs = matches
+                        .grammars()
+                        .iter()
+                        .map(|grammar| grammar.brackets_config.as_ref().unwrap())
+                        .collect::<Vec<_>>();
+
+                    let chunk_range = chunk_range.clone();
+                    let new_matches = iter::from_fn(move || {
+                        while let Some(mat) = matches.peek() {
+                            let mut open = None;
+                            let mut close = None;
+                            let depth = mat.depth;
+                            let config = configs[mat.grammar_index];
+                            let pattern = &config.patterns[mat.pattern_index];
+                            for capture in mat.captures {
+                                if capture.index == config.open_capture_ix {
+                                    open = Some(capture.node.byte_range());
+                                } else if capture.index == config.close_capture_ix {
+                                    close = Some(capture.node.byte_range());
+                                }
+                            }
+
+                            matches.advance();
+
+                            let Some((open_range, close_range)) = open.zip(close) else {
+                                continue;
+                            };
+
+                            let bracket_range = open_range.start..=close_range.end;
+                            if !bracket_range.overlaps(&chunk_range) {
+                                continue;
+                            }
+
+                            return Some((open_range, close_range, pattern, depth));
+                        }
+                        None
+                    })
+                    .sorted_by_key(|(open_range, _, _, _)| open_range.start)
+                    .map(|(open_range, close_range, pattern, syntax_layer_depth)| {
+                        while let Some(&last_bracket_end) = bracket_pairs_ends.last() {
+                            if last_bracket_end <= open_range.start {
+                                bracket_pairs_ends.pop();
+                            } else {
+                                break;
+                            }
+                        }
+
+                        let bracket_depth = bracket_pairs_ends.len();
+                        bracket_pairs_ends.push(close_range.end);
+
+                        BracketMatch {
+                            open_range,
+                            close_range,
+                            syntax_layer_depth,
+                            newline_only: pattern.newline_only,
+                            color_index: pattern.rainbow_exclude.not().then_some(bracket_depth),
+                        }
+                    })
+                    .collect::<Vec<_>>();
+
+                    new_bracket_matches.insert(chunk.id, new_matches.clone());
+                    new_matches
+                }
+            };
+            all_bracket_matches.insert(chunk.row_range(), bracket_matches);
+        }
+
+        let mut latest_tree_sitter_data = self.latest_tree_sitter_data();
+        if latest_tree_sitter_data.chunks.version() == &self.version {
+            for (chunk_id, new_matches) in new_bracket_matches {
+                let old_chunks = &mut latest_tree_sitter_data.brackets_by_chunks[chunk_id];
+                if old_chunks.is_none() {
+                    *old_chunks = Some(new_matches);
+                }
+            }
+        }
+
+        all_bracket_matches
+    }
+
+    fn latest_tree_sitter_data(&self) -> MutexGuard<'_, RawMutex, TreeSitterData> {
+        let mut tree_sitter_data = self.tree_sitter_data.lock();
+        if self
+            .version
+            .changed_since(tree_sitter_data.chunks.version())
+        {
+            *tree_sitter_data = TreeSitterData::new(self.text.clone());
+        }
+        tree_sitter_data
+    }
+
     pub fn all_bracket_ranges(
         &self,
         range: Range<usize>,
-    ) -> impl Iterator<Item = BracketMatch> + '_ {
-        let mut matches = self.syntax.matches(range.clone(), &self.text, |grammar| {
-            grammar.brackets_config.as_ref().map(|c| &c.query)
-        });
-        let configs = matches
-            .grammars()
-            .iter()
-            .map(|grammar| grammar.brackets_config.as_ref().unwrap())
-            .collect::<Vec<_>>();
-
-        iter::from_fn(move || {
-            while let Some(mat) = matches.peek() {
-                let mut open = None;
-                let mut close = None;
-                let depth = mat.depth;
-                let config = &configs[mat.grammar_index];
-                let pattern = &config.patterns[mat.pattern_index];
-                for capture in mat.captures {
-                    if capture.index == config.open_capture_ix {
-                        open = Some(capture.node.byte_range());
-                    } else if capture.index == config.close_capture_ix {
-                        close = Some(capture.node.byte_range());
-                    }
-                }
-
-                matches.advance();
-
-                let Some((open_range, close_range)) = open.zip(close) else {
-                    continue;
-                };
-
-                let bracket_range = open_range.start..=close_range.end;
-                if !bracket_range.overlaps(&range) {
-                    continue;
-                }
-
-                return Some(BracketMatch {
-                    open_range,
-                    close_range,
-                    newline_only: pattern.newline_only,
-                    depth,
-                });
-            }
-            None
-        })
+    ) -> impl Iterator<Item = BracketMatch<usize>> {
+        self.fetch_bracket_ranges(range.clone(), None)
+            .into_values()
+            .flatten()
+            .filter(move |bracket_match| {
+                let bracket_range = bracket_match.open_range.start..bracket_match.close_range.end;
+                bracket_range.overlaps(&range)
+            })
     }
 
     /// Returns bracket range pairs overlapping or adjacent to `range`
     pub fn bracket_ranges<T: ToOffset>(
         &self,
         range: Range<T>,
-    ) -> impl Iterator<Item = BracketMatch> + '_ {
+    ) -> impl Iterator<Item = BracketMatch<usize>> + '_ {
         // Find bracket pairs that *inclusively* contain the given range.
         let range = range.start.to_previous_offset(self)..range.end.to_next_offset(self);
         self.all_bracket_ranges(range)
@@ -4320,15 +4470,19 @@ impl BufferSnapshot {
     pub fn enclosing_bracket_ranges<T: ToOffset>(
         &self,
         range: Range<T>,
-    ) -> impl Iterator<Item = BracketMatch> + '_ {
+    ) -> impl Iterator<Item = BracketMatch<usize>> + '_ {
         let range = range.start.to_offset(self)..range.end.to_offset(self);
 
         let result: Vec<_> = self.bracket_ranges(range.clone()).collect();
-        let max_depth = result.iter().map(|mat| mat.depth).max().unwrap_or(0);
+        let max_depth = result
+            .iter()
+            .map(|mat| mat.syntax_layer_depth)
+            .max()
+            .unwrap_or(0);
         result.into_iter().filter(move |pair| {
             pair.open_range.start <= range.start
                 && pair.close_range.end >= range.end
-                && pair.depth == max_depth
+                && pair.syntax_layer_depth == max_depth
         })
     }
 
@@ -4815,6 +4969,7 @@ impl Clone for BufferSnapshot {
             remote_selections: self.remote_selections.clone(),
             diagnostics: self.diagnostics.clone(),
             language: self.language.clone(),
+            tree_sitter_data: self.tree_sitter_data.clone(),
             non_text_state_update_count: self.non_text_state_update_count,
         }
     }