inlay_map.rs

  1#![allow(unused)]
  2// TODO kb
  3
  4use std::{
  5    cmp::{self, Reverse},
  6    ops::{Add, AddAssign, Range, Sub},
  7    sync::atomic::{self, AtomicUsize},
  8};
  9
 10use crate::{Anchor, ExcerptId, InlayHintLocation, MultiBufferSnapshot, ToOffset, ToPoint};
 11
 12use super::{
 13    suggestion_map::{
 14        SuggestionBufferRows, SuggestionChunks, SuggestionEdit, SuggestionOffset, SuggestionPoint,
 15        SuggestionSnapshot,
 16    },
 17    TextHighlights,
 18};
 19use collections::{BTreeMap, HashMap, HashSet};
 20use gpui::fonts::HighlightStyle;
 21use language::{Chunk, Edit, Point, Rope, TextSummary};
 22use parking_lot::Mutex;
 23use project::InlayHint;
 24use rand::Rng;
 25use sum_tree::{Bias, Cursor, SumTree};
 26use util::post_inc;
 27
 28#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
 29pub struct InlayId(usize);
 30
 31pub struct InlayMap {
 32    snapshot: Mutex<InlaySnapshot>,
 33    next_inlay_id: usize,
 34    pub(super) inlays: HashMap<InlayId, (InlayHintLocation, Inlay)>,
 35}
 36
 37#[derive(Clone)]
 38pub struct InlaySnapshot {
 39    // TODO kb merge these two together?
 40    pub suggestion_snapshot: SuggestionSnapshot,
 41    transforms: SumTree<Transform>,
 42    pub version: usize,
 43}
 44
 45#[derive(Clone, Debug)]
 46enum Transform {
 47    Isomorphic(TextSummary),
 48    Inlay(Inlay),
 49}
 50
 51impl sum_tree::Item for Transform {
 52    type Summary = TransformSummary;
 53
 54    fn summary(&self) -> Self::Summary {
 55        match self {
 56            Transform::Isomorphic(summary) => TransformSummary {
 57                input: summary.clone(),
 58                output: summary.clone(),
 59            },
 60            Transform::Inlay(inlay) => TransformSummary {
 61                input: TextSummary::default(),
 62                output: inlay.properties.text.summary(),
 63            },
 64        }
 65    }
 66}
 67
 68#[derive(Clone, Debug, Default)]
 69struct TransformSummary {
 70    input: TextSummary,
 71    output: TextSummary,
 72}
 73
 74impl sum_tree::Summary for TransformSummary {
 75    type Context = ();
 76
 77    fn add_summary(&mut self, other: &Self, _: &()) {
 78        self.input += &other.input;
 79        self.output += &other.output;
 80    }
 81}
 82
 83pub type InlayEdit = Edit<InlayOffset>;
 84
 85#[derive(Copy, Clone, Debug, Default, Eq, Ord, PartialOrd, PartialEq)]
 86pub struct InlayOffset(pub usize);
 87
 88impl Add for InlayOffset {
 89    type Output = Self;
 90
 91    fn add(self, rhs: Self) -> Self::Output {
 92        Self(self.0 + rhs.0)
 93    }
 94}
 95
 96impl Sub for InlayOffset {
 97    type Output = Self;
 98
 99    fn sub(self, rhs: Self) -> Self::Output {
100        Self(self.0 - rhs.0)
101    }
102}
103
104impl AddAssign for InlayOffset {
105    fn add_assign(&mut self, rhs: Self) {
106        self.0 += rhs.0;
107    }
108}
109
110impl<'a> sum_tree::Dimension<'a, TransformSummary> for InlayOffset {
111    fn add_summary(&mut self, summary: &'a TransformSummary, _: &()) {
112        self.0 += &summary.output.len;
113    }
114}
115
116impl<'a> sum_tree::Dimension<'a, TransformSummary> for SuggestionOffset {
117    fn add_summary(&mut self, summary: &'a TransformSummary, _: &()) {
118        self.0 += &summary.input.len;
119    }
120}
121
122#[derive(Copy, Clone, Debug, Default, Eq, Ord, PartialOrd, PartialEq)]
123pub struct InlayPoint(pub Point);
124
125impl<'a> sum_tree::Dimension<'a, TransformSummary> for InlayPoint {
126    fn add_summary(&mut self, summary: &'a TransformSummary, _: &()) {
127        self.0 += &summary.output.lines;
128    }
129}
130
131impl<'a> sum_tree::Dimension<'a, TransformSummary> for SuggestionPoint {
132    fn add_summary(&mut self, summary: &'a TransformSummary, _: &()) {
133        self.0 += &summary.input.lines;
134    }
135}
136
137#[derive(Clone)]
138pub struct InlayBufferRows<'a> {
139    suggestion_rows: SuggestionBufferRows<'a>,
140}
141
142pub struct InlayChunks<'a> {
143    transforms: Cursor<'a, Transform, (InlayOffset, SuggestionOffset)>,
144    suggestion_chunks: SuggestionChunks<'a>,
145    suggestion_chunk: Option<Chunk<'a>>,
146    inlay_chunks: Option<text::Chunks<'a>>,
147    output_offset: InlayOffset,
148    max_output_offset: InlayOffset,
149}
150
151#[derive(Debug, Clone)]
152pub struct Inlay {
153    pub(super) id: InlayId,
154    pub(super) properties: InlayProperties,
155}
156
157#[derive(Debug, Clone)]
158pub struct InlayProperties {
159    pub(super) position: Anchor,
160    pub(super) text: Rope,
161}
162
163impl<'a> Iterator for InlayChunks<'a> {
164    type Item = Chunk<'a>;
165
166    fn next(&mut self) -> Option<Self::Item> {
167        if self.output_offset == self.max_output_offset {
168            return None;
169        }
170
171        let chunk = match self.transforms.item()? {
172            Transform::Isomorphic(transform) => {
173                let chunk = self
174                    .suggestion_chunk
175                    .get_or_insert_with(|| self.suggestion_chunks.next().unwrap());
176                if chunk.text.is_empty() {
177                    *chunk = self.suggestion_chunks.next().unwrap();
178                }
179
180                let (prefix, suffix) = chunk
181                    .text
182                    .split_at(cmp::min(transform.len, chunk.text.len()));
183                chunk.text = suffix;
184                self.output_offset.0 += prefix.len();
185                Chunk {
186                    text: prefix,
187                    ..chunk.clone()
188                }
189            }
190            Transform::Inlay(inlay) => {
191                let inlay_chunks = self.inlay_chunks.get_or_insert_with(|| {
192                    let start = self.output_offset - self.transforms.start().0;
193                    let end = cmp::min(self.max_output_offset, self.transforms.end(&()).0)
194                        - self.transforms.start().0;
195                    inlay.properties.text.chunks_in_range(start.0..end.0)
196                });
197
198                let chunk = inlay_chunks.next().unwrap();
199                self.output_offset.0 += chunk.len();
200                Chunk {
201                    text: chunk,
202                    ..Default::default()
203                }
204            }
205        };
206
207        if self.output_offset == self.transforms.end(&()).0 {
208            self.transforms.next(&());
209        }
210
211        Some(chunk)
212    }
213}
214
215impl<'a> Iterator for InlayBufferRows<'a> {
216    type Item = Option<u32>;
217
218    fn next(&mut self) -> Option<Self::Item> {
219        self.suggestion_rows.next()
220    }
221}
222
223impl InlayPoint {
224    pub fn new(row: u32, column: u32) -> Self {
225        Self(Point::new(row, column))
226    }
227
228    pub fn row(self) -> u32 {
229        self.0.row
230    }
231
232    pub fn column(self) -> u32 {
233        self.0.column
234    }
235}
236
237impl InlayMap {
238    pub fn new(suggestion_snapshot: SuggestionSnapshot) -> (Self, InlaySnapshot) {
239        let snapshot = InlaySnapshot {
240            suggestion_snapshot: suggestion_snapshot.clone(),
241            version: 0,
242            transforms: SumTree::from_item(
243                Transform::Isomorphic(suggestion_snapshot.text_summary()),
244                &(),
245            ),
246        };
247
248        (
249            Self {
250                snapshot: Mutex::new(snapshot.clone()),
251                next_inlay_id: 0,
252                inlays: HashMap::default(),
253            },
254            snapshot,
255        )
256    }
257
258    pub fn sync(
259        &self,
260        suggestion_snapshot: SuggestionSnapshot,
261        suggestion_edits: Vec<SuggestionEdit>,
262    ) -> (InlaySnapshot, Vec<InlayEdit>) {
263        let mut snapshot = self.snapshot.lock();
264
265        if snapshot.suggestion_snapshot.version != suggestion_snapshot.version {
266            snapshot.version += 1;
267        }
268
269        let mut new_transforms = SumTree::new();
270        let mut cursor = snapshot.transforms.cursor::<SuggestionOffset>();
271        let mut suggestion_edits = suggestion_edits.iter().peekable();
272
273        while let Some(suggestion_edit) = suggestion_edits.next() {
274            if suggestion_edit.old.start >= *cursor.start() {
275                new_transforms.push_tree(
276                    cursor.slice(&suggestion_edit.old.start, Bias::Right, &()),
277                    &(),
278                );
279            }
280
281            if suggestion_edit.old.end > cursor.end(&()) {
282                cursor.seek_forward(&suggestion_edit.old.end, Bias::Right, &());
283            }
284
285            let transform_start = SuggestionOffset(new_transforms.summary().input.len);
286            let mut transform_end = suggestion_edit.new.end;
287            if suggestion_edits
288                .peek()
289                .map_or(true, |edit| edit.old.start > cursor.end(&()))
290            {
291                transform_end += cursor.end(&()) - suggestion_edit.old.end;
292                cursor.next(&());
293            }
294            new_transforms.push(
295                Transform::Isomorphic(suggestion_snapshot.text_summary_for_range(
296                    suggestion_snapshot.to_point(transform_start)
297                        ..suggestion_snapshot.to_point(transform_end),
298                )),
299                &(),
300            );
301        }
302
303        new_transforms.push_tree(cursor.suffix(&()), &());
304        drop(cursor);
305
306        snapshot.suggestion_snapshot = suggestion_snapshot;
307        dbg!(new_transforms.items(&()));
308        snapshot.transforms = new_transforms;
309
310        (snapshot.clone(), Default::default())
311    }
312
313    pub fn splice(
314        &mut self,
315        to_remove: HashSet<InlayId>,
316        to_insert: Vec<(InlayHintLocation, InlayProperties)>,
317    ) -> (InlaySnapshot, Vec<InlayEdit>, Vec<InlayId>) {
318        let mut snapshot = self.snapshot.lock();
319
320        let mut inlays = BTreeMap::new();
321        let mut new_ids = Vec::new();
322        for (location, properties) in to_insert {
323            let inlay = Inlay {
324                id: InlayId(post_inc(&mut self.next_inlay_id)),
325                properties,
326            };
327            self.inlays.insert(inlay.id, (location, inlay.clone()));
328            new_ids.push(inlay.id);
329
330            let buffer_point = inlay
331                .properties
332                .position
333                .to_point(snapshot.buffer_snapshot());
334            let fold_point = snapshot
335                .suggestion_snapshot
336                .fold_snapshot
337                .to_fold_point(buffer_point, Bias::Left);
338            let suggestion_point = snapshot.suggestion_snapshot.to_suggestion_point(fold_point);
339            let inlay_point = snapshot.to_inlay_point(suggestion_point);
340
341            inlays.insert((inlay_point, Reverse(inlay.id)), Some(inlay));
342        }
343
344        for inlay_id in to_remove {
345            if let Some((_, inlay)) = self.inlays.remove(&inlay_id) {
346                let buffer_point = inlay
347                    .properties
348                    .position
349                    .to_point(snapshot.buffer_snapshot());
350                let fold_point = snapshot
351                    .suggestion_snapshot
352                    .fold_snapshot
353                    .to_fold_point(buffer_point, Bias::Left);
354                let suggestion_point = snapshot.suggestion_snapshot.to_suggestion_point(fold_point);
355                let inlay_point = snapshot.to_inlay_point(suggestion_point);
356                inlays.insert((inlay_point, Reverse(inlay.id)), None);
357            }
358        }
359
360        let mut new_transforms = SumTree::new();
361        let mut cursor = snapshot
362            .transforms
363            .cursor::<(InlayPoint, SuggestionPoint)>();
364        for ((inlay_point, inlay_id), inlay) in inlays {
365            new_transforms.push_tree(cursor.slice(&inlay_point, Bias::Right, &()), &());
366            while let Some(transform) = cursor.item() {
367                match transform {
368                    Transform::Isomorphic(_) => break,
369                    Transform::Inlay(inlay) => {
370                        if inlay.id > inlay_id.0 {
371                            new_transforms.push(transform.clone(), &());
372                            cursor.next(&());
373                        } else {
374                            if inlay.id == inlay_id.0 {
375                                cursor.next(&());
376                            }
377                            break;
378                        }
379                    }
380                }
381            }
382
383            if let Some(inlay) = inlay {
384                if let Some(Transform::Isomorphic(transform)) = cursor.item() {
385                    let prefix = inlay_point.0 - cursor.start().0 .0;
386                    if !prefix.is_zero() {
387                        let prefix_suggestion_start = cursor.start().1;
388                        let prefix_suggestion_end = SuggestionPoint(cursor.start().1 .0 + prefix);
389                        new_transforms.push(
390                            Transform::Isomorphic(
391                                snapshot.suggestion_snapshot.text_summary_for_range(
392                                    prefix_suggestion_start..prefix_suggestion_end,
393                                ),
394                            ),
395                            &(),
396                        );
397                    }
398
399                    new_transforms.push(Transform::Inlay(inlay), &());
400
401                    let suffix_suggestion_start = SuggestionPoint(cursor.start().1 .0 + prefix);
402                    let suffix_suggestion_end = cursor.end(&()).1;
403                    new_transforms.push(
404                        Transform::Isomorphic(snapshot.suggestion_snapshot.text_summary_for_range(
405                            suffix_suggestion_start..suffix_suggestion_end,
406                        )),
407                        &(),
408                    );
409
410                    cursor.next(&());
411                } else {
412                    new_transforms.push(Transform::Inlay(inlay), &());
413                }
414            }
415        }
416
417        new_transforms.push_tree(cursor.suffix(&()), &());
418        drop(cursor);
419        snapshot.transforms = new_transforms;
420        snapshot.version += 1;
421
422        (snapshot.clone(), Vec::new(), new_ids)
423    }
424}
425
426impl InlaySnapshot {
427    pub fn buffer_snapshot(&self) -> &MultiBufferSnapshot {
428        // TODO kb copied from suggestion_map
429        self.suggestion_snapshot.buffer_snapshot()
430    }
431
432    pub fn to_point(&self, offset: InlayOffset) -> InlayPoint {
433        // TODO kb copied from suggestion_map
434        self.to_inlay_point(
435            self.suggestion_snapshot
436                .to_point(super::suggestion_map::SuggestionOffset(offset.0)),
437        )
438    }
439
440    pub fn len(&self) -> InlayOffset {
441        InlayOffset(self.transforms.summary().output.len)
442    }
443
444    pub fn max_point(&self) -> InlayPoint {
445        InlayPoint(self.transforms.summary().output.lines)
446    }
447
448    pub fn to_offset(&self, point: InlayPoint) -> InlayOffset {
449        // TODO kb copied from suggestion_map
450        InlayOffset(
451            self.suggestion_snapshot
452                .to_offset(self.to_suggestion_point(point, Bias::Left))
453                .0,
454        )
455    }
456
457    pub fn chars_at(&self, start: InlayPoint) -> impl '_ + Iterator<Item = char> {
458        self.suggestion_snapshot
459            .chars_at(self.to_suggestion_point(start, Bias::Left))
460    }
461
462    // TODO kb what to do with bias?
463    pub fn to_suggestion_point(&self, point: InlayPoint, _: Bias) -> SuggestionPoint {
464        SuggestionPoint(point.0)
465    }
466
467    pub fn to_suggestion_offset(&self, offset: InlayOffset) -> SuggestionOffset {
468        let mut cursor = self.transforms.cursor::<(InlayOffset, SuggestionOffset)>();
469        cursor.seek(&offset, Bias::Right, &());
470        match cursor.item() {
471            Some(Transform::Isomorphic(transform)) => {
472                let overshoot = offset - cursor.start().0;
473                cursor.start().1 + SuggestionOffset(overshoot.0)
474            }
475            Some(Transform::Inlay(inlay)) => cursor.start().1,
476            None => self.suggestion_snapshot.len(),
477        }
478    }
479
480    pub fn to_inlay_point(&self, point: SuggestionPoint) -> InlayPoint {
481        InlayPoint(point.0)
482    }
483
484    pub fn clip_point(&self, point: InlayPoint, bias: Bias) -> InlayPoint {
485        // TODO kb copied from suggestion_map
486        self.to_inlay_point(
487            self.suggestion_snapshot
488                .clip_point(self.to_suggestion_point(point, bias), bias),
489        )
490    }
491
492    pub fn text_summary_for_range(&self, range: Range<InlayPoint>) -> TextSummary {
493        // TODO kb copied from suggestion_map
494        self.suggestion_snapshot.text_summary_for_range(
495            self.to_suggestion_point(range.start, Bias::Left)
496                ..self.to_suggestion_point(range.end, Bias::Left),
497        )
498    }
499
500    pub fn buffer_rows<'a>(&'a self, row: u32) -> InlayBufferRows<'a> {
501        InlayBufferRows {
502            suggestion_rows: self.suggestion_snapshot.buffer_rows(row),
503        }
504    }
505
506    pub fn line_len(&self, row: u32) -> u32 {
507        // TODO kb copied from suggestion_map
508        self.suggestion_snapshot.line_len(row)
509    }
510
511    pub fn chunks<'a>(
512        &'a self,
513        range: Range<InlayOffset>,
514        language_aware: bool,
515        text_highlights: Option<&'a TextHighlights>,
516        suggestion_highlight: Option<HighlightStyle>,
517    ) -> InlayChunks<'a> {
518        dbg!(self.transforms.items(&()));
519
520        let mut cursor = self.transforms.cursor::<(InlayOffset, SuggestionOffset)>();
521        cursor.seek(&range.start, Bias::Right, &());
522
523        let suggestion_range =
524            self.to_suggestion_offset(range.start)..self.to_suggestion_offset(range.end);
525        let suggestion_chunks = self.suggestion_snapshot.chunks(
526            suggestion_range,
527            language_aware,
528            text_highlights,
529            suggestion_highlight,
530        );
531
532        InlayChunks {
533            transforms: cursor,
534            suggestion_chunks,
535            inlay_chunks: None,
536            suggestion_chunk: None,
537            output_offset: range.start,
538            max_output_offset: range.end,
539        }
540    }
541
542    #[cfg(test)]
543    pub fn text(&self) -> String {
544        self.chunks(Default::default()..self.len(), false, None, None)
545            .map(|chunk| chunk.text)
546            .collect()
547    }
548}
549
550#[cfg(test)]
551mod tests {
552    use super::*;
553    use crate::{
554        display_map::{fold_map::FoldMap, suggestion_map::SuggestionMap},
555        MultiBuffer,
556    };
557    use gpui::AppContext;
558
559    #[gpui::test]
560    fn test_basic_inlays(cx: &mut AppContext) {
561        let buffer = MultiBuffer::build_simple("abcdefghi", cx);
562        let buffer_edits = buffer.update(cx, |buffer, _| buffer.subscribe());
563        let (mut fold_map, fold_snapshot) = FoldMap::new(buffer.read(cx).snapshot(cx));
564        let (suggestion_map, suggestion_snapshot) = SuggestionMap::new(fold_snapshot.clone());
565        let (mut inlay_map, inlay_snapshot) = InlayMap::new(suggestion_snapshot.clone());
566        assert_eq!(inlay_snapshot.text(), "abcdefghi");
567
568        let (inlay_snapshot, _, inlay_ids) = inlay_map.splice(
569            HashSet::default(),
570            vec![(
571                InlayHintLocation {
572                    buffer_id: 0,
573                    excerpt_id: ExcerptId::default(),
574                },
575                InlayProperties {
576                    position: buffer.read(cx).read(cx).anchor_before(3),
577                    text: "|123|".into(),
578                },
579            )],
580        );
581        assert_eq!(inlay_snapshot.text(), "abc|123|defghi");
582
583        buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "XYZ")], None, cx));
584        let (fold_snapshot, fold_edits) = fold_map.read(
585            buffer.read(cx).snapshot(cx),
586            buffer_edits.consume().into_inner(),
587        );
588        let (suggestion_snapshot, suggestion_edits) =
589            suggestion_map.sync(fold_snapshot.clone(), fold_edits);
590        let (inlay_snapshot, _) = inlay_map.sync(suggestion_snapshot.clone(), suggestion_edits);
591        assert_eq!(inlay_snapshot.text(), "XYZabc|123|defghi");
592
593        //////// case: folding and unfolding the text should hine and then return the hint back
594        let (mut fold_map_writer, _, _) = fold_map.write(
595            buffer.read(cx).snapshot(cx),
596            buffer_edits.consume().into_inner(),
597        );
598        let (fold_snapshot, fold_edits) = fold_map_writer.fold([4..8]);
599        let (suggestion_snapshot, suggestion_edits) =
600            suggestion_map.sync(fold_snapshot.clone(), fold_edits);
601        let (inlay_snapshot, _) = inlay_map.sync(suggestion_snapshot.clone(), suggestion_edits);
602        assert_eq!(inlay_snapshot.text(), "XYZa⋯fghi");
603
604        let (fold_snapshot, fold_edits) = fold_map_writer.unfold([4..8], false);
605        let (suggestion_snapshot, suggestion_edits) =
606            suggestion_map.sync(fold_snapshot.clone(), fold_edits);
607        let (inlay_snapshot, _) = inlay_map.sync(suggestion_snapshot.clone(), suggestion_edits);
608        assert_eq!(inlay_snapshot.text(), "XYZabc|123|defghi");
609
610        ////////// case: replacing the anchor that got the hint: it should disappear
611        buffer.update(cx, |buffer, cx| buffer.edit([(2..3, "C")], None, cx));
612        let (fold_snapshot, fold_edits) = fold_map.read(
613            buffer.read(cx).snapshot(cx),
614            buffer_edits.consume().into_inner(),
615        );
616        let (suggestion_snapshot, suggestion_edits) =
617            suggestion_map.sync(fold_snapshot.clone(), fold_edits);
618        let (inlay_snapshot, _) = inlay_map.sync(suggestion_snapshot.clone(), suggestion_edits);
619        assert_eq!(inlay_snapshot.text(), "XYZabCdefghi");
620    }
621}