Introduce fold_map::HighlightedChunks

zed/src/editor/buffer/mod.rs

@@ -752,12 +752,8 @@ impl Buffer {
         self.visible_text.chunks_in_range(start..end)
     }
 
-    pub fn highlighted_text_for_range<'a, T: ToOffset>(
+    pub fn highlighted_text_for_range<T: ToOffset>(&self, range: Range<T>) -> HighlightedChunks {
-        &'a self,
-        range: Range<T>,
-    ) -> impl Iterator<Item = (&'a str, Option<usize>)> {
         if let (Some(language), Some((tree, _))) = (&self.language, self.tree.as_ref()) {
-            let visible_text = &self.visible_text;
             let mut cursor = self
                 .query_cursor
                 .lock()
@@ -771,18 +767,14 @@ impl Buffer {
             let captures = cursor_ref.captures(
                 &language.highlight_query,
                 tree.root_node(),
-                move |node: tree_sitter::Node| {
+                TextProvider(&self.visible_text),
-                    visible_text
-                        .chunks_in_range(node.byte_range())
-                        .map(str::as_bytes)
-                },
             );
 
             HighlightedChunks {
                 captures: captures.peekable(),
                 chunks,
                 stack: Default::default(),
-                offset: start,
+                range: start..end,
                 query_cursor: Some(cursor),
                 buffer: self,
             }
@@ -2177,21 +2169,66 @@ impl<'a, F: Fn(&FragmentSummary) -> bool> Iterator for Edits<'a, F> {
     }
 }
 
-pub struct HighlightedChunks<'a, T: tree_sitter::TextProvider<'a>> {
+struct ByteChunks<'a>(rope::Chunks<'a>);
+
+impl<'a> Iterator for ByteChunks<'a> {
+    type Item = &'a [u8];
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.0.next().map(str::as_bytes)
+    }
+}
+
+struct TextProvider<'a>(&'a Rope);
+
+impl<'a> tree_sitter::TextProvider<'a> for TextProvider<'a> {
+    type I = ByteChunks<'a>;
+
+    fn text(&mut self, node: tree_sitter::Node) -> Self::I {
+        ByteChunks(self.0.chunks_in_range(node.byte_range()))
+    }
+}
+
+pub struct HighlightedChunks<'a> {
     chunks: Chunks<'a>,
-    captures: iter::Peekable<tree_sitter::QueryCaptures<'a, 'a, T>>,
+    captures: iter::Peekable<tree_sitter::QueryCaptures<'a, 'a, TextProvider<'a>>>,
     stack: Vec<(usize, usize)>,
-    offset: usize,
+    range: Range<usize>,
     query_cursor: Option<QueryCursor>,
     buffer: &'a Buffer,
 }
 
