Stop building semantic highlights on the main thread (#57264)

Trace on `main` that showed the hogged main thread: 
[Sample of Zed
Nightly.txt](https://github.com/user-attachments/files/28060310/Sample.of.Zed.Nightly.txt)


Release Notes:

- N/A

crates/language_tools/src/highlights_tree_view.rs

@@ -10,8 +10,9 @@ use gpui::{
 };
 use language::{BufferId, Point, ToOffset};
 use menu::{SelectNext, SelectPrevious};
-use std::{mem, ops::Range};
+use std::{mem, ops::Range, sync::Arc, time::Duration};
 use theme::ActiveTheme;
+use theme::SyntaxTheme;
 use ui::{
     ButtonCommon, ButtonLike, ButtonStyle, Color, ContextMenu, FluentBuilder as _, IconButton,
     IconName, IconPosition, IconSize, Label, LabelCommon, LabelSize, PopoverMenu,
@@ -147,6 +148,7 @@ pub struct HighlightsTreeView {
     show_syntax_tokens: bool,
     show_semantic_tokens: bool,
     skip_next_scroll: bool,
+    refresh_task: Task<()>,
 }
 
 pub struct HighlightsTreeToolbarItemView {
@@ -160,6 +162,19 @@ struct EditorState {
     _subscription: gpui::Subscription,
 }
 
+struct SemanticHighlightEntry {
+    range: Range<Anchor>,
+    style: HighlightStyle,
+    category: HighlightCategory,
+}
+
+struct HighlightRefreshInput {
+    multi_buffer_snapshot: MultiBufferSnapshot,
+    text_highlights: Vec<(HighlightKey, Arc<(HighlightStyle, Vec<Range<Anchor>>)>)>,
+    semantic_highlights: Vec<SemanticHighlightEntry>,
+    syntax_theme: Arc<SyntaxTheme>,
+}
+
 impl HighlightsTreeView {
     pub fn new(
         workspace_handle: WeakEntity<Workspace>,
@@ -181,6 +196,7 @@ impl HighlightsTreeView {
             show_syntax_tokens: true,
             show_semantic_tokens: true,
             skip_next_scroll: false,
+            refresh_task: Task::ready(()),
         };
 
         this.handle_item_updated(active_item, window, cx);
@@ -254,6 +270,7 @@ impl HighlightsTreeView {
     }
 
     fn clear(&mut self, cx: &mut Context<Self>) {
+        self.refresh_task = Task::ready(());
         self.cached_entries.clear();
         self.display_items.clear();
         self.selected_item_ix = None;
@@ -274,9 +291,9 @@ impl HighlightsTreeView {
 
         let subscription =
             cx.subscribe_in(&editor, window, |this, _, event, window, cx| match event {
-                editor::EditorEvent::Reparsed(_)
-                | editor::EditorEvent::SelectionsChanged { .. } => {
-                    this.refresh_highlights(window, cx);
+                editor::EditorEvent::Reparsed(_) => this.schedule_refresh_highlights(window, cx),
+                editor::EditorEvent::SelectionsChanged { .. } => {
+                    this.update_selection_from_editor(cx);
                 }
                 _ => return,
             });
@@ -285,68 +302,73 @@ impl HighlightsTreeView {
             editor,
             _subscription: subscription,
         });
-        self.refresh_highlights(window, cx);
+        self.refresh_task = Task::ready(());
+        self.cached_entries.clear();
+        self.display_items.clear();
+        self.selected_item_ix = None;
+        self.hovered_item_ix = None;
+        self.schedule_refresh_highlights(window, cx);
+        cx.notify();
     }
 
-    fn refresh_highlights(&mut self, _window: &mut Window, cx: &mut Context<Self>) {
-        let Some(editor_state) = self.editor.as_ref() else {
-            self.clear(cx);
-            return;
-        };
+    fn schedule_refresh_highlights(&mut self, window: &mut Window, cx: &mut Context<Self>) {
+        self.refresh_task = cx.spawn_in(window, async move |this, cx| {
+            cx.background_executor()
+                .timer(Duration::from_millis(30))
+                .await;
 
-        let (display_map, project, multi_buffer, cursor_position) = {
-            let editor = editor_state.editor.read(cx);
-            let cursor = editor.selections.newest_anchor().head();
-            (
-                editor.display_map.clone(),
-                editor.project().cloned(),
-                editor.buffer().clone(),
-                cursor,
-            )
-        };
-        let Some(project) = project else {
-            return;
-        };
-
-        let multi_buffer_snapshot = multi_buffer.read(cx).snapshot(cx);
-        let is_singleton = multi_buffer_snapshot.is_singleton();
-        self.is_singleton = is_singleton;
-
-        let mut entries = Vec::new();
-
-        let semantic_theme = cx.theme().syntax().clone();
-        display_map.update(cx, |display_map, cx| {
-            for (key, text_highlights) in display_map.all_text_highlights() {
-                for range in &text_highlights.1 {
-                    let Some((range_display, buffer_id, buffer_point_range)) =
-                        format_anchor_range(range, &multi_buffer_snapshot)
+            let Some(input) = this
+                .update(cx, |this, cx| this.highlight_refresh_input(cx))
+                .ok()
+                .flatten()
             else {
-                        continue;
+                return;
             };
-                    entries.push(HighlightEntry {
-                        range: range.clone(),
-                        buffer_id,
-                        range_display,
-                        style: text_highlights.0,
-                        category: HighlightCategory::Text(*key),
-                        buffer_point_range,
+
+            let new_highlights = cx
+                .background_spawn(async move { build_highlight_entries(input) })
+                .await;
+
+            this.update_in(cx, |this, _window, cx| {
+                this.apply_highlight_refresh(new_highlights, cx);
+            })
+            .ok();
         });
     }
-            }
 
+    fn highlight_refresh_input(&self, cx: &mut Context<Self>) -> Option<HighlightRefreshInput> {
+        let editor_state = self.editor.as_ref()?;
+        let editor = editor_state.editor.read(cx);
+        let display_map = editor.display_map.clone();
+        let project = editor.project().cloned()?;
+        let multi_buffer = editor.buffer().clone();
+        let multi_buffer_snapshot = multi_buffer.read(cx).snapshot(cx);
+        let syntax_theme = cx.theme().syntax().clone();
+
+        let (text_highlights, semantic_token_highlights) =
+            display_map.update(cx, |display_map, _| {
+                let text_highlights = display_map
+                    .all_text_highlights()
+                    .map(|(key, highlights)| (*key, highlights.clone()))
+                    .collect::<Vec<_>>();
+                let semantic_token_highlights = display_map
+                    .all_semantic_token_highlights()
+                    .map(|(buffer_id, (tokens, interner))| {
+                        (*buffer_id, tokens.clone(), interner.clone())
+                    })
+                    .collect::<Vec<_>>();
+                (text_highlights, semantic_token_highlights)
+            });
+
+        let mut semantic_highlights = Vec::new();
         project.read(cx).lsp_store().update(cx, |lsp_store, cx| {
-                for (buffer_id, (tokens, interner)) in display_map.all_semantic_token_highlights() {
+            for (buffer_id, tokens, interner) in semantic_token_highlights {
                 let language_name = multi_buffer
                     .read(cx)
-                        .buffer(*buffer_id)
-                        .and_then(|buf| buf.read(cx).language().map(|l| l.name()));
+                    .buffer(buffer_id)
+                    .and_then(|buffer| buffer.read(cx).language().map(|language| language.name()));
+
                 for token in tokens.iter() {
-                        let range = token.range.start..token.range.end;
-                        let Some((range_display, entry_buffer_id, buffer_point_range)) =
-                            format_anchor_range(&range, &multi_buffer_snapshot)
-                        else {
-                            continue;
-                        };
                     let Some(stylizer) = lsp_store.get_or_create_token_stylizer(
                         token.server_id,
                         language_name.as_ref(),
@@ -355,37 +377,29 @@ impl HighlightsTreeView {
                         continue;
                     };
 
-                        let theme_key =
-                            stylizer
+                    let theme_key = stylizer
                         .rules_for_token(token.token_type)
                         .and_then(|rules| {
                             rules
                                 .iter()
                                 .filter(|rule| {
                                     rule.token_modifiers.iter().all(|modifier| {
-                                                stylizer
-                                                    .has_modifier(token.token_modifiers, modifier)
+                                        stylizer.has_modifier(token.token_modifiers, modifier)
                                     })
                                 })
                                 .fold(None, |theme_key, rule| {
                                     rule.style
                                         .iter()
                                         .find(|style_name| {
-                                                    semantic_theme
-                                                        .style_for_name(style_name)
-                                                        .is_some()
-                                                })
-                                                .map(|style_name| {
-                                                    SharedString::from(style_name.clone())
+                                            syntax_theme.style_for_name(style_name).is_some()
                                         })
+                                        .map(|style_name| SharedString::from(style_name.clone()))
                                         .or(theme_key)
                                 })
                         });
 
-                        entries.push(HighlightEntry {
-                            range,
-                            buffer_id: entry_buffer_id,
-                            range_display,
+                    semantic_highlights.push(SemanticHighlightEntry {
+                        range: token.range.start..token.range.end,
                         style: interner[token.style],
                         category: HighlightCategory::SemanticToken {
                             token_type: stylizer.token_type_name(token.token_type).cloned(),
@@ -394,112 +408,46 @@ impl HighlightsTreeView {
                                 .map(SharedString::from),
                             theme_key,
                         },
-                            buffer_point_range,
                     });
                 }
             }
         });
-        });
 
-        let syntax_theme = cx.theme().syntax().clone();
-        for excerpt_range in multi_buffer_snapshot.excerpts() {
-            let Some(buffer_snapshot) =
-                multi_buffer_snapshot.buffer_for_id(excerpt_range.context.start.buffer_id)
-            else {
-                continue;
-            };
-
-            let start_offset = excerpt_range.context.start.to_offset(buffer_snapshot);
-            let end_offset = excerpt_range.context.end.to_offset(buffer_snapshot);
-            let range = start_offset..end_offset;
-
-            let captures = buffer_snapshot.captures(range, |grammar| {
-                grammar.highlights_config.as_ref().map(|c| &c.query)
-            });
-            let grammars: Vec<_> = captures.grammars().to_vec();
-            let highlight_maps: Vec<_> = grammars.iter().map(|g| g.highlight_map()).collect();
-
-            for capture in captures {
-                let Some(highlight_id) = highlight_maps[capture.grammar_index].get(capture.index)
-                else {
-                    continue;
-                };
-                let Some(style) = syntax_theme.get(highlight_id).cloned() else {
-                    continue;
-                };
-
-                let theme_key = syntax_theme
-                    .get_capture_name(highlight_id)
-                    .map(|theme_key| SharedString::from(theme_key.to_string()));
-
-                let capture_name = grammars[capture.grammar_index]
-                    .highlights_config
-                    .as_ref()
-                    .and_then(|config| config.query.capture_names().get(capture.index as usize))
-                    .map(|capture_name| SharedString::from((*capture_name).to_string()))
-                    .unwrap_or_else(|| SharedString::from("unknown"));
-
-                let start_anchor = buffer_snapshot.anchor_before(capture.node.start_byte());
-                let end_anchor = buffer_snapshot.anchor_after(capture.node.end_byte());
-
-                let start = multi_buffer_snapshot.anchor_in_excerpt(start_anchor);
-                let end = multi_buffer_snapshot.anchor_in_excerpt(end_anchor);
-
-                let (start, end) = match (start, end) {
-                    (Some(s), Some(e)) => (s, e),
-                    _ => continue,
-                };
-
-                let range = start..end;
-                let Some((range_display, buffer_id, buffer_point_range)) =
-                    format_anchor_range(&range, &multi_buffer_snapshot)
-                else {
-                    continue;
-                };
-
-                entries.push(HighlightEntry {
-                    range,
-                    buffer_id,
-                    range_display,
-                    style,
-                    category: HighlightCategory::SyntaxToken {
-                        capture_name,
-                        theme_key,
-                    },
-                    buffer_point_range,
-                });
-            }
+        Some(HighlightRefreshInput {
+            multi_buffer_snapshot,
+            text_highlights,
+            semantic_highlights,
+            syntax_theme,
+        })
     }
 
-        entries.sort_by(|a, b| {
-            a.buffer_id
-                .cmp(&b.buffer_id)
-                .then_with(|| a.buffer_point_range.start.cmp(&b.buffer_point_range.start))
-                .then_with(|| a.buffer_point_range.end.cmp(&b.buffer_point_range.end))
-                .then_with(|| a.category.cmp(&b.category))
-        });
-        entries.dedup_by(|a, b| {
-            a.buffer_id == b.buffer_id
-                && a.buffer_point_range == b.buffer_point_range
-                && a.category == b.category
-        });
+    fn apply_highlight_refresh(
+        &mut self,
+        new_highlights: Vec<HighlightEntry>,
+        cx: &mut Context<Self>,
+    ) {
+        let Some(editor) = self.editor.as_ref().map(|state| state.editor.clone()) else {
+            return;
+        };
+        let editor = editor.read(cx);
+        let cursor_position = editor.selections.newest_anchor().head();
+        let multi_buffer_snapshot = editor.buffer().read(cx).snapshot(cx);
 
-        self.cached_entries = entries;
+        self.is_singleton = multi_buffer_snapshot.is_singleton();
+        self.cached_entries = new_highlights;
+        self.hovered_item_ix = None;
         self.rebuild_display_items(&multi_buffer_snapshot, cx);
 
-        if self.skip_next_scroll {
-            self.skip_next_scroll = false;
-        } else {
-            self.scroll_to_cursor_position(&cursor_position, &multi_buffer_snapshot);
-        }
+        let should_scroll = !mem::take(&mut self.skip_next_scroll);
+        self.update_selection_for_cursor(&cursor_position, &multi_buffer_snapshot, should_scroll);
+        self.sync_editor_highlight_for_selected_entry(cx);
         cx.notify();
     }
 
     fn rebuild_display_items(&mut self, snapshot: &MultiBufferSnapshot, cx: &App) {
         self.display_items.clear();
 
-        let mut last_range_end: Option<Anchor> = None;
-
+        let mut last_range_end = None;
         for (entry_ix, entry) in self.cached_entries.iter().enumerate() {
             if !self.should_show_entry(entry) {
                 continue;
@@ -531,7 +479,47 @@ impl HighlightsTreeView {
         }
     }
 
-    fn scroll_to_cursor_position(&mut self, cursor: &Anchor, snapshot: &MultiBufferSnapshot) {
+    fn update_selection_from_editor(&mut self, cx: &mut Context<Self>) {
+        let Some(editor_state) = self.editor.as_ref() else {
+            return;
+        };
+
+        let (cursor_position, multi_buffer_snapshot) = {
+            let editor = editor_state.editor.read(cx);
+            (
+                editor.selections.newest_anchor().head(),
+                editor.buffer().read(cx).snapshot(cx),
+            )
+        };
+
+        let should_update_selection = !mem::take(&mut self.skip_next_scroll);
+        if !should_update_selection {
+            self.sync_editor_highlight_for_selected_entry(cx);
+            return;
+        }
+
+        let should_notify =
+            self.update_selection_for_cursor(&cursor_position, &multi_buffer_snapshot, true);
+        self.sync_editor_highlight_for_selected_entry(cx);
+        if should_notify {
+            cx.notify();
+        }
+    }
+
+    fn update_selection_for_cursor(
+        &mut self,
+        cursor: &Anchor,
+        snapshot: &MultiBufferSnapshot,
+        should_scroll: bool,
+    ) -> bool {
+        let cursor_point = cursor.to_point(snapshot);
+        let Some((cursor_buffer, cursor_point)) = snapshot.point_to_buffer_point(cursor_point)
+        else {
+            let changed = self.selected_item_ix.take().is_some();
+            return changed;
+        };
+        let cursor_buffer_id = cursor_buffer.remote_id();
+
         let best = self
             .display_items
             .iter()
@@ -544,12 +532,17 @@ impl HighlightsTreeView {
                 _ => None,
             })
             .filter(|(_, _, entry)| {
-                entry.range.start.cmp(&cursor, snapshot).is_le()
-                    && cursor.cmp(&entry.range.end, snapshot).is_lt()
+                entry.buffer_id == cursor_buffer_id
+                    && entry.buffer_point_range.start <= cursor_point
+                    && cursor_point < entry.buffer_point_range.end
             })
             .min_by_key(|(_, _, entry)| {
                 (
-                    entry.buffer_point_range.end.row - entry.buffer_point_range.start.row,
+                    entry
+                        .buffer_point_range
+                        .end
+                        .row
+                        .saturating_sub(entry.buffer_point_range.start.row),
                     entry
                         .buffer_point_range
                         .end
@@ -559,11 +552,17 @@ impl HighlightsTreeView {
             })
             .map(|(display_ix, entry_ix, _)| (display_ix, entry_ix));
 
-        if let Some((display_ix, entry_ix)) = best {
-            self.selected_item_ix = Some(entry_ix);
+        let selected_item_ix = best.map(|(_, entry_ix)| entry_ix);
+        let changed = self.selected_item_ix != selected_item_ix;
+        self.selected_item_ix = selected_item_ix;
+
+        if should_scroll && let Some((display_ix, _)) = best {
             self.list_scroll_handle
                 .scroll_to_item(display_ix, ScrollStrategy::Center);
+            return true;
         }
+
+        changed
     }
 
     fn update_editor_with_range_for_entry(
@@ -707,6 +706,25 @@ impl HighlightsTreeView {
         );
     }
 
+    fn sync_editor_highlight_for_selected_entry(&self, cx: &mut Context<Self>) {
+        let Some(editor_state) = self.editor.as_ref() else {
+            return;
+        };
+        let key = cx.entity_id().as_u64() as usize;
+        let range = self
+            .selected_item_ix
+            .and_then(|entry_ix| self.cached_entries.get(entry_ix))
+            .map(|entry| entry.range.clone());
+
+        editor_state.editor.update(cx, |editor, cx| {
+            if let Some(range) = range {
+                Self::set_editor_highlights(editor, key, &[range], cx);
+            } else {
+                editor.clear_background_highlights(HighlightKey::HighlightsTreeView(key), cx);
+            }
+        });
+    }
+
     fn clear_editor_highlights(editor: &Entity<Editor>, cx: &mut Context<Self>) {
         let highlight_key = HighlightKey::HighlightsTreeView(cx.entity_id().as_u64() as usize);
         editor.update(cx, |editor, cx| {
@@ -1105,6 +1123,142 @@ fn excerpt_label_for(
     path_label.into()
 }
 
+fn build_highlight_entries(
+    HighlightRefreshInput {
+        multi_buffer_snapshot,
+        text_highlights,
+        semantic_highlights,
+        syntax_theme,
+    }: HighlightRefreshInput,
+) -> Vec<HighlightEntry> {
+    let mut entries = Vec::new();
+
+    for (key, text_highlights) in text_highlights {
+        for range in &text_highlights.1 {
+            let Some((range_display, buffer_id, buffer_point_range)) =
+                format_anchor_range(range, &multi_buffer_snapshot)
+            else {
+                continue;
+            };
+            entries.push(HighlightEntry {
+                range: range.clone(),
+                buffer_id,
+                range_display,
+                style: text_highlights.0,
+                category: HighlightCategory::Text(key),
+                buffer_point_range,
+            });
+        }
+    }
+
+    for highlight in semantic_highlights {
+        let Some((range_display, buffer_id, buffer_point_range)) =
+            format_anchor_range(&highlight.range, &multi_buffer_snapshot)
+        else {
+            continue;
+        };
+
+        entries.push(HighlightEntry {
+            range: highlight.range,
+            buffer_id,
+            range_display,
+            style: highlight.style,
+            category: highlight.category,
+            buffer_point_range,
+        });
+    }
+
+    for excerpt_range in multi_buffer_snapshot.excerpts() {
+        let Some(buffer_snapshot) =
+            multi_buffer_snapshot.buffer_for_id(excerpt_range.context.start.buffer_id)
+        else {
+            continue;
+        };
+
+        let start_offset = excerpt_range.context.start.to_offset(buffer_snapshot);
+        let end_offset = excerpt_range.context.end.to_offset(buffer_snapshot);
+        let range = start_offset..end_offset;
+
+        let captures = buffer_snapshot.captures(range, |grammar| {
+            grammar
+                .highlights_config
+                .as_ref()
+                .map(|config| &config.query)
+        });
+        let grammars = captures.grammars().to_vec();
+        let highlight_maps = grammars
+            .iter()
+            .map(|grammar| grammar.highlight_map())
+            .collect::<Vec<_>>();
+
+        for capture in captures {
+            let Some(highlight_id) = highlight_maps[capture.grammar_index].get(capture.index)
+            else {
+                continue;
+            };
+            let Some(style) = syntax_theme.get(highlight_id).cloned() else {
+                continue;
+            };
+
+            let theme_key = syntax_theme
+                .get_capture_name(highlight_id)
+                .map(|theme_key| SharedString::from(theme_key.to_string()));
+
+            let capture_name = grammars[capture.grammar_index]
+                .highlights_config
+                .as_ref()
+                .and_then(|config| config.query.capture_names().get(capture.index as usize))
+                .map(|capture_name| SharedString::from((*capture_name).to_string()))
+                .unwrap_or_else(|| SharedString::from("unknown"));
+
+            let start_anchor = buffer_snapshot.anchor_before(capture.node.start_byte());
+            let end_anchor = buffer_snapshot.anchor_after(capture.node.end_byte());
+
+            let start = multi_buffer_snapshot.anchor_in_excerpt(start_anchor);
+            let end = multi_buffer_snapshot.anchor_in_excerpt(end_anchor);
+
+            let (start, end) = match (start, end) {
+                (Some(start), Some(end)) => (start, end),
+                _ => continue,
+            };
+
+            let range = start..end;
+            let Some((range_display, buffer_id, buffer_point_range)) =
+                format_anchor_range(&range, &multi_buffer_snapshot)
+            else {
+                continue;
+            };
+
+            entries.push(HighlightEntry {
+                range,
+                buffer_id,
+                range_display,
+                style,
+                category: HighlightCategory::SyntaxToken {
+                    capture_name,
+                    theme_key,
+                },
+                buffer_point_range,
+            });
+        }
+    }
+
+    entries.sort_by(|a, b| {
+        a.buffer_id
+            .cmp(&b.buffer_id)
+            .then_with(|| a.buffer_point_range.start.cmp(&b.buffer_point_range.start))
+            .then_with(|| a.buffer_point_range.end.cmp(&b.buffer_point_range.end))
+            .then_with(|| a.category.cmp(&b.category))
+    });
+    entries.dedup_by(|a, b| {
+        a.buffer_id == b.buffer_id
+            && a.buffer_point_range == b.buffer_point_range
+            && a.category == b.category
+    });
+
+    entries
+}
+
 fn format_anchor_range(
     range: &Range<Anchor>,
     snapshot: &MultiBufferSnapshot,