editor: Do treesitter bracket colorization work on the background (#50068)

Lukas Wirth created

This is pure computation work that is disconnected from main thread
entity state yet it can still block for a couple milliseconds depending
on the file

Release Notes:

- N/A *or* Added/Fixed/Improved ...

Change summary

crates/editor/src/bracket_colorization.rs | 277 +++++++++++++++---------
crates/editor/src/editor.rs               |   8 
crates/editor/src/editor_tests.rs         |   1 
crates/editor/src/inlays/inlay_hints.rs   |   4 
4 files changed, 175 insertions(+), 115 deletions(-)

Detailed changes

crates/editor/src/bracket_colorization.rs 🔗

@@ -5,10 +5,10 @@
 use std::ops::Range;
 
 use crate::{Editor, HighlightKey};
-use collections::HashMap;
-use gpui::{Context, HighlightStyle};
+use collections::{HashMap, HashSet};
+use gpui::{AppContext as _, Context, HighlightStyle};
 use itertools::Itertools;
-use language::language_settings;
+use language::{BufferRow, BufferSnapshot, language_settings};
 use multi_buffer::{Anchor, ExcerptId};
 use ui::{ActiveTheme, utils::ensure_minimum_contrast};
 
@@ -19,22 +19,16 @@ impl Editor {
         }
 
         if invalidate {
-            self.fetched_tree_sitter_chunks.clear();
+            self.bracket_fetched_tree_sitter_chunks.clear();
         }
 
         let accents_count = cx.theme().accents().0.len();
         let multi_buffer_snapshot = self.buffer().read(cx).snapshot(cx);
-        let anchors_in_multi_buffer = |current_excerpt: ExcerptId,
-                                       text_anchors: [text::Anchor; 4]|
-         -> Option<[Option<_>; 4]> {
-            multi_buffer_snapshot
-                .anchors_in_excerpt(current_excerpt, text_anchors)?
-                .collect_array()
-        };
-
-        let bracket_matches_by_accent = self.visible_excerpts(false, cx).into_iter().fold(
-            HashMap::default(),
-            |mut acc, (excerpt_id, (buffer, _, buffer_range))| {
+
+        let visible_excerpts = self.visible_excerpts(false, cx);
+        let excerpt_data: Vec<(ExcerptId, BufferSnapshot, Range<usize>)> = visible_excerpts
+            .into_iter()
+            .filter_map(|(excerpt_id, (buffer, _, buffer_range))| {
                 let buffer_snapshot = buffer.read(cx).snapshot();
                 if language_settings::language_settings(
                     buffer_snapshot.language().map(|language| language.name()),
@@ -43,112 +37,173 @@ impl Editor {
                 )
                 .colorize_brackets
                 {
-                    let fetched_chunks = self
-                        .fetched_tree_sitter_chunks
-                        .entry(excerpt_id)
-                        .or_default();
-
-                    let brackets_by_accent = buffer_snapshot
-                        .fetch_bracket_ranges(
-                            buffer_range.start..buffer_range.end,
-                            Some(fetched_chunks),
-                        )
-                        .into_iter()
-                        .flat_map(|(chunk_range, pairs)| {
-                            if fetched_chunks.insert(chunk_range) {
-                                pairs
-                            } else {
-                                Vec::new()
-                            }
-                        })
-                        .filter_map(|pair| {
-                            let color_index = pair.color_index?;
-
-                            let buffer_open_range =
-                                buffer_snapshot.anchor_range_around(pair.open_range);
-                            let buffer_close_range =
-                                buffer_snapshot.anchor_range_around(pair.close_range);
-                            let [
-                                buffer_open_range_start,
-                                buffer_open_range_end,
-                                buffer_close_range_start,
-                                buffer_close_range_end,
-                            ] = anchors_in_multi_buffer(
-                                excerpt_id,
-                                [
-                                    buffer_open_range.start,
-                                    buffer_open_range.end,
-                                    buffer_close_range.start,
-                                    buffer_close_range.end,
-                                ],
-                            )?;
-                            let multi_buffer_open_range =
-                                buffer_open_range_start.zip(buffer_open_range_end);
-                            let multi_buffer_close_range =
-                                buffer_close_range_start.zip(buffer_close_range_end);
-
-                            let mut ranges = Vec::with_capacity(2);
-                            if let Some((open_start, open_end)) = multi_buffer_open_range {
-                                ranges.push(open_start..open_end);
-                            }
-                            if let Some((close_start, close_end)) = multi_buffer_close_range {
-                                ranges.push(close_start..close_end);
-                            }
-                            if ranges.is_empty() {
-                                None
-                            } else {
-                                Some((color_index % accents_count, ranges))
-                            }
-                        });
+                    Some((excerpt_id, buffer_snapshot, buffer_range))
+                } else {
+                    None
+                }
+            })
+            .collect();
 
-                    for (accent_number, new_ranges) in brackets_by_accent {
-                        let ranges = acc
-                            .entry(accent_number)
-                            .or_insert_with(Vec::<Range<Anchor>>::new);
+        let mut fetched_tree_sitter_chunks = excerpt_data
+            .iter()
+            .filter_map(|(excerpt_id, ..)| {
+                Some((
+                    *excerpt_id,
+                    self.bracket_fetched_tree_sitter_chunks
+                        .get(excerpt_id)
+                        .cloned()?,
+                ))
+            })
+            .collect::<HashMap<ExcerptId, HashSet<Range<BufferRow>>>>();
+
+        let bracket_matches_by_accent = cx.background_spawn(async move {
+            let anchors_in_multi_buffer = |current_excerpt: ExcerptId,
+                                           text_anchors: [text::Anchor; 4]|
+             -> Option<[Option<_>; 4]> {
+                multi_buffer_snapshot
+                    .anchors_in_excerpt(current_excerpt, text_anchors)?
+                    .collect_array()
+            };
 
-                        for new_range in new_ranges {
-                            let i = ranges
-                                .binary_search_by(|probe| {
-                                    probe.start.cmp(&new_range.start, &multi_buffer_snapshot)
-                                })
-                                .unwrap_or_else(|i| i);
-                            ranges.insert(i, new_range);
+            let bracket_matches_by_accent: HashMap<usize, Vec<Range<Anchor>>> =
+                excerpt_data.into_iter().fold(
+                    HashMap::default(),
+                    |mut acc, (excerpt_id, buffer_snapshot, buffer_range)| {
+                        let fetched_chunks =
+                            fetched_tree_sitter_chunks.entry(excerpt_id).or_default();
+
+                        let brackets_by_accent = compute_bracket_ranges(
+                            &buffer_snapshot,
+                            buffer_range,
+                            fetched_chunks,
+                            excerpt_id,
+                            accents_count,
+                            &anchors_in_multi_buffer,
+                        );
+
+                        for (accent_number, new_ranges) in brackets_by_accent {
+                            let ranges = acc
+                                .entry(accent_number)
+                                .or_insert_with(Vec::<Range<Anchor>>::new);
+
+                            for new_range in new_ranges {
+                                let i = ranges
+                                    .binary_search_by(|probe| {
+                                        probe.start.cmp(&new_range.start, &multi_buffer_snapshot)
+                                    })
+                                    .unwrap_or_else(|i| i);
+                                ranges.insert(i, new_range);
+                            }
                         }
-                    }
-                }
 
-                acc
-            },
-        );
+                        acc
+                    },
+                );
 
-        if invalidate {
-            self.clear_highlights_with(
-                &mut |key| matches!(key, HighlightKey::ColorizeBracket(_)),
-                cx,
-            );
-        }
+            (bracket_matches_by_accent, fetched_tree_sitter_chunks)
+        });
 
         let editor_background = cx.theme().colors().editor_background;
         let accents = cx.theme().accents().clone();
-        for (accent_number, bracket_highlights) in bracket_matches_by_accent {
-            let bracket_color = accents.color_for_index(accent_number as u32);
-            let adjusted_color = ensure_minimum_contrast(bracket_color, editor_background, 55.0);
-            let style = HighlightStyle {
-                color: Some(adjusted_color),
-                ..HighlightStyle::default()
-            };
 
-            self.highlight_text_key(
-                HighlightKey::ColorizeBracket(accent_number),
-                bracket_highlights,
-                style,
-                true,
-                cx,
-            );
-        }
+        self.colorize_brackets_task = cx.spawn(async move |editor, cx| {
+            if invalidate {
+                editor
+                    .update(cx, |editor, cx| {
+                        editor.clear_highlights_with(
+                            &mut |key| matches!(key, HighlightKey::ColorizeBracket(_)),
+                            cx,
+                        );
+                    })
+                    .ok();
+            }
+
+            let (bracket_matches_by_accent, updated_chunks) = bracket_matches_by_accent.await;
+
+            editor
+                .update(cx, |editor, cx| {
+                    editor
+                        .bracket_fetched_tree_sitter_chunks
+                        .extend(updated_chunks);
+                    for (accent_number, bracket_highlights) in bracket_matches_by_accent {
+                        let bracket_color = accents.color_for_index(accent_number as u32);
+                        let adjusted_color =
+                            ensure_minimum_contrast(bracket_color, editor_background, 55.0);
+                        let style = HighlightStyle {
+                            color: Some(adjusted_color),
+                            ..HighlightStyle::default()
+                        };
+
+                        editor.highlight_text_key(
+                            HighlightKey::ColorizeBracket(accent_number),
+                            bracket_highlights,
+                            style,
+                            true,
+                            cx,
+                        );
+                    }
+                })
+                .ok();
+        });
     }
 }
 
