highlighter: Use sum_tree to store the highlight cache. (#1198)

Cargo.lock

@@ -2583,15 +2583,6 @@ dependencies = [
  "libc",
 ]
 
-[[package]]
-name = "ftree"
-version = "1.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9ae0379499242d3b9355c5069b43b9417def8c9b09903b930db1fe49318dc9e9"
-dependencies = [
- "serde",
-]
-
 [[package]]
 name = "funty"
 version = "2.0.0"
@@ -3159,7 +3150,6 @@ dependencies = [
  "gpui",
  "gpui-component-macros",
  "html5ever 0.27.0",
- "indexset",
  "indoc",
  "itertools 0.13.0",
  "markdown",
@@ -3177,6 +3167,7 @@ dependencies = [
  "serde_repr",
  "smallvec",
  "smol 1.3.0",
+ "sum_tree",
  "tracing",
  "tree-sitter",
  "tree-sitter-bash",
@@ -3849,15 +3840,6 @@ dependencies = [
  "serde",
 ]
 
-[[package]]
-name = "indexset"
-version = "0.12.3"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ff794cab64c942437d60272e215f923d466b23dfa6c999cdd0cafe5b6d170805"
-dependencies = [
- "ftree",
-]
-
 [[package]]
 name = "indicatif"
 version = "0.17.11"

Cargo.toml

@@ -17,6 +17,7 @@ resolver = "2"
 [workspace.dependencies]
 gpui = { git = "https://github.com/zed-industries/zed.git" }
 reqwest_client = { git = "https://github.com/zed-industries/zed.git" }
+sum_tree = { git = "https://github.com/zed-industries/zed.git" }
 gpui-component = { path = "crates/ui" }
 gpui-component-macros = { path = "crates/macros" }
 story = { path = "crates/story" }

crates/ui/Cargo.toml

@@ -49,6 +49,7 @@ tree-sitter-languages = [
 
 [dependencies]
 gpui.workspace = true
+sum_tree.workspace = true
 gpui-component-macros.workspace = true
 rust-i18n.workspace = true
 schemars.workspace = true
@@ -87,7 +88,6 @@ markup5ever_rcdom = "0.3.0"
 chrono = "0.4.38"
 
 # Code Editor
-indexset = "0.12.2"
 tree-sitter = "0.25.4"
 tree-sitter-json = "0.24.8"
 tree-sitter-bash = { version = "0.23.3", optional = true }

crates/ui/src/highlighter/highlighter.rs

@@ -1,12 +1,9 @@
 use super::HighlightTheme;
 use crate::highlighter::LanguageRegistry;
+
 use anyhow::{anyhow, Context, Result};
 use gpui::{App, HighlightStyle, SharedString};
-use indexset::BTreeMap;
+use std::{collections::HashMap, ops::Range, usize};
-use std::{
-    collections::HashMap,
-    ops::{Bound, Range},
-};
 use tree_sitter::{
     InputEdit, Node, Parser, Point, Query, QueryCursor, QueryMatch, StreamingIterator, Tree,
 };
@@ -39,7 +36,85 @@ pub struct SyntaxHighlighter {
     ///
     /// - The `key` is the `start` of the range.
     /// -The `value` is a tuple of the range (in the entire text) and the highlight name.
-    cache: BTreeMap<usize, (Range<usize>, SharedString)>,
+    cache: sum_tree::SumTree<HighlightItem>,
+}
+
+#[derive(Debug, Default, Clone)]
+struct HighlightSummary {
+    count: usize,
+    start: usize,
+    end: usize,
+    min_start: usize,
+    max_end: usize,
+}
+
+#[derive(Debug, Default, Clone)]
+struct HighlightItem {
+    range: Range<usize>,
+    /// The highlight name, like `function`, `string`, `comment`, etc.
+    name: SharedString,
+}
+
+impl HighlightItem {
+    pub fn new(range: Range<usize>, name: impl Into<SharedString>) -> Self {
+        Self {
+            range,
+            name: name.into(),
+        }
+    }
+}
+
+impl sum_tree::Item for HighlightItem {
+    type Summary = HighlightSummary;
+    fn summary(&self, _cx: &()) -> Self::Summary {
+        HighlightSummary {
+            count: 1,
+            start: self.range.start,
+            end: self.range.end,
+            min_start: self.range.start,
+            max_end: self.range.end,
+        }
+    }
+}
+
+impl sum_tree::Summary for HighlightSummary {
+    type Context = ();
+    fn zero(_: &Self::Context) -> Self {
+        HighlightSummary {
+            count: 0,
+            start: usize::MIN,
+            end: usize::MAX,
+            min_start: usize::MAX,
+            max_end: usize::MIN,
+        }
+    }
+
+    fn add_summary(&mut self, other: &Self, _: &Self::Context) {
+        self.min_start = self.min_start.min(other.min_start);
+        self.max_end = self.max_end.max(other.max_end);
+        self.start = other.start;
+        self.end = other.end;
+        self.count += other.count;
+    }
+}
+
+impl<'a> sum_tree::Dimension<'a, HighlightSummary> for usize {
+    fn zero(_: &()) -> Self {
+        0
+    }
+
+    fn add_summary(&mut self, _: &'a HighlightSummary, _: &()) {}
+}
+
+impl<'a> sum_tree::Dimension<'a, HighlightSummary> for Range<usize> {
+    fn zero(_: &()) -> Self {
+        Default::default()
+    }
+
+    fn add_summary(&mut self, summary: &'a HighlightSummary, _: &()) {
+        self.start = summary.start;
+        self.end = summary.end;
+    }
 }
 
 impl SyntaxHighlighter {
@@ -180,7 +255,7 @@ impl SyntaxHighlighter {
             parser,
             old_tree: None,
             text: SharedString::new(""),
-            cache: BTreeMap::new(),
+            cache: sum_tree::SumTree::new(&()),
             locals_pattern_index,
             highlights_pattern_index,
             non_local_variable_patterns,
@@ -261,7 +336,7 @@ impl SyntaxHighlighter {
         let source = self.text.as_bytes();
         let mut query_cursor = QueryCursor::new();
         let root_node = tree.root_node();
-        self.cache.clear();
+        self.cache = sum_tree::SumTree::new(&());
 
         let mut matches = query_cursor.matches(&query, root_node, source);
         while let Some(m) = matches.next() {
@@ -273,7 +348,8 @@ impl SyntaxHighlighter {
                     let styles = self.handle_injection(&language_name, content_node, source, cx);
                     for (node_range, highlight_name) in styles {
                         self.cache
-                            .insert(node_range.start, (node_range, highlight_name.into()));
+                            .push(HighlightItem::new(node_range.clone(), highlight_name), &());
+                        // .insert(node_range.start, (node_range, highlight_name.into()));
                     }
                 }
 
@@ -291,28 +367,34 @@ impl SyntaxHighlighter {
                 let highlight_name = SharedString::from(highlight_name.to_string());
 
                 // Merge near range and same highlight name
-                let last_item = self.cache.last_key_value().map(|kv| kv.1);
+                let last_item = self.cache.last();
-                let last_range = last_item.map(|(range, _)| range).unwrap_or(&(0..0));
+                let last_range = last_item.map(|item| &item.range).unwrap_or(&(0..0));
-                let last_highlight_name = last_item.map(|(_, name)| name.clone());
+                let last_highlight_name = last_item.map(|item| item.name.clone());
 
                 if last_range.end <= node_range.start
                     && last_highlight_name.as_ref() == Some(&highlight_name)
                 {
-                    self.cache.insert(
+                    self.cache.push(
-                        last_range.start,
+                        HighlightItem::new(
-                        (last_range.start..node_range.end, highlight_name.clone()),
+                            last_range.start..node_range.end,
+                            highlight_name.clone(),
+                        ),
+                        &(),
                     );
                 } else if last_range == &node_range {
                     // case:
                     // last_range: 213..220, last_highlight_name: Some("property")
                     // last_range: 213..220, last_highlight_name: Some("string")
-                    self.cache.insert(
+                    self.cache.push(
-                        node_range.start,
+                        HighlightItem::new(
-                        (node_range, last_highlight_name.unwrap_or(highlight_name)),
+                            node_range,
+                            last_highlight_name.unwrap_or(highlight_name),
+                        ),
+                        &(),
                     );
                 } else {
                     self.cache
-                        .insert(node_range.start, (node_range, highlight_name.clone()));
+                        .push(HighlightItem::new(node_range, highlight_name.clone()), &());
                 }
             }
         }
@@ -473,28 +555,29 @@ impl SyntaxHighlighter {
     ) -> Vec<(Range<usize>, HighlightStyle)> {
         let mut styles = vec![];
         let start_offset = range.start;
+
+        let mut cursor = self.cache.cursor::<usize>(&());
+        let bias = if start_offset == 0 {
+            sum_tree::Bias::Right
+        } else {
+            sum_tree::Bias::Left
+        };
+
+        let left_items = cursor.slice(&start_offset, bias);
+        let mut filter = left_items.filter::<_, Range<usize>>(&(), move |sum| {
+            range.start <= sum.max_end && range.end >= sum.min_start
+        });
+        filter.next();
+
         let mut last_range = start_offset..start_offset;
+        // let mut iter_count = 0;
+        while let Some(item) = filter.item() {
+            // iter_count += 1;
+            let node_range = &item.range;
+            let name = &item.name;
 
-        // NOTE: Iterate over the cache and print the range and style for each item.
-        // for (_, (range, style)) in self.cache.iter() {
-        //     println!("-- range: {:?}, style: {:?}", range, style);
-        // }
-
-        let mut cursor = self.cache.lower_bound(Bound::Included(&range.start));
-        // Move to the previous item if the current item is not the start of the range.
-        // This is for case like JsDoc, where token may contains multiple lines.
-        if cursor.key() != Some(&range.start) {
-            cursor.move_prev();
-        }
-
-        while let Some((node_range, name)) = cursor.value() {
-            // Break loop if the node_range is out of the range
-            if node_range.start > range.end {
-                break;
-            }
-
-            let mut node_range = node_range.start.max(range.start)..node_range.end.min(range.end);
             // Avoid start larger than end
+            let mut node_range = node_range.start.max(range.start)..node_range.end.min(range.end);
             if node_range.start > node_range.end {
                 node_range.end = node_range.start;
             }
@@ -510,8 +593,9 @@ impl SyntaxHighlighter {
                 theme.style(name.as_ref()).unwrap_or_default(),
             ));
 
-            cursor.move_next();
+            filter.next();
         }
+        // dbg!(iter_count);
 
         // If the matched styles is empty, return a default range.
         if styles.len() == 0 {