custom_highlights.rs

  1use collections::BTreeMap;
  2use gpui::HighlightStyle;
  3use language::Chunk;
  4use multi_buffer::{MultiBufferChunks, MultiBufferOffset, MultiBufferSnapshot, ToOffset as _};
  5use std::{
  6    cmp,
  7    iter::{self, Peekable},
  8    ops::Range,
  9    vec,
 10};
 11
 12use crate::display_map::{HighlightKey, TextHighlights};
 13
 14pub struct CustomHighlightsChunks<'a> {
 15    buffer_chunks: MultiBufferChunks<'a>,
 16    buffer_chunk: Option<Chunk<'a>>,
 17    offset: MultiBufferOffset,
 18    multibuffer_snapshot: &'a MultiBufferSnapshot,
 19
 20    highlight_endpoints: Peekable<vec::IntoIter<HighlightEndpoint>>,
 21    active_highlights: BTreeMap<HighlightKey, HighlightStyle>,
 22    text_highlights: Option<&'a TextHighlights>,
 23}
 24
 25#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 26struct HighlightEndpoint {
 27    offset: MultiBufferOffset,
 28    tag: HighlightKey,
 29    style: Option<HighlightStyle>,
 30}
 31
 32impl<'a> CustomHighlightsChunks<'a> {
 33    #[ztracing::instrument(skip_all)]
 34    pub fn new(
 35        range: Range<MultiBufferOffset>,
 36        language_aware: bool,
 37        text_highlights: Option<&'a TextHighlights>,
 38        multibuffer_snapshot: &'a MultiBufferSnapshot,
 39    ) -> Self {
 40        Self {
 41            buffer_chunks: multibuffer_snapshot.chunks(range.clone(), language_aware),
 42            buffer_chunk: None,
 43            offset: range.start,
 44            text_highlights,
 45            highlight_endpoints: create_highlight_endpoints(
 46                &range,
 47                text_highlights,
 48                multibuffer_snapshot,
 49            ),
 50            active_highlights: Default::default(),
 51            multibuffer_snapshot,
 52        }
 53    }
 54
 55    #[ztracing::instrument(skip_all)]
 56    pub fn seek(&mut self, new_range: Range<MultiBufferOffset>) {
 57        self.highlight_endpoints =
 58            create_highlight_endpoints(&new_range, self.text_highlights, self.multibuffer_snapshot);
 59        self.offset = new_range.start;
 60        self.buffer_chunks.seek(new_range);
 61        self.buffer_chunk.take();
 62        self.active_highlights.clear()
 63    }
 64}
 65
 66fn create_highlight_endpoints(
 67    range: &Range<MultiBufferOffset>,
 68    text_highlights: Option<&TextHighlights>,
 69    buffer: &MultiBufferSnapshot,
 70) -> iter::Peekable<vec::IntoIter<HighlightEndpoint>> {
 71    let mut highlight_endpoints = Vec::new();
 72    if let Some(text_highlights) = text_highlights {
 73        let start = buffer.anchor_after(range.start);
 74        let end = buffer.anchor_after(range.end);
 75        for (&tag, text_highlights) in text_highlights.iter() {
 76            let style = text_highlights.0;
 77            let ranges = &text_highlights.1;
 78
 79            let start_ix = ranges
 80                .binary_search_by(|probe| probe.end.cmp(&start, buffer).then(cmp::Ordering::Less))
 81                .unwrap_or_else(|i| i);
 82
 83            for range in &ranges[start_ix..] {
 84                if range.start.cmp(&end, buffer).is_ge() {
 85                    break;
 86                }
 87
 88                let start = range.start.to_offset(buffer);
 89                let end = range.end.to_offset(buffer);
 90                if start == end {
 91                    continue;
 92                }
 93                highlight_endpoints.push(HighlightEndpoint {
 94                    offset: start,
 95                    tag,
 96                    style: Some(style),
 97                });
 98                highlight_endpoints.push(HighlightEndpoint {
 99                    offset: end,
100                    tag,
101                    style: None,
102                });
103            }
104        }
105        highlight_endpoints.sort();
106    }
107    highlight_endpoints.into_iter().peekable()
108}
109
110impl<'a> Iterator for CustomHighlightsChunks<'a> {
111    type Item = Chunk<'a>;
112
113    #[ztracing::instrument(skip_all)]
114    fn next(&mut self) -> Option<Self::Item> {
115        let mut next_highlight_endpoint = MultiBufferOffset(usize::MAX);
116        while let Some(endpoint) = self.highlight_endpoints.peek().copied() {
117            if endpoint.offset <= self.offset {
118                if let Some(style) = endpoint.style {
119                    self.active_highlights.insert(endpoint.tag, style);
120                } else {
121                    self.active_highlights.remove(&endpoint.tag);
122                }
123                self.highlight_endpoints.next();
124            } else {
125                next_highlight_endpoint = endpoint.offset;
126                break;
127            }
128        }
129
130        let chunk = match &mut self.buffer_chunk {
131            Some(it) => it,
132            slot => slot.insert(self.buffer_chunks.next()?),
133        };
134        while chunk.text.is_empty() {
135            *chunk = self.buffer_chunks.next()?;
136        }
137
138        let split_idx = chunk.text.len().min(next_highlight_endpoint - self.offset);
139        let (prefix, suffix) = chunk.text.split_at(split_idx);
140        self.offset += prefix.len();
141
142        let mask = 1u128.unbounded_shl(split_idx as u32).wrapping_sub(1);
143        let chars = chunk.chars & mask;
144        let tabs = chunk.tabs & mask;
145        let mut prefix = Chunk {
146            text: prefix,
147            chars,
148            tabs,
149            ..chunk.clone()
150        };
151
152        chunk.chars = chunk.chars.unbounded_shr(split_idx as u32);
153        chunk.tabs = chunk.tabs.unbounded_shr(split_idx as u32);
154        chunk.text = suffix;
155        if !self.active_highlights.is_empty() {
156            prefix.highlight_style = self
157                .active_highlights
158                .values()
159                .copied()
160                .reduce(|acc, active_highlight| acc.highlight(active_highlight));
161        }
162        Some(prefix)
163    }
164}
165
166impl PartialOrd for HighlightEndpoint {
167    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
168        Some(self.cmp(other))
169    }
170}
171
172impl Ord for HighlightEndpoint {
173    fn cmp(&self, other: &Self) -> cmp::Ordering {
174        self.offset
175            .cmp(&other.offset)
176            .then_with(|| self.style.is_some().cmp(&other.style.is_some()))
177    }
178}
179
180#[cfg(test)]
181mod tests {
182    use std::{any::TypeId, sync::Arc};
183
184    use super::*;
185    use crate::MultiBuffer;
186    use gpui::App;
187    use rand::prelude::*;
188    use util::RandomCharIter;
189
190    #[gpui::test(iterations = 100)]
191    fn test_random_chunk_bitmaps(cx: &mut App, mut rng: StdRng) {
192        // Generate random buffer using existing test infrastructure
193        let len = rng.random_range(10..10000);
194        let buffer = if rng.random() {
195            let text = RandomCharIter::new(&mut rng).take(len).collect::<String>();
196            MultiBuffer::build_simple(&text, cx)
197        } else {
198            MultiBuffer::build_random(&mut rng, cx)
199        };
200
201        let buffer_snapshot = buffer.read(cx).snapshot(cx);
202
203        // Create random highlights
204        let mut highlights = sum_tree::TreeMap::default();
205        let highlight_count = rng.random_range(1..10);
206
207        for _i in 0..highlight_count {
208            let style = HighlightStyle {
209                color: Some(gpui::Hsla {
210                    h: rng.random::<f32>(),
211                    s: rng.random::<f32>(),
212                    l: rng.random::<f32>(),
213                    a: 1.0,
214                }),
215                ..Default::default()
216            };
217
218            let mut ranges = Vec::new();
219            let range_count = rng.random_range(1..10);
220            let text = buffer_snapshot.text();
221            for _ in 0..range_count {
222                if buffer_snapshot.len() == MultiBufferOffset(0) {
223                    continue;
224                }
225
226                let mut start = rng.random_range(
227                    MultiBufferOffset(0)..=buffer_snapshot.len().saturating_sub_usize(10),
228                );
229
230                while !text.is_char_boundary(start.0) {
231                    start = start.saturating_sub_usize(1);
232                }
233
234                let end_end = buffer_snapshot.len().min(start + 100usize);
235                let mut end = rng.random_range(start..=end_end);
236                while !text.is_char_boundary(end.0) {
237                    end = end.saturating_sub_usize(1);
238                }
239
240                if start < end {
241                    start = end;
242                }
243                let start_anchor = buffer_snapshot.anchor_before(start);
244                let end_anchor = buffer_snapshot.anchor_after(end);
245                ranges.push(start_anchor..end_anchor);
246            }
247
248            let type_id = TypeId::of::<()>(); // Simple type ID for testing
249            highlights.insert(HighlightKey::Type(type_id), Arc::new((style, ranges)));
250        }
251
252        // Get all chunks and verify their bitmaps
253        let chunks = CustomHighlightsChunks::new(
254            MultiBufferOffset(0)..buffer_snapshot.len(),
255            false,
256            None,
257            &buffer_snapshot,
258        );
259
260        for chunk in chunks {
261            let chunk_text = chunk.text;
262            let chars_bitmap = chunk.chars;
263            let tabs_bitmap = chunk.tabs;
264
265            // Check empty chunks have empty bitmaps
266            if chunk_text.is_empty() {
267                assert_eq!(
268                    chars_bitmap, 0,
269                    "Empty chunk should have empty chars bitmap"
270                );
271                assert_eq!(tabs_bitmap, 0, "Empty chunk should have empty tabs bitmap");
272                continue;
273            }
274
275            // Verify that chunk text doesn't exceed 128 bytes
276            assert!(
277                chunk_text.len() <= 128,
278                "Chunk text length {} exceeds 128 bytes",
279                chunk_text.len()
280            );
281
282            // Verify chars bitmap
283            let char_indices = chunk_text
284                .char_indices()
285                .map(|(i, _)| i)
286                .collect::<Vec<_>>();
287
288            for byte_idx in 0..chunk_text.len() {
289                let should_have_bit = char_indices.contains(&byte_idx);
290                let has_bit = chars_bitmap & (1 << byte_idx) != 0;
291
292                if has_bit != should_have_bit {
293                    eprintln!("Chunk text bytes: {:?}", chunk_text.as_bytes());
294                    eprintln!("Char indices: {:?}", char_indices);
295                    eprintln!("Chars bitmap: {:#b}", chars_bitmap);
296                    assert_eq!(
297                        has_bit, should_have_bit,
298                        "Chars bitmap mismatch at byte index {} in chunk {:?}. Expected bit: {}, Got bit: {}",
299                        byte_idx, chunk_text, should_have_bit, has_bit
300                    );
301                }
302            }
303
304            // Verify tabs bitmap
305            for (byte_idx, byte) in chunk_text.bytes().enumerate() {
306                let is_tab = byte == b'\t';
307                let has_bit = tabs_bitmap & (1 << byte_idx) != 0;
308
309                if has_bit != is_tab {
310                    eprintln!("Chunk text bytes: {:?}", chunk_text.as_bytes());
311                    eprintln!("Tabs bitmap: {:#b}", tabs_bitmap);
312                    assert_eq!(
313                        has_bit, is_tab,
314                        "Tabs bitmap mismatch at byte index {} in chunk {:?}. Byte: {:?}, Expected bit: {}, Got bit: {}",
315                        byte_idx, chunk_text, byte as char, is_tab, has_bit
316                    );
317                }
318            }
319        }
320    }
321}