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,44 +80,62 @@ pub fn highlight_with_tree(
         return (plain(source), None);
     };
 
-    let mut parser = tree_sitter::Parser::new();
-    if parser.set_language(&grammar).is_err() {
-        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 None;
+                }
+                let query = tree_sitter::Query::new(&grammar, query_source).ok()?;
+                let roles = query
+                    .capture_names()
+                    .iter()
+                    .map(|name| role_for_capture(name))
+                    .collect();
+                Some(Highlighter {
+                    parser,
+                    query,
+                    roles,
+                })
+            })
+            .as_mut();
 
-    let Some(tree) = parser.parse(source, old_tree) else {
-        return (plain(source), None);
-    };
+        let Some(highlighter) = highlighter else {
+            return (plain(source), None);
+        };
 
-    let Ok(query) = tree_sitter::Query::new(&grammar, query_source) else {
-        return (plain(source), None);
-    };
+        let Some(tree) = highlighter.parser.parse(source, old_tree) else {
+            return (plain(source), None);
+        };
 
-    let roles: Vec<Option<Role>> = query
-        .capture_names()
-        .iter()
-        .map(|name| role_for_capture(name))
-        .collect();
+        // 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(&highlighter.query, tree.root_node(), source.as_bytes());
 
-    // 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());
-
-    while let Some(matched) = matches.next() {
-        for capture in matched.captures {
-            let Some(role) = roles.get(capture.index as usize).copied().flatten() else {
-                continue;
-            };
-            for byte in capture.node.byte_range() {
-                if let Some(slot) = role_at.get_mut(byte) {
-                    *slot = Some(role);
+        while let Some(matched) = matches.next() {
+            for capture in matched.captures {
+                let Some(role) = highlighter
+                    .roles
+                    .get(capture.index as usize)
+                    .copied()
+                    .flatten()
+                else {
+                    continue;
+                };
+                for byte in capture.node.byte_range() {
+                    if let Some(slot) = role_at.get_mut(byte) {
+                        *slot = Some(role);
+                    }
                 }
             }
         }
-    }
 
-    (build_lines(source, &role_at), Some(tree))
+        (build_lines(source, &role_at), Some(tree))
+    })
 }
 
 /// Source rendered without any grammar: plain text roles.