1use collections::BTreeMap;
  2use gpui::HighlightStyle;
  3use language::Chunk;
  4use multi_buffer::{MultiBufferChunks, MultiBufferSnapshot, ToOffset as _};
  5use std::{
  6    any::TypeId,
  7    cmp,
  8    iter::{self, Peekable},
  9    ops::Range,
 10    vec,
 11};
 12
 13use crate::display_map::TextHighlights;
 14
 15pub struct CustomHighlightsChunks<'a> {
 16    buffer_chunks: MultiBufferChunks<'a>,
 17    buffer_chunk: Option<Chunk<'a>>,
 18    offset: usize,
 19    multibuffer_snapshot: &'a MultiBufferSnapshot,
 20
 21    highlight_endpoints: Peekable<vec::IntoIter<HighlightEndpoint>>,
 22    active_highlights: BTreeMap<(TypeId, usize), HighlightStyle>,
 23    text_highlights: Option<&'a TextHighlights>,
 24}
 25
 26#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 27struct HighlightEndpoint {
 28    offset: usize,
 29    is_start: bool,
 30    tag: (TypeId, usize),
 31    style: HighlightStyle,
 32}
 33
 34impl<'a> CustomHighlightsChunks<'a> {
 35    pub fn new(
 36        range: Range<usize>,
 37        language_aware: bool,
 38        text_highlights: Option<&'a TextHighlights>,
 39        multibuffer_snapshot: &'a MultiBufferSnapshot,
 40    ) -> Self {
 41        Self {
 42            buffer_chunks: multibuffer_snapshot.chunks(range.clone(), language_aware),
 43            buffer_chunk: None,
 44            offset: range.start,
 45
 46            text_highlights,
 47            highlight_endpoints: create_highlight_endpoints(
 48                &range,
 49                text_highlights,
 50                multibuffer_snapshot,
 51            ),
 52            active_highlights: Default::default(),
 53            multibuffer_snapshot,
 54        }
 55    }
 56
 57    pub fn seek(&mut self, new_range: Range<usize>) {
 58        self.highlight_endpoints =
 59            create_highlight_endpoints(&new_range, self.text_highlights, self.multibuffer_snapshot);
 60        self.offset = new_range.start;
 61        self.buffer_chunks.seek(new_range);
 62        self.buffer_chunk.take();
 63        self.active_highlights.clear()
 64    }
 65}
 66
 67fn create_highlight_endpoints(
 68    range: &Range<usize>,
 69    text_highlights: Option<&TextHighlights>,
 70    buffer: &MultiBufferSnapshot,
 71) -> iter::Peekable<vec::IntoIter<HighlightEndpoint>> {
 72    let mut highlight_endpoints = Vec::new();
 73    if let Some(text_highlights) = text_highlights {
 74        let start = buffer.anchor_after(range.start);
 75        let end = buffer.anchor_after(range.end);
 76        for (&tag, text_highlights) in text_highlights.iter() {
 77            let start_ix = match text_highlights.binary_search_by(|(probe, _)| {
 78                let cmp = probe.end.cmp(&start, &buffer);
 79                if cmp.is_gt() {
 80                    cmp::Ordering::Greater
 81                } else {
 82                    cmp::Ordering::Less
 83                }
 84            }) {
 85                Ok(i) | Err(i) => i,
 86            };
 87
 88            for (ix, (range, style)) in text_highlights[start_ix..].iter().enumerate() {
 89                if range.start.cmp(&end, &buffer).is_ge() {
 90                    break;
 91                }
 92
 93                highlight_endpoints.push(HighlightEndpoint {
 94                    offset: range.start.to_offset(&buffer),
 95                    is_start: true,
 96                    tag: (tag, ix),
 97                    style: *style,
 98                });
 99                highlight_endpoints.push(HighlightEndpoint {
100                    offset: range.end.to_offset(&buffer),
101                    is_start: false,
102                    tag: (tag, ix),
103                    style: *style,
104                });
105            }
106        }
107        highlight_endpoints.sort();
108    }
109    highlight_endpoints.into_iter().peekable()
110}
111
112impl<'a> Iterator for CustomHighlightsChunks<'a> {
113    type Item = Chunk<'a>;
114
115    fn next(&mut self) -> Option<Self::Item> {
116        let mut next_highlight_endpoint = usize::MAX;
117        while let Some(endpoint) = self.highlight_endpoints.peek().copied() {
118            if endpoint.offset <= self.offset {
119                if endpoint.is_start {
120                    self.active_highlights.insert(endpoint.tag, endpoint.style);
121                } else {
122                    self.active_highlights.remove(&endpoint.tag);
123                }
124                self.highlight_endpoints.next();
125            } else {
126                next_highlight_endpoint = endpoint.offset;
127                break;
128            }
129        }
130
131        let chunk = self
132            .buffer_chunk
133            .get_or_insert_with(|| self.buffer_chunks.next().unwrap());
134        if chunk.text.is_empty() {
135            *chunk = self.buffer_chunks.next().unwrap();
136        }
137
138        let (prefix, suffix) = chunk
139            .text
140            .split_at(chunk.text.len().min(next_highlight_endpoint - self.offset));
141
142        chunk.text = suffix;
143        self.offset += prefix.len();
144        let mut prefix = Chunk {
145            text: prefix,
146            ..chunk.clone()
147        };
148        if !self.active_highlights.is_empty() {
149            let mut highlight_style = HighlightStyle::default();
150            for active_highlight in self.active_highlights.values() {
151                highlight_style.highlight(*active_highlight);
152            }
153            prefix.highlight_style = Some(highlight_style);
154        }
155        Some(prefix)
156    }
157}
158
159impl PartialOrd for HighlightEndpoint {
160    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
161        Some(self.cmp(other))
162    }
163}
164
165impl Ord for HighlightEndpoint {
166    fn cmp(&self, other: &Self) -> cmp::Ordering {
167        self.offset
168            .cmp(&other.offset)
169            .then_with(|| other.is_start.cmp(&self.is_start))
170    }
171}