Add `language::BufferSnapshot::highlighted_text_for_range` (#24060)

Michael Sloan created

In support of work on

https://github.com/zed-industries/zed/tree/new-ui-for-edit-prediction-with-lsp-completions,
where we want to be able to extract a range of the buffer as
`HighlightedText`.

Release Notes:

- N/A

Change summary

crates/language/src/buffer.rs | 211 +++++++++++++++++++++---------------
1 file changed, 125 insertions(+), 86 deletions(-)

Detailed changes

crates/language/src/buffer.rs 🔗

@@ -588,6 +588,99 @@ pub struct Runnable {
     pub buffer: BufferId,
 }
 
+#[derive(Default, Clone, Debug)]
+pub struct HighlightedText {
+    pub text: SharedString,
+    pub highlights: Vec<(Range<usize>, HighlightStyle)>,
+}
+
+#[derive(Default, Debug)]
+struct HighlightedTextBuilder {
+    pub text: String,
+    pub highlights: Vec<(Range<usize>, HighlightStyle)>,
+}
+
+impl HighlightedText {
+    pub fn from_buffer_range<T: ToOffset>(
+        range: Range<T>,
+        snapshot: &text::BufferSnapshot,
+        syntax_snapshot: &SyntaxSnapshot,
+        override_style: Option<HighlightStyle>,
+        syntax_theme: &SyntaxTheme,
+    ) -> Self {
+        let mut highlighted_text = HighlightedTextBuilder::default();
+        highlighted_text.add_text_from_buffer_range(
+            range,
+            snapshot,
+            syntax_snapshot,
+            override_style,
+            syntax_theme,
+        );
+        highlighted_text.build()
+    }
+}
+
+impl HighlightedTextBuilder {
+    pub fn build(self) -> HighlightedText {
+        HighlightedText {
+            text: self.text.into(),
+            highlights: self.highlights,
+        }
+    }
+
+    pub fn add_text_from_buffer_range<T: ToOffset>(
+        &mut self,
+        range: Range<T>,
+        snapshot: &text::BufferSnapshot,
+        syntax_snapshot: &SyntaxSnapshot,
+        override_style: Option<HighlightStyle>,
+        syntax_theme: &SyntaxTheme,
+    ) {
+        let range = range.to_offset(snapshot);
+        for chunk in Self::highlighted_chunks(range, snapshot, syntax_snapshot) {
+            let start = self.text.len();
+            self.text.push_str(chunk.text);
+            let end = self.text.len();
+
+            if let Some(mut highlight_style) = chunk
+                .syntax_highlight_id
+                .and_then(|id| id.style(syntax_theme))
+            {
+                if let Some(override_style) = override_style {
+                    highlight_style.highlight(override_style);
+                }
+                self.highlights.push((start..end, highlight_style));
+            } else if let Some(override_style) = override_style {
+                self.highlights.push((start..end, override_style));
+            }
+        }
+    }
+
+    fn highlighted_chunks<'a>(
+        range: Range<usize>,
+        snapshot: &'a text::BufferSnapshot,
+        syntax_snapshot: &'a SyntaxSnapshot,
+    ) -> BufferChunks<'a> {
+        let captures = syntax_snapshot.captures(range.clone(), snapshot, |grammar| {
+            grammar.highlights_query.as_ref()
+        });
+
+        let highlight_maps = captures
+            .grammars()
+            .iter()
+            .map(|grammar| grammar.highlight_map())
+            .collect();
+
+        BufferChunks::new(
+            snapshot.as_rope(),
+            range,
+            Some((captures, highlight_maps)),
+            false,
+            None,
+        )
+    }
+}
+
 #[derive(Clone)]
 pub struct EditPreview {
     old_snapshot: text::BufferSnapshot,
@@ -595,12 +688,6 @@ pub struct EditPreview {
     syntax_snapshot: SyntaxSnapshot,
 }
 
-#[derive(Default, Clone, Debug)]
-pub struct HighlightedText {
-    pub text: SharedString,
-    pub highlights: Vec<(Range<usize>, HighlightStyle)>,
-}
-
 impl EditPreview {
     pub fn highlight_edits(
         &self,
@@ -613,8 +700,7 @@ impl EditPreview {
             return HighlightedText::default();
         };
 
-        let mut text = String::new();
-        let mut highlights = Vec::new();
+        let mut highlighted_text = HighlightedTextBuilder::default();
 
         let mut offset_in_preview_snapshot = visible_range_in_preview_snapshot.start;
 
@@ -626,6 +712,7 @@ impl EditPreview {
             background_color: Some(cx.theme().status().deleted_background),
             ..Default::default()
         };
+        let syntax_theme = cx.theme().syntax();
 
         for (range, edit_text) in edits {
             let edit_new_end_in_preview_snapshot = range
@@ -637,111 +724,48 @@ impl EditPreview {
             let unchanged_range_in_preview_snapshot =
                 offset_in_preview_snapshot..edit_start_in_preview_snapshot;
             if !unchanged_range_in_preview_snapshot.is_empty() {
-                Self::highlight_text(
-                    unchanged_range_in_preview_snapshot.clone(),
-                    &mut text,
-                    &mut highlights,
-                    None,
+                highlighted_text.add_text_from_buffer_range(
+                    unchanged_range_in_preview_snapshot,
                     &self.applied_edits_snapshot,
                     &self.syntax_snapshot,
-                    cx,
+                    None,
+                    &syntax_theme,
                 );
             }
 
             let range_in_current_snapshot = range.to_offset(current_snapshot);
             if include_deletions && !range_in_current_snapshot.is_empty() {
-                Self::highlight_text(
-                    range_in_current_snapshot.clone(),
-                    &mut text,
-                    &mut highlights,
-                    Some(deletion_highlight_style),
+                highlighted_text.add_text_from_buffer_range(
+                    range_in_current_snapshot,
                     &current_snapshot.text,
                     &current_snapshot.syntax,
-                    cx,
+                    Some(deletion_highlight_style),
+                    &syntax_theme,
                 );
             }
 
             if !edit_text.is_empty() {
-                Self::highlight_text(
+                highlighted_text.add_text_from_buffer_range(
                     edit_start_in_preview_snapshot..edit_new_end_in_preview_snapshot,
-                    &mut text,
-                    &mut highlights,
-                    Some(insertion_highlight_style),
                     &self.applied_edits_snapshot,
                     &self.syntax_snapshot,
-                    cx,
+                    Some(insertion_highlight_style),
+                    &syntax_theme,
                 );
             }
 
             offset_in_preview_snapshot = edit_new_end_in_preview_snapshot;
         }
 
-        Self::highlight_text(
+        highlighted_text.add_text_from_buffer_range(
             offset_in_preview_snapshot..visible_range_in_preview_snapshot.end,
-            &mut text,
-            &mut highlights,
-            None,
             &self.applied_edits_snapshot,
             &self.syntax_snapshot,
-            cx,
+            None,
+            &syntax_theme,
         );
 
-        HighlightedText {
-            text: text.into(),
-            highlights,
-        }
-    }
-
-    fn highlight_text(
-        range: Range<usize>,
-        text: &mut String,
-        highlights: &mut Vec<(Range<usize>, HighlightStyle)>,
-        override_style: Option<HighlightStyle>,
-        snapshot: &text::BufferSnapshot,
-        syntax_snapshot: &SyntaxSnapshot,
-        cx: &App,
-    ) {
-        for chunk in Self::highlighted_chunks(range, snapshot, syntax_snapshot) {
-            let start = text.len();
-            text.push_str(chunk.text);
-            let end = text.len();
-
-            if let Some(mut highlight_style) = chunk
-                .syntax_highlight_id
-                .and_then(|id| id.style(cx.theme().syntax()))
-            {
-                if let Some(override_style) = override_style {
-                    highlight_style.highlight(override_style);
-                }
-                highlights.push((start..end, highlight_style));
-            } else if let Some(override_style) = override_style {
-                highlights.push((start..end, override_style));
-            }
-        }
-    }
-
-    fn highlighted_chunks<'a>(
-        range: Range<usize>,
-        snapshot: &'a text::BufferSnapshot,
-        syntax_snapshot: &'a SyntaxSnapshot,
-    ) -> BufferChunks<'a> {
-        let captures = syntax_snapshot.captures(range.clone(), snapshot, |grammar| {
-            grammar.highlights_query.as_ref()
-        });
-
-        let highlight_maps = captures
-            .grammars()
-            .iter()
-            .map(|grammar| grammar.highlight_map())
-            .collect();
-
-        BufferChunks::new(
-            snapshot.as_rope(),
-            range,
-            Some((captures, highlight_maps)),
-            false,
-            None,
-        )
+        highlighted_text.build()
     }
 
     fn compute_visible_range(&self, edits: &[(Range<Anchor>, String)]) -> Option<Range<usize>> {
@@ -2982,6 +3006,21 @@ impl BufferSnapshot {
         BufferChunks::new(self.text.as_rope(), range, syntax, diagnostics, Some(self))
     }
 
+    pub fn highlighted_text_for_range<T: ToOffset>(
+        &self,
+        range: Range<T>,
+        override_style: Option<HighlightStyle>,
+        syntax_theme: &SyntaxTheme,
+    ) -> HighlightedText {
+        HighlightedText::from_buffer_range(
+            range,
+            &self.text,
+            &self.syntax,
+            override_style,
+            syntax_theme,
+        )
+    }
+
     /// Invokes the given callback for each line of text in the given range of the buffer.
     /// Uses callback to avoid allocating a string for each line.
     fn for_each_line(&self, range: Range<Point>, mut callback: impl FnMut(u32, &str)) {