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