1use collections::BTreeMap;
  2use gpui::HighlightStyle;
  3use language::Chunk;
  4use multi_buffer::{Anchor, MultiBufferChunks, MultiBufferSnapshot, ToOffset as _};
  5use std::{
  6    any::TypeId,
  7    cmp,
  8    iter::{self, Peekable},
  9    ops::Range,
 10    sync::Arc,
 11    vec,
 12};
 13use sum_tree::TreeMap;
 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, HighlightStyle>,
 23    text_highlights: Option<&'a TreeMap<TypeId, Arc<(HighlightStyle, Vec<Range<Anchor>>)>>>,
 24}
 25
 26#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 27struct HighlightEndpoint {
 28    offset: usize,
 29    is_start: bool,
 30    tag: TypeId,
 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 TreeMap<TypeId, Arc<(HighlightStyle, Vec<Range<Anchor>>)>>>,
 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<&TreeMap<TypeId, Arc<(HighlightStyle, Vec<Range<Anchor>>)>>>,
 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 style = text_highlights.0;
 78            let ranges = &text_highlights.1;
 79
 80            let start_ix = match ranges.binary_search_by(|probe| {
 81                let cmp = probe.end.cmp(&start, &buffer);
 82                if cmp.is_gt() {
 83                    cmp::Ordering::Greater
 84                } else {
 85                    cmp::Ordering::Less
 86                }
 87            }) {
 88                Ok(i) | Err(i) => i,
 89            };
 90
 91            for range in &ranges[start_ix..] {
 92                if range.start.cmp(&end, &buffer).is_ge() {
 93                    break;
 94                }
 95
 96                highlight_endpoints.push(HighlightEndpoint {
 97                    offset: range.start.to_offset(&buffer),
 98                    is_start: true,
 99                    tag,
100                    style,
101                });
102                highlight_endpoints.push(HighlightEndpoint {
103                    offset: range.end.to_offset(&buffer),
104                    is_start: false,
105                    tag,
106                    style,
107                });
108            }
109        }
110        highlight_endpoints.sort();
111    }
112    highlight_endpoints.into_iter().peekable()
113}
114
115impl<'a> Iterator for CustomHighlightsChunks<'a> {
116    type Item = Chunk<'a>;
117
118    fn next(&mut self) -> Option<Self::Item> {
119        let mut next_highlight_endpoint = usize::MAX;
120        while let Some(endpoint) = self.highlight_endpoints.peek().copied() {
121            if endpoint.offset <= self.offset {
122                if endpoint.is_start {
123                    self.active_highlights.insert(endpoint.tag, endpoint.style);
124                } else {
125                    self.active_highlights.remove(&endpoint.tag);
126                }
127                self.highlight_endpoints.next();
128            } else {
129                next_highlight_endpoint = endpoint.offset;
130                break;
131            }
132        }
133
134        let chunk = self
135            .buffer_chunk
136            .get_or_insert_with(|| self.buffer_chunks.next().unwrap());
137        if chunk.text.is_empty() {
138            *chunk = self.buffer_chunks.next().unwrap();
139        }
140
141        let (prefix, suffix) = chunk
142            .text
143            .split_at(chunk.text.len().min(next_highlight_endpoint - self.offset));
144
145        chunk.text = suffix;
146        self.offset += prefix.len();
147        let mut prefix = Chunk {
148            text: prefix,
149            ..chunk.clone()
150        };
151        if !self.active_highlights.is_empty() {
152            let mut highlight_style = HighlightStyle::default();
153            for active_highlight in self.active_highlights.values() {
154                highlight_style.highlight(*active_highlight);
155            }
156            prefix.highlight_style = Some(highlight_style);
157        }
158        Some(prefix)
159    }
160}
161
162impl PartialOrd for HighlightEndpoint {
163    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
164        Some(self.cmp(other))
165    }
166}
167
168impl Ord for HighlightEndpoint {
169    fn cmp(&self, other: &Self) -> cmp::Ordering {
170        self.offset
171            .cmp(&other.offset)
172            .then_with(|| other.is_start.cmp(&self.is_start))
173    }
174}