tenant.openagents/omega

1use collections::BTreeMap;
2use gpui::HighlightStyle;
3use language::{Chunk, LanguageAwareStyling};
4use multi_buffer::{MultiBufferChunks, MultiBufferOffset, MultiBufferSnapshot};
5use std::{cmp, ops::Range};
6
7use crate::display_map::{HighlightKey, SemanticTokensHighlights, TextHighlights};
8
9pub struct CustomHighlightsChunks<'a> {
10    buffer_chunks: MultiBufferChunks<'a>,
11    buffer_chunk: Option<Chunk<'a>>,
12    offset: MultiBufferOffset,
13    multibuffer_snapshot: &'a MultiBufferSnapshot,
14
15    highlight_endpoints: Vec<HighlightEndpoint>,
16    active_highlights: BTreeMap<HighlightKey, HighlightStyle>,
17    text_highlights: Option<&'a TextHighlights>,
18    semantic_token_highlights: Option<&'a SemanticTokensHighlights>,
19}
20
21#[derive(Debug, Copy, Clone, Eq, PartialEq)]
22struct HighlightEndpoint {
23    offset: MultiBufferOffset,
24    tag: HighlightKey,
25    style: Option<HighlightStyle>,
26}
27
28impl<'a> CustomHighlightsChunks<'a> {
29    #[ztracing::instrument(skip_all)]
30    pub fn new(
31        range: Range<MultiBufferOffset>,
32        language_aware: LanguageAwareStyling,
33        text_highlights: Option<&'a TextHighlights>,
34        semantic_token_highlights: Option<&'a SemanticTokensHighlights>,
35        multibuffer_snapshot: &'a MultiBufferSnapshot,
36    ) -> Self {
37        let mut highlight_endpoints = Vec::new();
38        create_highlight_endpoints(
39            &range,
40            text_highlights,
41            semantic_token_highlights,
42            multibuffer_snapshot,
43            &mut highlight_endpoints,
44        );
45        Self {
46            buffer_chunks: multibuffer_snapshot.chunks(range.clone(), language_aware),
47            buffer_chunk: None,
48            offset: range.start,
49            text_highlights,
50            highlight_endpoints,
51            active_highlights: Default::default(),
52            multibuffer_snapshot,
53            semantic_token_highlights,
54        }
55    }
56
57    #[ztracing::instrument(skip_all)]
58    pub fn seek(&mut self, new_range: Range<MultiBufferOffset>) {
59        create_highlight_endpoints(
60            &new_range,
61            self.text_highlights,
62            self.semantic_token_highlights,
63            self.multibuffer_snapshot,
64            &mut self.highlight_endpoints,
65        );
66        self.offset = new_range.start;
67        self.buffer_chunks.seek(new_range);
68        self.buffer_chunk.take();
69        self.active_highlights.clear()
70    }
71}
72
73fn create_highlight_endpoints(
74    range: &Range<MultiBufferOffset>,
75    text_highlights: Option<&TextHighlights>,
76    semantic_token_highlights: Option<&SemanticTokensHighlights>,
77    buffer: &MultiBufferSnapshot,
78    highlight_endpoints: &mut Vec<HighlightEndpoint>,
79) {
80    highlight_endpoints.clear();
81    if let Some(text_highlights) = text_highlights {
82        let start = buffer.anchor_after(range.start);
83        let end = buffer.anchor_after(range.end);
84        let mut text_highlights_scratch = Vec::new();
85
86        for (&tag, text_highlights) in text_highlights.iter() {
87            let style = text_highlights.0;
88            let ranges = &text_highlights.1;
89
90            let start_ix = ranges
91                .binary_search_by(|probe| probe.end.cmp(&start, buffer).then(cmp::Ordering::Less))
92                .unwrap_or_else(|i| i);
93            let end_ix = ranges[start_ix..]
94                .binary_search_by(|probe| {
95                    probe.start.cmp(&end, buffer).then(cmp::Ordering::Greater)
96                })
97                .unwrap_or_else(|i| i);
98
99            let ranges_ = &ranges[start_ix..][..end_ix];
100            text_highlights_scratch.clear();
101            text_highlights_scratch.reserve(ranges_.len());
102            highlight_endpoints.reserve(2 * ranges_.len());
103
104            let mut iter = ranges_.iter();
105            buffer.summaries_for_anchors_cb(
106                ranges_.iter().map(|r| &r.start),
107                |start: MultiBufferOffset| {
108                    text_highlights_scratch.push((start, iter.next().unwrap().end));
109                },
110            );
111            text_highlights_scratch.sort_by(|a, b| a.1.cmp(&b.1, buffer));
112            let mut iter = text_highlights_scratch.iter();
113            buffer.summaries_for_anchors_cb(
114                text_highlights_scratch.iter().map(|(_, end)| end),
115                |end: MultiBufferOffset| {
116                    let start = iter.next().unwrap().0;
117                    if start == end {
118                        return;
119                    }
120                    highlight_endpoints.push(HighlightEndpoint {
121                        offset: start,
122                        tag,
123                        style: Some(style),
124                    });
125                    highlight_endpoints.push(HighlightEndpoint {
126                        offset: end,
127                        tag,
128                        style: None,
129                    });
130                },
131            );
132        }
133    }
134    if let Some(semantic_token_highlights) = semantic_token_highlights {
135        let start = buffer.anchor_after(range.start);
136        let end = buffer.anchor_after(range.end);
137        let mut semantic_highlights_scratch = Vec::new();
138        for buffer_id in buffer.buffer_ids_for_range(range.clone()) {
139            let Some((semantic_token_highlights, interner)) =
140                semantic_token_highlights.get(&buffer_id)
141            else {
142                continue;
143            };
144            let start_ix = semantic_token_highlights
145                .binary_search_by(|probe| {
146                    probe
147                        .range
148                        .end
149                        .cmp(&start, buffer)
150                        .then(cmp::Ordering::Less)
151                })
152                .unwrap_or_else(|i| i);
153            let end_ix = semantic_token_highlights[start_ix..]
154                .binary_search_by(|probe| {
155                    probe
156                        .range
157                        .start
158                        .cmp(&end, buffer)
159                        .then(cmp::Ordering::Greater)
160                })
161                .unwrap_or_else(|i| i);
162
163            let ranges_ = &semantic_token_highlights[start_ix..][..end_ix];
164            semantic_highlights_scratch.clear();
165            semantic_highlights_scratch.reserve(ranges_.len());
166            highlight_endpoints.reserve(2 * ranges_.len());
167
168            let mut iter = ranges_.iter();
169            buffer.summaries_for_anchors_cb(
170                ranges_.iter().map(|token| &token.range.start),
171                |start: MultiBufferOffset| {
172                    semantic_highlights_scratch.push((start, iter.next().unwrap()));
173                },
174            );
175            semantic_highlights_scratch.sort_by(|a, b| a.1.range.end.cmp(&b.1.range.end, buffer));
176            let mut iter = semantic_highlights_scratch.iter();
177            buffer.summaries_for_anchors_cb(
178                semantic_highlights_scratch
179                    .iter()
180                    .map(|(_, token)| &token.range.end),
181                |end: MultiBufferOffset| {
182                    let (start, token) = iter.next().unwrap();
183                    if *start == end {
184                        return;
185                    }
186                    highlight_endpoints.push(HighlightEndpoint {
187                        offset: *start,
188                        tag: HighlightKey::SemanticToken,
189                        style: Some(interner[token.style]),
190                    });
191                    highlight_endpoints.push(HighlightEndpoint {
192                        offset: end,
193                        tag: HighlightKey::SemanticToken,
194                        style: None,
195                    });
196                },
197            );
198        }
199    }
200    highlight_endpoints.sort_by(|a, b| a.cmp(b).reverse());
201}
202
203impl<'a> Iterator for CustomHighlightsChunks<'a> {
204    type Item = Chunk<'a>;
205
206    #[ztracing::instrument(skip_all)]
207    fn next(&mut self) -> Option<Self::Item> {
208        let mut next_highlight_endpoint = MultiBufferOffset(usize::MAX);
209        while let Some(endpoint) = self.highlight_endpoints.last().copied() {
210            if endpoint.offset <= self.offset {
211                if let Some(style) = endpoint.style {
212                    self.active_highlights.insert(endpoint.tag, style);
213                } else {
214                    self.active_highlights.remove(&endpoint.tag);
215                }
216                self.highlight_endpoints.pop();
217            } else {
218                next_highlight_endpoint = endpoint.offset;
219                break;
220            }
221        }
222
223        let chunk = match &mut self.buffer_chunk {
224            Some(it) => it,
225            slot => slot.insert(self.buffer_chunks.next()?),
226        };
227        while chunk.text.is_empty() {
228            *chunk = self.buffer_chunks.next()?;
229        }
230
231        let split_idx = chunk.text.len().min(next_highlight_endpoint - self.offset);
232        let (prefix, suffix) = chunk.text.split_at(split_idx);
233        self.offset += prefix.len();
234
235        let mask = 1u128.unbounded_shl(split_idx as u32).wrapping_sub(1);
236        let chars = chunk.chars & mask;
237        let tabs = chunk.tabs & mask;
238        let newlines = chunk.newlines & mask;
239        let mut prefix = Chunk {
240            text: prefix,
241            chars,
242            tabs,
243            newlines,
244            ..chunk.clone()
245        };
246
247        chunk.chars = chunk.chars.unbounded_shr(split_idx as u32);
248        chunk.tabs = chunk.tabs.unbounded_shr(split_idx as u32);
249        chunk.newlines = chunk.newlines.unbounded_shr(split_idx as u32);
250        chunk.text = suffix;
251        if !self.active_highlights.is_empty() {
252            prefix.highlight_style = self
253                .active_highlights
254                .values()
255                .copied()
256                .reduce(|acc, active_highlight| acc.highlight(active_highlight));
257        }
258        Some(prefix)
259    }
260}
261
262impl PartialOrd for HighlightEndpoint {
263    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
264        Some(self.cmp(other))
265    }
266}
267
268impl Ord for HighlightEndpoint {
269    fn cmp(&self, other: &Self) -> cmp::Ordering {
270        self.offset
271            .cmp(&other.offset)
272            .then_with(|| self.style.is_some().cmp(&other.style.is_some()))
273            .then_with(|| self.tag.cmp(&other.tag))
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use std::sync::Arc;
280
281    use super::*;
282    use crate::MultiBuffer;
283    use gpui::App;
284    use rand::prelude::*;
285    use util::RandomCharIter;
286
287    #[gpui::test(iterations = 100)]
288    fn test_random_chunk_bitmaps(cx: &mut App, mut rng: StdRng) {
289        // Generate random buffer using existing test infrastructure
290        let len = rng.random_range(10..10000);
291        let buffer = if rng.random() {
292            let text = RandomCharIter::new(&mut rng).take(len).collect::<String>();
293            MultiBuffer::build_simple(&text, cx)
294        } else {
295            MultiBuffer::build_random(&mut rng, cx)
296        };
297
298        let buffer_snapshot = buffer.read(cx).snapshot(cx);
299
300        // Create random highlights
301        let mut highlights = sum_tree::TreeMap::default();
302        let highlight_count = rng.random_range(1..10);
303
304        for _i in 0..highlight_count {
305            let style = HighlightStyle {
306                color: Some(gpui::Hsla {
307                    h: rng.random::<f32>(),
308                    s: rng.random::<f32>(),
309                    l: rng.random::<f32>(),
310                    a: 1.0,
311                }),
312                ..Default::default()
313            };
314
315            let mut ranges = Vec::new();
316            let range_count = rng.random_range(1..10);
317            let text = buffer_snapshot.text();
318            for _ in 0..range_count {
319                if buffer_snapshot.len() == MultiBufferOffset(0) {
320                    continue;
321                }
322
323                let mut start = rng.random_range(
324                    MultiBufferOffset(0)..=buffer_snapshot.len().saturating_sub_usize(10),
325                );
326
327                while !text.is_char_boundary(start.0) {
328                    start = start.saturating_sub_usize(1);
329                }
330
331                let end_end = buffer_snapshot.len().min(start + 100usize);
332                let mut end = rng.random_range(start..=end_end);
333                while !text.is_char_boundary(end.0) {
334                    end = end.saturating_sub_usize(1);
335                }
336
337                if start < end {
338                    start = end;
339                }
340                let start_anchor = buffer_snapshot.anchor_before(start);
341                let end_anchor = buffer_snapshot.anchor_after(end);
342                ranges.push(start_anchor..end_anchor);
343            }
344
345            highlights.insert(HighlightKey::Editor, Arc::new((style, ranges)));
346        }
347
348        // Get all chunks and verify their bitmaps
349        let chunks = CustomHighlightsChunks::new(
350            MultiBufferOffset(0)..buffer_snapshot.len(),
351            LanguageAwareStyling {
352                tree_sitter: false,
353                diagnostics: false,
354            },
355            None,
356            None,
357            &buffer_snapshot,
358        );
359
360        for chunk in chunks {
361            let chunk_text = chunk.text;
362            let chars_bitmap = chunk.chars;
363            let tabs_bitmap = chunk.tabs;
364
365            // Check empty chunks have empty bitmaps
366            if chunk_text.is_empty() {
367                assert_eq!(
368                    chars_bitmap, 0,
369                    "Empty chunk should have empty chars bitmap"
370                );
371                assert_eq!(tabs_bitmap, 0, "Empty chunk should have empty tabs bitmap");
372                continue;
373            }
374
375            // Verify that chunk text doesn't exceed 128 bytes
376            assert!(
377                chunk_text.len() <= 128,
378                "Chunk text length {} exceeds 128 bytes",
379                chunk_text.len()
380            );
381
382            // Verify chars bitmap
383            let char_indices = chunk_text
384                .char_indices()
385                .map(|(i, _)| i)
386                .collect::<Vec<_>>();
387
388            for byte_idx in 0..chunk_text.len() {
389                let should_have_bit = char_indices.contains(&byte_idx);
390                let has_bit = chars_bitmap & (1 << byte_idx) != 0;
391
392                if has_bit != should_have_bit {
393                    eprintln!("Chunk text bytes: {:?}", chunk_text.as_bytes());
394                    eprintln!("Char indices: {:?}", char_indices);
395                    eprintln!("Chars bitmap: {:#b}", chars_bitmap);
396                    assert_eq!(
397                        has_bit, should_have_bit,
398                        "Chars bitmap mismatch at byte index {} in chunk {:?}. Expected bit: {}, Got bit: {}",
399                        byte_idx, chunk_text, should_have_bit, has_bit
400                    );
401                }
402            }
403
404            // Verify tabs bitmap
405            for (byte_idx, byte) in chunk_text.bytes().enumerate() {
406                let is_tab = byte == b'\t';
407                let has_bit = tabs_bitmap & (1 << byte_idx) != 0;
408
409                if has_bit != is_tab {
410                    eprintln!("Chunk text bytes: {:?}", chunk_text.as_bytes());
411                    eprintln!("Tabs bitmap: {:#b}", tabs_bitmap);
412                    assert_eq!(
413                        has_bit, is_tab,
414                        "Tabs bitmap mismatch at byte index {} in chunk {:?}. Byte: {:?}, Expected bit: {}, Got bit: {}",
415                        byte_idx, chunk_text, byte as char, is_tab, has_bit
416                    );
417                }
418            }
419        }
420    }
421}
422