+fn compute_bracket_ranges(
+    buffer_snapshot: &BufferSnapshot,
+    buffer_range: Range<usize>,
+    fetched_chunks: &mut HashSet<Range<BufferRow>>,
+    excerpt_id: ExcerptId,
+    accents_count: usize,
+    anchors_in_multi_buffer: &impl Fn(ExcerptId, [text::Anchor; 4]) -> Option<[Option<Anchor>; 4]>,
+) -> Vec<(usize, Vec<Range<Anchor>>)> {
+    buffer_snapshot
+        .fetch_bracket_ranges(buffer_range.start..buffer_range.end, Some(fetched_chunks))
+        .into_iter()
+        .flat_map(|(chunk_range, pairs)| {
+            if fetched_chunks.insert(chunk_range) {
+                pairs
+            } else {
+                Vec::new()
+            }
+        })
+        .filter_map(|pair| {
+            let color_index = pair.color_index?;
+
+            let buffer_open_range = buffer_snapshot.anchor_range_around(pair.open_range);
+            let buffer_close_range = buffer_snapshot.anchor_range_around(pair.close_range);
+            let [
+                buffer_open_range_start,
+                buffer_open_range_end,
+                buffer_close_range_start,
+                buffer_close_range_end,
+            ] = anchors_in_multi_buffer(
+                excerpt_id,
+                [
+                    buffer_open_range.start,
+                    buffer_open_range.end,
+                    buffer_close_range.start,
+                    buffer_close_range.end,
+                ],
+            )?;
+            let multi_buffer_open_range = buffer_open_range_start.zip(buffer_open_range_end);
+            let multi_buffer_close_range = buffer_close_range_start.zip(buffer_close_range_end);
+
+            let mut ranges = Vec::with_capacity(2);
+            if let Some((open_start, open_end)) = multi_buffer_open_range {
+                ranges.push(open_start..open_end);
+            }
+            if let Some((close_start, close_end)) = multi_buffer_close_range {
+                ranges.push(close_start..close_end);
+            }
+            if ranges.is_empty() {
+                None
+            } else {
+                Some((color_index % accents_count, ranges))
+            }
+        })
+        .collect()
+}
+
 #[cfg(test)]
 mod tests {
     use std::{cmp, sync::Arc, time::Duration};
@@ -164,7 +219,7 @@ mod tests {
     };
     use collections::HashSet;
     use fs::FakeFs;
-    use gpui::{AppContext as _, UpdateGlobal as _};
+    use gpui::UpdateGlobal as _;
     use indoc::indoc;
     use itertools::Itertools;
     use language::{Capability, markdown_lang};
@@ -749,6 +804,7 @@ mod foo «1{
                 });
             });
         });