-impl<'a, T: tree_sitter::TextProvider<'a>> Iterator for HighlightedChunks<'a, T> {
+impl<'a> HighlightedChunks<'a> {
+    pub fn seek(&mut self, offset: usize) {
+        let language = self.buffer.language.as_ref().unwrap();
+        let tree = &self.buffer.tree.as_ref().unwrap().0;
+        let mut cursor = self.query_cursor.as_mut().unwrap();
+        let cursor_ref = unsafe { mem::transmute::<_, &'static mut QueryCursor>(&mut cursor) };
+
+        self.stack.clear();
+        self.range.start = offset;
+        self.chunks.seek(offset);
+        cursor.set_byte_range(self.range.start, self.range.end);
+        self.captures = cursor_ref
+            .captures(
+                &language.highlight_query,
+                tree.root_node(),
+                TextProvider(&self.buffer.visible_text),
+            )
+            .peekable();
+    }
+
+    pub fn offset(&self) -> usize {
+        self.range.start
+    }
+}
+
+impl<'a> Iterator for HighlightedChunks<'a> {
     type Item = (&'a str, Option<usize>);
 
     fn next(&mut self) -> Option<Self::Item> {
         while let Some((parent_capture_end, _)) = self.stack.last() {
-            if *parent_capture_end <= self.offset {
+            if *parent_capture_end <= self.range.start {
                 self.stack.pop();
             } else {
                 break;
@@ -2201,7 +2238,7 @@ impl<'a, T: tree_sitter::TextProvider<'a>> Iterator for HighlightedChunks<'a, T>
         let mut next_capture_start = usize::MAX;
         while let Some((mat, capture_ix)) = self.captures.peek() {
             let capture = mat.captures[*capture_ix as usize];
-            if self.offset < capture.node.start_byte() {
+            if self.range.start < capture.node.start_byte() {
                 next_capture_start = capture.node.start_byte();
                 break;
             } else {
@@ -2212,7 +2249,7 @@ impl<'a, T: tree_sitter::TextProvider<'a>> Iterator for HighlightedChunks<'a, T>
         }
 
         if let Some(chunk) = self.chunks.peek() {
-            let chunk_start = self.offset;
+            let chunk_start = self.range.start;
             let mut chunk_end = (self.chunks.offset() + chunk.len()).min(next_capture_start);
             let mut capture_ix = None;
             if let Some((parent_capture_end, parent_capture_ix)) = self.stack.last() {
@@ -2222,8 +2259,8 @@ impl<'a, T: tree_sitter::TextProvider<'a>> Iterator for HighlightedChunks<'a, T>
 
             let slice =
                 &chunk[chunk_start - self.chunks.offset()..chunk_end - self.chunks.offset()];
-            self.offset = chunk_end;
+            self.range.start = chunk_end;
-            if self.offset == self.chunks.offset() + chunk.len() {
+            if self.range.start == self.chunks.offset() + chunk.len() {
                 self.chunks.next().unwrap();
             }
 
@@ -2234,7 +2271,7 @@ impl<'a, T: tree_sitter::TextProvider<'a>> Iterator for HighlightedChunks<'a, T>
     }
 }
 
-impl<'a, T: tree_sitter::TextProvider<'a>> Drop for HighlightedChunks<'a, T> {
+impl<'a> Drop for HighlightedChunks<'a> {
     fn drop(&mut self) {
         let query_cursor = self.query_cursor.take().unwrap();
         let mut buffer_cursor = self.buffer.query_cursor.lock();

zed/src/editor/display_map/fold_map.rs

@@ -426,6 +426,24 @@ impl FoldMapSnapshot {
         }
     }
 
+    pub fn highlighted_chunks_at<'a>(
+        &'a self,
+        offset: DisplayOffset,
+        ctx: &'a AppContext,
+    ) -> HighlightedChunks<'a> {
+        let mut transform_cursor = self.transforms.cursor::<DisplayOffset, TransformSummary>();
+        transform_cursor.seek(&offset, SeekBias::Right, &());
+        let overshoot = offset.0 - transform_cursor.start().display.bytes;
+        let buffer_offset = transform_cursor.start().buffer.bytes + overshoot;
+        let buffer = self.buffer.read(ctx);
+        HighlightedChunks {
+            transform_cursor,
+            buffer_offset,
+            buffer_chunks: buffer.highlighted_text_for_range(buffer_offset..buffer.len()),
+            buffer_chunk: None,
+        }
+    }
+
     pub fn chars_at<'a>(
         &'a self,
         point: DisplayPoint,
@@ -720,6 +738,66 @@ impl<'a> Iterator for Chunks<'a> {
     }
 }
 
+pub struct HighlightedChunks<'a> {
+    transform_cursor: Cursor<'a, Transform, DisplayOffset, TransformSummary>,
+    buffer_chunks: buffer::HighlightedChunks<'a>,
+    buffer_chunk: Option<(&'a str, Option<usize>)>,
+    buffer_offset: usize,
+}
+
+impl<'a> Iterator for HighlightedChunks<'a> {
+    type Item = (&'a str, Option<usize>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let transform = if let Some(item) = self.transform_cursor.item() {
+            item
+        } else {
+            return None;
+        };
+
+        // If we're in a fold, then return the fold's display text and
+        // advance the transform and buffer cursors to the end of the fold.
+        if let Some(display_text) = transform.display_text {
+            self.buffer_chunk.take();
+            self.buffer_offset += transform.summary.buffer.bytes;
+            self.buffer_chunks.seek(self.buffer_offset);
+
+            while self.buffer_offset >= self.transform_cursor.end().buffer.bytes
+                && self.transform_cursor.item().is_some()
+            {
+                self.transform_cursor.next();
+            }
+
+            return Some((display_text, None));
+        }
+
+        // Retrieve a chunk from the current buffer cursor's location.
+        if self.buffer_chunk.is_none() {
+            self.buffer_chunk = self.buffer_chunks.next();
+        }
+
+        // Otherwise, take a chunk from the buffer's text.
+        if let Some((mut chunk, capture_ix)) = self.buffer_chunk {
+            let offset_in_chunk = self.buffer_offset - self.buffer_chunks.offset();
+            chunk = &chunk[offset_in_chunk..];
+
+            // Truncate the chunk so that it ends at the next fold.
+            let region_end = self.transform_cursor.end().buffer.bytes - self.buffer_offset;
+            if chunk.len() >= region_end {
+                chunk = &chunk[0..region_end];
+                self.transform_cursor.next();
+            } else {
+                self.buffer_chunk.take();
+            }
+
+            self.buffer_offset += chunk.len();
+            return Some((chunk, capture_ix));
+        }
+
+        None
+    }
+}
+
 impl<'a> sum_tree::Dimension<'a, TransformSummary> for DisplayPoint {
     fn add_summary(&mut self, summary: &'a TransformSummary) {
         self.0 += &summary.display.lines;