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