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