bracket_colorization.rs

  1use crate::Editor;
  2use collections::HashMap;
  3use gpui::{Context, HighlightStyle};
  4use language::language_settings;
  5use ui::{ActiveTheme, utils::ensure_minimum_contrast};
  6
  7struct RainbowBracketHighlight;
  8
  9impl Editor {
 10    pub(crate) fn colorize_brackets(&mut self, invalidate: bool, cx: &mut Context<Editor>) {
 11        if !self.mode.is_full() {
 12            return;
 13        }
 14
 15        if invalidate {
 16            self.fetched_tree_sitter_chunks.clear();
 17        }
 18
 19        let accents_count = cx.theme().accents().0.len();
 20        let multi_buffer_snapshot = self.buffer().read(cx).snapshot(cx);
 21        let bracket_matches_by_accent = self.visible_excerpts(cx).into_iter().fold(
 22            HashMap::default(),
 23            |mut acc, (excerpt_id, (buffer, buffer_version, buffer_range))| {
 24                let buffer_snapshot = buffer.read(cx).snapshot();
 25                if language_settings::language_settings(
 26                    buffer_snapshot.language().map(|language| language.name()),
 27                    buffer_snapshot.file(),
 28                    cx,
 29                )
 30                .colorize_brackets
 31                {
 32                    let fetched_chunks = self
 33                        .fetched_tree_sitter_chunks
 34                        .entry(excerpt_id)
 35                        .or_default();
 36
 37                    let brackets_by_accent = buffer_snapshot
 38                        .fetch_bracket_ranges(
 39                            buffer_range.start..buffer_range.end,
 40                            Some((&buffer_version, fetched_chunks)),
 41                        )
 42                        .into_iter()
 43                        .flat_map(|(chunk_range, pairs)| {
 44                            if fetched_chunks.insert(chunk_range) {
 45                                pairs
 46                            } else {
 47                                Vec::new()
 48                            }
 49                        })
 50                        .filter_map(|pair| {
 51                            let color_index = pair.color_index?;
 52                            let buffer_open_range = buffer_snapshot
 53                                .anchor_before(pair.open_range.start)
 54                                ..buffer_snapshot.anchor_after(pair.open_range.end);
 55                            let multi_buffer_open_range = multi_buffer_snapshot
 56                                .anchor_in_excerpt(excerpt_id, buffer_open_range.start)?
 57                                ..multi_buffer_snapshot
 58                                    .anchor_in_excerpt(excerpt_id, buffer_open_range.end)?;
 59                            let buffer_close_range = buffer_snapshot
 60                                .anchor_before(pair.close_range.start)
 61                                ..buffer_snapshot.anchor_after(pair.close_range.end);
 62                            let multi_buffer_close_range = multi_buffer_snapshot
 63                                .anchor_in_excerpt(excerpt_id, buffer_close_range.start)?
 64                                ..multi_buffer_snapshot
 65                                    .anchor_in_excerpt(excerpt_id, buffer_close_range.end)?;
 66                            Some((
 67                                color_index % accents_count,
 68                                multi_buffer_open_range,
 69                                multi_buffer_close_range,
 70                            ))
 71                        });
 72
 73                    for (accent_number, open_range, close_range) in brackets_by_accent {
 74                        let ranges = acc.entry(accent_number).or_insert_with(Vec::new);
 75                        ranges.push(open_range);
 76                        ranges.push(close_range);
 77                    }
 78                }
 79
 80                acc
 81            },
 82        );
 83
 84        if invalidate {
 85            self.clear_highlights::<RainbowBracketHighlight>(cx);
 86        }
 87
 88        let editor_background = cx.theme().colors().editor_background;
 89        for (accent_number, bracket_highlights) in bracket_matches_by_accent {
 90            let bracket_color = cx.theme().accents().color_for_index(accent_number as u32);
 91            let adjusted_color = ensure_minimum_contrast(bracket_color, editor_background, 55.0);
 92            let style = HighlightStyle {
 93                color: Some(adjusted_color),
 94                ..HighlightStyle::default()
 95            };
 96
 97            self.highlight_text_key::<RainbowBracketHighlight>(
 98                accent_number,
 99                bracket_highlights,
100                style,
101                true,
102                cx,
103            );
104        }
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use std::{cmp, sync::Arc, time::Duration};
111
112    use super::*;
113    use crate::{
114        DisplayPoint,
115        display_map::{DisplayRow, ToDisplayPoint},
116        editor_tests::init_test,
117        test::{
118            editor_lsp_test_context::EditorLspTestContext, editor_test_context::EditorTestContext,
119        },
120    };
121    use collections::HashSet;
122    use indoc::indoc;
123    use itertools::Itertools;
124    use languages::rust_lang;
125    use pretty_assertions::assert_eq;
126    use rope::Point;
127    use text::OffsetRangeExt;
128    use util::post_inc;
129
130    #[gpui::test]
131    async fn test_basic_bracket_colorization(cx: &mut gpui::TestAppContext) {
132        init_test(cx, |language_settings| {
133            language_settings.defaults.colorize_brackets = Some(true);
134        });
135        let mut cx = EditorLspTestContext::new(
136            Arc::into_inner(rust_lang()).unwrap(),
137            lsp::ServerCapabilities::default(),
138            cx,
139        )
140        .await;
141
142        cx.set_state(indoc! {r#"ˇuse std::{collections::HashMap, future::Future};
143
144fn main() {
145    let a = one((), { () }, ());
146    println!("{a}");
147    println!("{a}");
148    for i in 0..a {
149        println!("{i}");
150    }
151
152    let b = {
153        {
154            {
155                [([([([([([([([([([((), ())])])])])])])])])])]
156            }
157        }
158    };
159}
160
161#[rustfmt::skip]
162fn one(a: (), (): (), c: ()) -> usize { 1 }
163
164fn two<T>(a: HashMap<String, Vec<Option<T>>>) -> usize
165where
166    T: Future<Output = HashMap<String, Vec<Option<Box<()>>>>>,
167{
168    2
169}
170"#});
171        cx.executor().advance_clock(Duration::from_millis(100));
172        cx.executor().run_until_parked();
173
174        assert_bracket_colors(
175            r#"use std::«1{collections::HashMap, future::Future}1»;
176
177fn main«1()1» «1{
178    let a = one«2(«3()3», «3{ «4()4» }3», «3()3»)2»;
179    println!«2("{a}")2»;
180    println!«2("{a}")2»;
181    for i in 0..a «2{
182        println!«3("{i}")3»;
183    }2»
184
185    let b = «2{
186        «3{
187            «4{
188                «5[«6(«7[«1(«2[«3(«4[«5(«6[«7(«1[«2(«3[«4(«5[«6(«7[«1(«2[«3(«4()4», «4()4»)3»]2»)1»]7»)6»]5»)4»]3»)2»]1»)7»]6»)5»]4»)3»]2»)1»]7»)6»]5»
189            }4»
190        }3»
191    }2»;
192}1»
193
194#«1[rustfmt::skip]1»
195fn one«1(a: «2()2», «2()2»: «2()2», c: «2()2»)1» -> usize «1{ 1 }1»
196
197fn two«1<T>1»«1(a: HashMap«2<String, Vec«3<Option«4<T>4»>3»>2»)1» -> usize
198where
199    T: Future«1<Output = HashMap«2<String, Vec«3<Option«4<Box«5<«6()6»>5»>4»>3»>2»>1»,
200«1{
201    2
202}1»
203
2041 hsla(207.80, 16.20%, 69.19%, 1.00)
2052 hsla(29.00, 54.00%, 65.88%, 1.00)
2063 hsla(286.00, 51.00%, 75.25%, 1.00)
2074 hsla(187.00, 47.00%, 59.22%, 1.00)
2085 hsla(355.00, 65.00%, 75.94%, 1.00)
2096 hsla(95.00, 38.00%, 62.00%, 1.00)
2107 hsla(39.00, 67.00%, 69.00%, 1.00)
211"#,
212            &mut cx,
213        );
214    }
215
216    #[track_caller]
217    fn assert_bracket_colors(expected_markup: &str, cx: &mut EditorTestContext) {
218        let result = cx.update_editor(|editor, window, cx| {
219            let snapshot = editor.snapshot(window, cx);
220            let actual_ranges = snapshot.all_text_highlight_ranges::<RainbowBracketHighlight>();
221            let editor_text = snapshot.text();
222
223            let mut next_index = 1;
224            let mut color_to_index = HashMap::default();
225            let mut annotations = Vec::new();
226            for (color, range) in &actual_ranges {
227                let color_index = *color_to_index
228                    .entry(*color)
229                    .or_insert_with(|| post_inc(&mut next_index));
230                let start_offset = snapshot.buffer_snapshot().point_to_offset(range.start);
231                let end_offset = snapshot.buffer_snapshot().point_to_offset(range.end);
232                let bracket_text = &editor_text[start_offset..end_offset];
233                let bracket_char = bracket_text.chars().next().unwrap();
234
235                if matches!(bracket_char, '{' | '[' | '(' | '<') {
236                    annotations.push((start_offset, format!("«{color_index}")));
237                } else {
238                    annotations.push((end_offset, format!("{color_index}»")));
239                }
240            }
241
242            annotations.sort_by(|(pos_a, text_a), (pos_b, text_b)| {
243                pos_a.cmp(pos_b).reverse().then_with(|| {
244                    let a_is_opening = text_a.starts_with('«');
245                    let b_is_opening = text_b.starts_with('«');
246                    match (a_is_opening, b_is_opening) {
247                        (true, false) => cmp::Ordering::Less,
248                        (false, true) => cmp::Ordering::Greater,
249                        _ => cmp::Ordering::Equal,
250                    }
251                })
252            });
253
254            let mut text_with_annotations = editor_text;
255            for (pos, text) in annotations {
256                text_with_annotations.insert_str(pos, &text);
257            }
258
259            text_with_annotations.push_str("\n");
260            for (index, color) in color_to_index
261                .iter()
262                .map(|(color, index)| (*index, *color))
263                .sorted_by_key(|(index, _)| *index)
264            {
265                text_with_annotations.push_str(&format!("{index} {color}\n"));
266            }
267
268            text_with_annotations
269        });
270        assert_eq!(expected_markup, result);
271    }
272
273    #[gpui::test]
274    async fn test_rainbow_bracket_highlights(cx: &mut gpui::TestAppContext) {
275        init_test(cx, |language_settings| {
276            language_settings.defaults.colorize_brackets = Some(true);
277        });
278        let mut cx = EditorLspTestContext::new(
279            Arc::into_inner(rust_lang()).unwrap(),
280            lsp::ServerCapabilities::default(),
281            cx,
282        )
283        .await;
284
285        // taken from r-a https://github.com/rust-lang/rust-analyzer/blob/d733c07552a2dc0ec0cc8f4df3f0ca969a93fd90/crates/ide/src/inlay_hints.rs#L81-L297
286        cx.set_state(indoc! {r#287            pub(crate) fn inlay_hints(
288                db: &RootDatabase,
289                file_id: FileId,
290                range_limit: Option<TextRange>,
291                config: &InlayHintsConfig,
292            ) -> Vec<InlayHint> {
293                let _p = tracing::info_span!("inlay_hints").entered();
294                let sema = Semantics::new(db);
295                let file_id = sema
296                    .attach_first_edition(file_id)
297                    .unwrap_or_else(|| EditionedFileId::current_edition(db, file_id));
298                let file = sema.parse(file_id);
299                let file = file.syntax();
300
301                let mut acc = Vec::new();
302
303                let Some(scope) = sema.scope(file) else {
304                    return acc;
305                };
306                let famous_defs = FamousDefs(&sema, scope.krate());
307                let display_target = famous_defs.1.to_display_target(sema.db);
308
309                let ctx = &mut InlayHintCtx::default();
310                let mut hints = |event| {
311                    if let Some(node) = handle_event(ctx, event) {
312                        hints(&mut acc, ctx, &famous_defs, config, file_id, display_target, node);
313                    }
314                };
315                let mut preorder = file.preorder();
316                salsa::attach(sema.db, || {
317                    while let Some(event) = preorder.next() {
318                        if matches!((&event, range_limit), (WalkEvent::Enter(node), Some(range)) if range.intersect(node.text_range()).is_none())
319                        {
320                            preorder.skip_subtree();
321                            continue;
322                        }
323                        hints(event);
324                    }
325                });
326                if let Some(range_limit) = range_limit {
327                    acc.retain(|hint| range_limit.contains_range(hint.range));
328                }
329                acc
330            }
331
332            #[derive(Default)]
333            struct InlayHintCtx {
334                lifetime_stacks: Vec<Vec<SmolStr>>,
335                extern_block_parent: Option<ast::ExternBlock>,
336            }
337
338            pub(crate) fn inlay_hints_resolve(
339                db: &RootDatabase,
340                file_id: FileId,
341                resolve_range: TextRange,
342                hash: u64,
343                config: &InlayHintsConfig,
344                hasher: impl Fn(&InlayHint) -> u64,
345            ) -> Option<InlayHint> {
346                let _p = tracing::info_span!("inlay_hints_resolve").entered();
347                let sema = Semantics::new(db);
348                let file_id = sema
349                    .attach_first_edition(file_id)
350                    .unwrap_or_else(|| EditionedFileId::current_edition(db, file_id));
351                let file = sema.parse(file_id);
352                let file = file.syntax();
353
354                let scope = sema.scope(file)?;
355                let famous_defs = FamousDefs(&sema, scope.krate());
356                let mut acc = Vec::new();
357
358                let display_target = famous_defs.1.to_display_target(sema.db);
359
360                let ctx = &mut InlayHintCtx::default();
361                let mut hints = |event| {
362                    if let Some(node) = handle_event(ctx, event) {
363                        hints(&mut acc, ctx, &famous_defs, config, file_id, display_target, node);
364                    }
365                };
366
367                let mut preorder = file.preorder();
368                while let Some(event) = preorder.next() {
369                    // FIXME: This can miss some hints that require the parent of the range to calculate
370                    if matches!(&event, WalkEvent::Enter(node) if resolve_range.intersect(node.text_range()).is_none())
371                    {
372                        preorder.skip_subtree();
373                        continue;
374                    }
375                    hints(event);
376                }
377                acc.into_iter().find(|hint| hasher(hint) == hash)
378            }
379
380            fn handle_event(ctx: &mut InlayHintCtx, node: WalkEvent<SyntaxNode>) -> Option<SyntaxNode> {
381                match node {
382                    WalkEvent::Enter(node) => {
383                        if let Some(node) = ast::AnyHasGenericParams::cast(node.clone()) {
384                            let params = node
385                                .generic_param_list()
386                                .map(|it| {
387                                    it.lifetime_params()
388                                        .filter_map(|it| {
389                                            it.lifetime().map(|it| format_smolstr!("{}", &it.text()[1..]))
390                                        })
391                                        .collect()
392                                })
393                                .unwrap_or_default();
394                            ctx.lifetime_stacks.push(params);
395                        }
396                        if let Some(node) = ast::ExternBlock::cast(node.clone()) {
397                            ctx.extern_block_parent = Some(node);
398                        }
399                        Some(node)
400                    }
401                    WalkEvent::Leave(n) => {
402                        if ast::AnyHasGenericParams::can_cast(n.kind()) {
403                            ctx.lifetime_stacks.pop();
404                        }
405                        if ast::ExternBlock::can_cast(n.kind()) {
406                            ctx.extern_block_parent = None;
407                        }
408                        None
409                    }
410                }
411            }
412
413            // FIXME: At some point when our hir infra is fleshed out enough we should flip this and traverse the
414            // HIR instead of the syntax tree.
415            fn hints(
416                hints: &mut Vec<InlayHint>,
417                ctx: &mut InlayHintCtx,
418                famous_defs @ FamousDefs(sema, _krate): &FamousDefs<'_, '_>,
419                config: &InlayHintsConfig,
420                file_id: EditionedFileId,
421                display_target: DisplayTarget,
422                node: SyntaxNode,
423            ) {
424                closing_brace::hints(
425                    hints,
426                    sema,
427                    config,
428                    display_target,
429                    InRealFile { file_id, value: node.clone() },
430                );
431                if let Some(any_has_generic_args) = ast::AnyHasGenericArgs::cast(node.clone()) {
432                    generic_param::hints(hints, famous_defs, config, any_has_generic_args);
433                }
434
435                match_ast! {
436                    match node {
437                        ast::Expr(expr) => {
438                            chaining::hints(hints, famous_defs, config, display_target, &expr);
439                            adjustment::hints(hints, famous_defs, config, display_target, &expr);
440                            match expr {
441                                ast::Expr::CallExpr(it) => param_name::hints(hints, famous_defs, config, file_id, ast::Expr::from(it)),
442                                ast::Expr::MethodCallExpr(it) => {
443                                    param_name::hints(hints, famous_defs, config, file_id, ast::Expr::from(it))
444                                }
445                                ast::Expr::ClosureExpr(it) => {
446                                    closure_captures::hints(hints, famous_defs, config, it.clone());
447                                    closure_ret::hints(hints, famous_defs, config, display_target, it)
448                                },
449                                ast::Expr::RangeExpr(it) => range_exclusive::hints(hints, famous_defs, config, it),
450                                _ => Some(()),
451                            }
452                        },
453                        ast::Pat(it) => {
454                            binding_mode::hints(hints, famous_defs, config, &it);
455                            match it {
456                                ast::Pat::IdentPat(it) => {
457                                    bind_pat::hints(hints, famous_defs, config, display_target, &it);
458                                }
459                                ast::Pat::RangePat(it) => {
460                                    range_exclusive::hints(hints, famous_defs, config, it);
461                                }
462                                _ => {}
463                            }
464                            Some(())
465                        },
466                        ast::Item(it) => match it {
467                            ast::Item::Fn(it) => {
468                                implicit_drop::hints(hints, famous_defs, config, display_target, &it);
469                                if let Some(extern_block) = &ctx.extern_block_parent {
470                                    extern_block::fn_hints(hints, famous_defs, config, &it, extern_block);
471                                }
472                                lifetime::fn_hints(hints, ctx, famous_defs, config,  it)
473                            },
474                            ast::Item::Static(it) => {
475                                if let Some(extern_block) = &ctx.extern_block_parent {
476                                    extern_block::static_hints(hints, famous_defs, config, &it, extern_block);
477                                }
478                                implicit_static::hints(hints, famous_defs, config,  Either::Left(it))
479                            },
480                            ast::Item::Const(it) => implicit_static::hints(hints, famous_defs, config, Either::Right(it)),
481                            ast::Item::Enum(it) => discriminant::enum_hints(hints, famous_defs, config, it),
482                            ast::Item::ExternBlock(it) => extern_block::extern_block_hints(hints, famous_defs, config, it),
483                            _ => None,
484                        },
485                        // FIXME: trait object type elisions
486                        ast::Type(ty) => match ty {
487                            ast::Type::FnPtrType(ptr) => lifetime::fn_ptr_hints(hints, ctx, famous_defs, config,  ptr),
488                            ast::Type::PathType(path) => {
489                                lifetime::fn_path_hints(hints, ctx, famous_defs, config, &path);
490                                implied_dyn_trait::hints(hints, famous_defs, config, Either::Left(path));
491                                Some(())
492                            },
493                            ast::Type::DynTraitType(dyn_) => {
494                                implied_dyn_trait::hints(hints, famous_defs, config, Either::Right(dyn_));
495                                Some(())
496                            },
497                            _ => Some(()),
498                        },
499                        ast::GenericParamList(it) => bounds::hints(hints, famous_defs, config,  it),
500                        _ => Some(()),
501                    }
502                };
503            }
504        "#});
505        cx.executor().advance_clock(Duration::from_millis(100));
506        cx.executor().run_until_parked();
507
508        let actual_ranges = cx.update_editor(|editor, window, cx| {
509            editor
510                .snapshot(window, cx)
511                .all_text_highlight_ranges::<RainbowBracketHighlight>()
512        });
513
514        let mut highlighted_brackets = HashMap::default();
515        for (color, range) in actual_ranges.iter().cloned() {
516            highlighted_brackets.insert(range, color);
517        }
518
519        let last_bracket = actual_ranges
520            .iter()
521            .max_by_key(|(_, p)| p.end.row)
522            .unwrap()
523            .clone();
524
525        cx.update_editor(|editor, window, cx| {
526            let was_scrolled = editor.set_scroll_position(
527                gpui::Point::new(0.0, last_bracket.1.end.row as f64 * 2.0),
528                window,
529                cx,
530            );
531            assert!(was_scrolled.0);
532        });
533        cx.executor().advance_clock(Duration::from_millis(100));
534        cx.executor().run_until_parked();
535
536        let ranges_after_scrolling = cx.update_editor(|editor, window, cx| {
537            editor
538                .snapshot(window, cx)
539                .all_text_highlight_ranges::<RainbowBracketHighlight>()
540        });
541        let new_last_bracket = ranges_after_scrolling
542            .iter()
543            .max_by_key(|(_, p)| p.end.row)
544            .unwrap()
545            .clone();
546
547        assert_ne!(
548            last_bracket, new_last_bracket,
549            "After scrolling down, we should have highlighted more brackets"
550        );
551
552        cx.update_editor(|editor, window, cx| {
553            let was_scrolled = editor.set_scroll_position(gpui::Point::default(), window, cx);
554            assert!(was_scrolled.0);
555        });
556
557        for _ in 0..200 {
558            cx.update_editor(|editor, window, cx| {
559                editor.apply_scroll_delta(gpui::Point::new(0.0, 0.25), window, cx);
560            });
561            cx.executor().run_until_parked();
562
563            let colored_brackets = cx.update_editor(|editor, window, cx| {
564                editor
565                    .snapshot(window, cx)
566                    .all_text_highlight_ranges::<RainbowBracketHighlight>()
567            });
568            for (color, range) in colored_brackets.clone() {
569                assert!(
570                    highlighted_brackets.entry(range).or_insert(color) == &color,
571                    "Colors should stay consistent while scrolling!"
572                );
573            }
574
575            let snapshot = cx.update_editor(|editor, window, cx| editor.snapshot(window, cx));
576            let scroll_position = snapshot.scroll_position();
577            let visible_lines =
578                cx.update_editor(|editor, _, _| editor.visible_line_count().unwrap());
579            let visible_range = DisplayRow(scroll_position.y as u32)
580                ..DisplayRow((scroll_position.y + visible_lines) as u32);
581
582            let current_highlighted_bracket_set: HashSet<Point> = HashSet::from_iter(
583                colored_brackets
584                    .iter()
585                    .flat_map(|(_, range)| [range.start, range.end]),
586            );
587
588            for highlight_range in highlighted_brackets.keys().filter(|bracket_range| {
589                visible_range.contains(&bracket_range.start.to_display_point(&snapshot).row())
590                    || visible_range.contains(&bracket_range.end.to_display_point(&snapshot).row())
591            }) {
592                assert!(
593                    current_highlighted_bracket_set.contains(&highlight_range.start)
594                        || current_highlighted_bracket_set.contains(&highlight_range.end),
595                    "Should not lose highlights while scrolling in the visible range!"
596                );
597            }
598
599            let buffer_snapshot = snapshot.buffer().as_singleton().unwrap().2;
600            for (start, end) in snapshot
601                .bracket_ranges(
602                    DisplayPoint::new(visible_range.start, Default::default()).to_point(&snapshot)
603                        ..DisplayPoint::new(
604                            visible_range.end,
605                            snapshot.line_len(visible_range.end),
606                        )
607                        .to_point(&snapshot),
608                )
609                .into_iter()
610                .flatten()
611            {
612                let start_bracket = colored_brackets
613                    .iter()
614                    .find(|(_, range)| range.to_offset(buffer_snapshot) == start);
615                assert!(
616                    start_bracket.is_some(),
617                    "Existing bracket start in the visible range should be highlighted"
618                );
619
620                let end_bracket = colored_brackets
621                    .iter()
622                    .find(|(_, range)| range.to_offset(buffer_snapshot) == end);
623                assert!(
624                    end_bracket.is_some(),
625                    "Existing bracket end in the visible range should be highlighted"
626                );
627
628                assert_eq!(
629                    start_bracket.unwrap().0,
630                    end_bracket.unwrap().0,
631                    "Bracket pair should be highlighted the same color!"
632                )
633            }
634        }
635
636        // todo! more tests, check no brackets missing in range, settings toggle
637    }
638}