Re-use allocations in create_highlight_endpoints

crates/editor/src/display_map/custom_highlights.rs

@@ -2,12 +2,7 @@ use collections::BTreeMap;
 use gpui::HighlightStyle;
 use language::{Chunk, LanguageAwareStyling};
 use multi_buffer::{MultiBufferChunks, MultiBufferOffset, MultiBufferSnapshot};
-use std::{
-    cmp,
-    iter::{self, Peekable},
-    ops::Range,
-    vec,
-};
+use std::{cmp, ops::Range};
 
 use crate::display_map::{HighlightKey, SemanticTokensHighlights, TextHighlights};
 
@@ -17,7 +12,7 @@ pub struct CustomHighlightsChunks<'a> {
     offset: MultiBufferOffset,
     multibuffer_snapshot: &'a MultiBufferSnapshot,
 
-    highlight_endpoints: Peekable<vec::IntoIter<HighlightEndpoint>>,
+    highlight_endpoints: Vec<HighlightEndpoint>,
     active_highlights: BTreeMap<HighlightKey, HighlightStyle>,
     text_highlights: Option<&'a TextHighlights>,
     semantic_token_highlights: Option<&'a SemanticTokensHighlights>,
@@ -39,17 +34,20 @@ impl<'a> CustomHighlightsChunks<'a> {
         semantic_token_highlights: Option<&'a SemanticTokensHighlights>,
         multibuffer_snapshot: &'a MultiBufferSnapshot,
     ) -> Self {
+        let mut highlight_endpoints = Vec::new();
+        create_highlight_endpoints(
+            &range,
+            text_highlights,
+            semantic_token_highlights,
+            multibuffer_snapshot,
+            &mut highlight_endpoints,
+        );
         Self {
             buffer_chunks: multibuffer_snapshot.chunks(range.clone(), language_aware),
             buffer_chunk: None,
             offset: range.start,
             text_highlights,
-            highlight_endpoints: create_highlight_endpoints(
-                &range,
-                text_highlights,
-                semantic_token_highlights,
-                multibuffer_snapshot,
-            ),
+            highlight_endpoints,
             active_highlights: Default::default(),
             multibuffer_snapshot,
             semantic_token_highlights,
@@ -58,11 +56,12 @@ impl<'a> CustomHighlightsChunks<'a> {
 
     #[ztracing::instrument(skip_all)]
     pub fn seek(&mut self, new_range: Range<MultiBufferOffset>) {
-        self.highlight_endpoints = create_highlight_endpoints(
+        create_highlight_endpoints(
             &new_range,
             self.text_highlights,
             self.semantic_token_highlights,
             self.multibuffer_snapshot,
+            &mut self.highlight_endpoints,
         );
         self.offset = new_range.start;
         self.buffer_chunks.seek(new_range);
@@ -76,12 +75,13 @@ fn create_highlight_endpoints(
     text_highlights: Option<&TextHighlights>,
     semantic_token_highlights: Option<&SemanticTokensHighlights>,
     buffer: &MultiBufferSnapshot,
-) -> iter::Peekable<vec::IntoIter<HighlightEndpoint>> {
-    let mut highlight_endpoints = Vec::new();
+    highlight_endpoints: &mut Vec<HighlightEndpoint>,
+) {
+    highlight_endpoints.clear();
     if let Some(text_highlights) = text_highlights {
         let start = buffer.anchor_after(range.start);
         let end = buffer.anchor_after(range.end);
-        let mut v = Vec::new();
+        let mut text_highlights_scratch = Vec::new();
 
         for (&tag, text_highlights) in text_highlights.iter() {
             let style = text_highlights.0;
@@ -97,21 +97,21 @@ fn create_highlight_endpoints(
                 .unwrap_or_else(|i| i);
 
             let ranges_ = &ranges[start_ix..][..end_ix];
-            v.clear();
-            v.reserve(ranges_.len());
+            text_highlights_scratch.clear();
+            text_highlights_scratch.reserve(ranges_.len());
             highlight_endpoints.reserve(2 * ranges_.len());
 
             let mut iter = ranges_.iter();
             buffer.summaries_for_anchors_cb(
                 ranges_.iter().map(|r| &r.start),
                 |start: MultiBufferOffset| {
-                    v.push((start, iter.next().unwrap().end));
+                    text_highlights_scratch.push((start, iter.next().unwrap().end));
                 },
             );
-            v.sort_by(|a, b| a.1.cmp(&b.1, buffer));
-            let mut iter = v.iter();
+            text_highlights_scratch.sort_by(|a, b| a.1.cmp(&b.1, buffer));
+            let mut iter = text_highlights_scratch.iter();
             buffer.summaries_for_anchors_cb(
-                v.iter().map(|(_, end)| end),
+                text_highlights_scratch.iter().map(|(_, end)| end),
                 |end: MultiBufferOffset| {
                     let start = iter.next().unwrap().0;
                     if start == end {
@@ -134,6 +134,7 @@ fn create_highlight_endpoints(
     if let Some(semantic_token_highlights) = semantic_token_highlights {
         let start = buffer.anchor_after(range.start);
         let end = buffer.anchor_after(range.end);
+        let mut semantic_highlights_scratch = Vec::new();
         for buffer_id in buffer.buffer_ids_for_range(range.clone()) {
             let Some((semantic_token_highlights, interner)) =
                 semantic_token_highlights.get(&buffer_id)
@@ -160,20 +161,21 @@ fn create_highlight_endpoints(
                 .unwrap_or_else(|i| i);
 
             let ranges_ = &semantic_token_highlights[start_ix..][..end_ix];
-            let mut ranges_with_offsets = Vec::with_capacity(ranges_.len());
+            semantic_highlights_scratch.clear();
+            semantic_highlights_scratch.reserve(ranges_.len());
             highlight_endpoints.reserve(2 * ranges_.len());
 
             let mut iter = ranges_.iter();
             buffer.summaries_for_anchors_cb(
                 ranges_.iter().map(|token| &token.range.start),
                 |start: MultiBufferOffset| {
-                    ranges_with_offsets.push((start, iter.next().unwrap()));
+                    semantic_highlights_scratch.push((start, iter.next().unwrap()));
                 },
             );
-            ranges_with_offsets.sort_by(|a, b| a.1.range.end.cmp(&b.1.range.end, buffer));
-            let mut iter = ranges_with_offsets.iter();
+            semantic_highlights_scratch.sort_by(|a, b| a.1.range.end.cmp(&b.1.range.end, buffer));
+            let mut iter = semantic_highlights_scratch.iter();
             buffer.summaries_for_anchors_cb(
-                ranges_with_offsets
+                semantic_highlights_scratch
                     .iter()
                     .map(|(_, token)| &token.range.end),
                 |end: MultiBufferOffset| {
@@ -195,8 +197,7 @@ fn create_highlight_endpoints(
             );
         }
     }
-    highlight_endpoints.sort();
-    highlight_endpoints.into_iter().peekable()
+    highlight_endpoints.sort_by(|a, b| a.cmp(b).reverse());
 }
 
 impl<'a> Iterator for CustomHighlightsChunks<'a> {
@@ -205,14 +206,14 @@ impl<'a> Iterator for CustomHighlightsChunks<'a> {
     #[ztracing::instrument(skip_all)]
     fn next(&mut self) -> Option<Self::Item> {
         let mut next_highlight_endpoint = MultiBufferOffset(usize::MAX);
-        while let Some(endpoint) = self.highlight_endpoints.peek().copied() {
+        while let Some(endpoint) = self.highlight_endpoints.last().copied() {
             if endpoint.offset <= self.offset {
                 if let Some(style) = endpoint.style {
                     self.active_highlights.insert(endpoint.tag, style);
                 } else {
                     self.active_highlights.remove(&endpoint.tag);
                 }
-                self.highlight_endpoints.next();
+                self.highlight_endpoints.pop();
             } else {
                 next_highlight_endpoint = endpoint.offset;
                 break;