perf: speed up large markdown rendering

crates/editor/src/highlight.rs

@@ -3,9 +3,24 @@
 //! code without re-parsing it.
 
 use gpui::SharedString;
+use std::cell::RefCell;
+use std::collections::HashMap;
+use tree_sitter::StreamingIterator as _;
+
+thread_local! {
+    static HIGHLIGHTERS: RefCell<HashMap<Language, Option<Highlighter>>> =
+        RefCell::new(HashMap::new());
+}
+
+/// A grammar's parser and compiled highlight query.
+struct Highlighter {
+    parser: tree_sitter::Parser,
+    query: tree_sitter::Query,
+    roles: Vec<Option<Role>>,
+}
 
 /// A language source can highlight as.
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub enum Language {
     Rust,
     JavaScript,
@@ -50,8 +65,6 @@ pub struct HighlightedCode {
     pub lines: Vec<Line>,
 }
 
-use tree_sitter::StreamingIterator as _;
-
 /// Highlights `source` with `language`'s tree-sitter grammar.
 pub fn highlight(source: &str, language: Language) -> HighlightedCode {
     highlight_with_tree(source, language, None).0
@@ -67,33 +80,50 @@ pub fn highlight_with_tree(
         return (plain(source), None);
     };
 
+    HIGHLIGHTERS.with(|cache| {
+        let mut cache = cache.borrow_mut();
+        let highlighter = cache
+            .entry(language)
+            .or_insert_with(|| {
                 let mut parser = tree_sitter::Parser::new();
                 if parser.set_language(&grammar).is_err() {
-        return (plain(source), None);
+                    return None;
                 }
-
+                let query = tree_sitter::Query::new(&grammar, query_source).ok()?;
-    let Some(tree) = parser.parse(source, old_tree) else {
+                let roles = query
-        return (plain(source), None);
-    };
-
-    let Ok(query) = tree_sitter::Query::new(&grammar, query_source) else {
-        return (plain(source), None);
-    };
-
-    let roles: Vec<Option<Role>> = query
                     .capture_names()
                     .iter()
                     .map(|name| role_for_capture(name))
                     .collect();
+                Some(Highlighter {
+                    parser,
+                    query,
+                    roles,
+                })
+            })
+            .as_mut();
+
+        let Some(highlighter) = highlighter else {
+            return (plain(source), None);
+        };
+
+        let Some(tree) = highlighter.parser.parse(source, old_tree) else {
+            return (plain(source), None);
+        };
 
         // One role per byte; matches arrive in query order
         let mut role_at = vec![None; source.len()];
         let mut cursor = tree_sitter::QueryCursor::new();
-    let mut matches = cursor.matches(&query, tree.root_node(), source.as_bytes());
+        let mut matches = cursor.matches(&highlighter.query, tree.root_node(), source.as_bytes());
 
         while let Some(matched) = matches.next() {
             for capture in matched.captures {
-            let Some(role) = roles.get(capture.index as usize).copied().flatten() else {
+                let Some(role) = highlighter
+                    .roles
+                    .get(capture.index as usize)
+                    .copied()
+                    .flatten()
+                else {
                     continue;
                 };
                 for byte in capture.node.byte_range() {
@@ -105,6 +135,7 @@ pub fn highlight_with_tree(
         }
 
         (build_lines(source, &role_at), Some(tree))
+    })
 }
 
 /// Source rendered without any grammar: plain text roles.

crates/editor/src/markdown.rs

@@ -7,7 +7,7 @@
 use gpui::prelude::FluentBuilder as _;
 use gpui::{
     AnyElement, App, Context, Div, ElementId, FontStyle, FontWeight, HighlightStyle,
-    InteractiveElement as _, InteractiveText, IntoElement, ListAlignment, ListState,
+    InteractiveElement as _, InteractiveText, IntoElement, ListAlignment, ListOffset, ListState,
     ParentElement as _, Pixels, Render, SharedString, StatefulInteractiveElement as _,
     StrikethroughStyle, Styled as _, StyledText, UnderlineStyle, Window, div, list, px, relative,
 };
@@ -190,7 +190,7 @@ impl MarkdownViewState {
             source,
             blocks,
             highlight_code: true,
-            handle: ListState::new(count, ListAlignment::Top, px(400.)).measure_all(),
+            handle: ListState::new(count, ListAlignment::Top, px(400.)),
             a11y_label: None,
         }
     }
@@ -207,7 +207,11 @@ impl MarkdownViewState {
     where
         S: Into<SharedString>,
     {
-        self.source = source.into();
+        let source = source.into();
+        if self.source == source {
+            return;
+        }
+        self.source = source;
         self.reparse();
         cx.notify();
     }
@@ -233,10 +237,13 @@ impl MarkdownViewState {
         self.highlight_code
     }
 
-    /// Scrolls `block` into view.
+    /// Scrolls `block` to the top of the view.
     pub fn scroll_to_block(&self, block: usize) {
-        self.handle
+        let item_ix = block.min(self.blocks.len().saturating_sub(1));
-            .scroll_to_reveal_item(block.min(self.blocks.len().saturating_sub(1)));
+        self.handle.scroll_to(ListOffset {
+            item_ix,
+            offset_in_item: px(0.),
+        });
     }
 
     /// Overrides the accessibility label of the view.

crates/editor/tests/markdown.rs

@@ -447,3 +447,50 @@ fn parse_body(source: &str) -> Vec<MarkdownBlock> {
     let state = MarkdownViewState::new("parse-md", source);
     state.blocks().to_vec()
 }
+/// Mirrors the gallery's "Large file" sample: five blocks per section.
+fn large_markdown(sections: usize) -> String {
+    let mut source = String::new();
+    for section in 0..sections {
+        source.push_str(&format!("## Section {section}\n\n"));
+        source.push_str(&format!(
+            "Handler {section} retries stale sockets and logs every outcome.\n\n"
+        ));
+        source.push_str(&format!(
+            "- [ ] Audit handler {section}\n- See [issue {section}](https://example.com/issues/{section})\n\n"
+        ));
+        source.push_str("```rust\nfn handle() -> bool { true }\n```\n\n");
+    }
+    source
+}
+
+#[test]
+fn a_large_document_swaps_and_renders_without_freezing() {
+    let source = large_markdown(1000);
+    let started = std::time::Instant::now();
+
+    let (_cx, mut visual, state) = view_window("# small\n");
+    state.update_in(&mut visual, |state, _window, cx| {
+        state.set_source(source, cx)
+    });
+    visual.update(|window, cx| {
+        window.draw(cx).clear(cx);
+    });
+
+    let elapsed = started.elapsed();
+    assert!(
+        elapsed < std::time::Duration::from_secs(15),
+        "swapping to a large document took {elapsed:?}; the frame must not
+         measure every block"
+    );
+
+    let count = _cx.update(|cx| state.read(cx).blocks().len());
+    assert_eq!(count, 5000, "five blocks per section");
+    assert!(
+        visual.debug_bounds("test-md-block-0").is_some(),
+        "the first block renders after the swap"
+    );
+    assert!(
+        visual.debug_bounds("test-md-block-4999").is_none(),
+        "blocks past the viewport never render"
+    );
+}