Account for greedy tree-sitter bracket matches (#43607)

Kirill Bulatov created

Current approach is to colorize brackets based on their depth, which was
broken for markdown:

<img width="388" height="50" alt="image"
src="https://github.com/user-attachments/assets/bd8b6c2f-5a26-4d6b-a301-88675bf05920"
/>

Markdown grammar, for bracket queries
https://github.com/zed-industries/zed/blob/00e93bfa113a3daed6e4a97a7244ad04d58453ee/crates/languages/src/markdown/brackets.scm#L1-L8

and markdown document `[LLM-powered features](./ai/overview.md), [bring
and configure your own API
keys](./ai/llm-providers.md#use-your-own-keys)`, matches first bracket
(offset 0) with two different ones:

* `[LLM-powered features]`

* `[LLM-powered features](./ai/overview.md), [bring and configure your
own API keys]`

which mix and add different color markers.

Now, in case multiple pairs exist for the same first bracket, Zed will
only colorize the shortest one:
<img width="373" height="33" alt="image"
src="https://github.com/user-attachments/assets/04b3f7af-8927-4a8b-8f52-de8b5bb063ac"
/>


Release Notes:

- Fixed bracket colorization mixing colors in markdown files

Change summary

crates/editor/src/bracket_colorization.rs | 27 ++++++++
crates/language/src/buffer.rs             | 77 +++++++++++++++++-------
2 files changed, 81 insertions(+), 23 deletions(-)

Detailed changes

crates/editor/src/bracket_colorization.rs 🔗

@@ -161,7 +161,7 @@ mod tests {
     use gpui::{AppContext as _, UpdateGlobal as _};
     use indoc::indoc;
     use itertools::Itertools;
-    use language::Capability;
+    use language::{Capability, markdown_lang};
     use languages::rust_lang;
     use multi_buffer::{ExcerptRange, MultiBuffer};
     use pretty_assertions::assert_eq;
@@ -261,6 +261,31 @@ where
         );
     }
 
+    #[gpui::test]
+    async fn test_markdown_bracket_colorization(cx: &mut gpui::TestAppContext) {
+        init_test(cx, |language_settings| {
+            language_settings.defaults.colorize_brackets = Some(true);
+        });
+        let mut cx = EditorLspTestContext::new(
+            Arc::into_inner(markdown_lang()).unwrap(),
+            lsp::ServerCapabilities::default(),
+            cx,
+        )
+        .await;
+
+        cx.set_state(indoc! {r#"ˇ[LLM-powered features](./ai/overview.md), [bring and configure your own API keys](./ai/llm-providers.md#use-your-own-keys)"#});
+        cx.executor().advance_clock(Duration::from_millis(100));
+        cx.executor().run_until_parked();
+
+        assert_eq!(
+            r#"«1[LLM-powered features]1»«1(./ai/overview.md)1», «1[bring and configure your own API keys]1»«1(./ai/llm-providers.md#use-your-own-keys)1»
+1 hsla(207.80, 16.20%, 69.19%, 1.00)
+"#,
+            &bracket_colors_markup(&mut cx),
+            "All markdown brackets should be colored based on their depth"
+        );
+    }
+
     #[gpui::test]
     async fn test_bracket_colorization_when_editing(cx: &mut gpui::TestAppContext) {
         init_test(cx, |language_settings| {

crates/language/src/buffer.rs 🔗

@@ -45,12 +45,12 @@ use std::{
     borrow::Cow,
     cell::Cell,
     cmp::{self, Ordering, Reverse},
-    collections::{BTreeMap, BTreeSet},
+    collections::{BTreeMap, BTreeSet, hash_map},
     future::Future,
     iter::{self, Iterator, Peekable},
     mem,
     num::NonZeroU32,
-    ops::{Deref, Not, Range},
+    ops::{Deref, Range},
     path::PathBuf,
     rc,
     sync::{Arc, LazyLock},
@@ -4236,6 +4236,7 @@ impl BufferSnapshot {
 
         let mut new_bracket_matches = HashMap::default();
         let mut all_bracket_matches = HashMap::default();
+        let mut bracket_matches_to_color = HashMap::default();
 
         for chunk in tree_sitter_data
             .chunks
@@ -4265,7 +4266,7 @@ impl BufferSnapshot {
                         .collect::<Vec<_>>();
 
                     let chunk_range = chunk_range.clone();
-                    let new_matches = iter::from_fn(move || {
+                    let tree_sitter_matches = iter::from_fn(|| {
                         while let Some(mat) = matches.peek() {
                             let mut open = None;
                             let mut close = None;
@@ -4291,32 +4292,64 @@ impl BufferSnapshot {
                                 continue;
                             }
 
+                            if !pattern.rainbow_exclude {
+                                // Certain tree-sitter grammars may return more bracket pairs than needed:
+                                // see `test_markdown_bracket_colorization` for a set-up that returns pairs with the same start bracket and different end one.
+                                // Pick the pair with the shortest range in case of ambiguity.
+                                match bracket_matches_to_color.entry(open_range.clone()) {
+                                    hash_map::Entry::Vacant(v) => {
+                                        v.insert(close_range.clone());
+                                    }
+                                    hash_map::Entry::Occupied(mut o) => {
+                                        let previous_close_range = o.get();
+                                        let previous_length =
+                                            previous_close_range.end - open_range.start;
+                                        let new_length = close_range.end - open_range.start;
+                                        if new_length < previous_length {
+                                            o.insert(close_range.clone());
+                                        }
+                                    }
+                                }
+                            }
                             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;
-                            }
-                        }
+                    .collect::<Vec<_>>();
 
-                        let bracket_depth = bracket_pairs_ends.len();
-                        bracket_pairs_ends.push(close_range.end);
+                    let new_matches = tree_sitter_matches
+                        .into_iter()
+                        .map(|(open_range, close_range, pattern, syntax_layer_depth)| {
+                            let participates_in_coloring =
+                                bracket_matches_to_color.get(&open_range).is_some_and(
+                                    |close_range_to_color| close_range_to_color == &close_range,
+                                );
+                            let color_index = if participates_in_coloring {
+                                while let Some(&last_bracket_end) = bracket_pairs_ends.last() {
+                                    if last_bracket_end <= open_range.start {
+                                        bracket_pairs_ends.pop();
+                                    } else {
+                                        break;
+                                    }
+                                }
 
-                        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<_>>();
+                                let bracket_depth = bracket_pairs_ends.len();
+                                bracket_pairs_ends.push(close_range.end);
+                                Some(bracket_depth)
+                            } else {
+                                None
+                            };
+
+                            BracketMatch {
+                                open_range,
+                                close_range,
+                                syntax_layer_depth,
+                                newline_only: pattern.newline_only,
+                                color_index,
+                            }
+                        })
+                        .collect::<Vec<_>>();
 
                     new_bracket_matches.insert(chunk.id, new_matches.clone());
                     new_matches