+        cx.executor().run_until_parked();
         assert_eq!(
             &separate_with_comment_lines(
                 indoc! {r#"
@@ -776,6 +832,7 @@ mod foo {
                 });
             });
         });
+        cx.executor().run_until_parked();
         assert_eq!(
             &separate_with_comment_lines(
                 indoc! {r#"

crates/editor/src/editor.rs 🔗

@@ -1347,7 +1347,7 @@ pub struct Editor {
     suppress_selection_callback: bool,
     applicable_language_settings: HashMap<Option<LanguageName>, LanguageSettings>,
     accent_data: Option<AccentData>,
-    fetched_tree_sitter_chunks: HashMap<ExcerptId, HashSet<Range<BufferRow>>>,
+    bracket_fetched_tree_sitter_chunks: HashMap<ExcerptId, HashSet<Range<BufferRow>>>,
     semantic_token_state: SemanticTokenState,
     pub(crate) refresh_matching_bracket_highlights_task: Task<()>,
     refresh_document_symbols_task: Shared<Task<()>>,
@@ -1356,6 +1356,7 @@ pub struct Editor {
     outline_symbols_at_cursor: Option<(BufferId, Vec<OutlineItem<Anchor>>)>,
     sticky_headers_task: Task<()>,
     sticky_headers: Option<Vec<OutlineItem<Anchor>>>,
+    pub(crate) colorize_brackets_task: Task<()>,
 }
 
 #[derive(Debug, PartialEq)]
@@ -2600,7 +2601,7 @@ impl Editor {
             applicable_language_settings: HashMap::default(),
             semantic_token_state: SemanticTokenState::new(cx, full_mode),
             accent_data: None,
-            fetched_tree_sitter_chunks: HashMap::default(),
+            bracket_fetched_tree_sitter_chunks: HashMap::default(),
             number_deleted_lines: false,
             refresh_matching_bracket_highlights_task: Task::ready(()),
             refresh_document_symbols_task: Task::ready(()).shared(),
@@ -2609,6 +2610,7 @@ impl Editor {
             outline_symbols_at_cursor: None,
             sticky_headers_task: Task::ready(()),
             sticky_headers: None,
+            colorize_brackets_task: Task::ready(()),
         };
 
         if is_minimap {
@@ -24165,7 +24167,7 @@ impl Editor {
                 self.refresh_document_highlights(cx);
                 let snapshot = multibuffer.read(cx).snapshot(cx);
                 for id in ids {
-                    self.fetched_tree_sitter_chunks.remove(id);
+                    self.bracket_fetched_tree_sitter_chunks.remove(id);
                     if let Some(buffer) = snapshot.buffer_for_excerpt(*id) {
                         self.semantic_token_state
                             .invalidate_buffer(&buffer.remote_id());

crates/editor/src/editor_tests.rs 🔗

@@ -17284,6 +17284,7 @@ async fn test_no_duplicated_completion_requests(cx: &mut TestAppContext) {
         }
     });
 
+    cx.executor().run_until_parked();
     cx.condition(|editor, _| editor.context_menu_visible())
         .await;
     cx.assert_editor_state("fn main() { let a = 2.ˇ; }");

crates/editor/src/inlays/inlay_hints.rs 🔗

@@ -4501,9 +4501,9 @@ let c = 3;"#
             },
         );
 
-        let buffer = project
+        let (buffer, _buffer_handle) = project
             .update(cx, |project, cx| {
-                project.open_local_buffer(path!("/a/main.rs"), cx)
+                project.open_local_buffer_with_lsp(path!("/a/main.rs"), cx)
             })
             .await
             .unwrap();