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    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.text.split_at(transform.len);
181                chunk.text = suffix;
182                self.output_offset.0 += prefix.len();
183                Chunk {
184                    text: prefix,
185                    ..chunk.clone()
186                }
187            }
188            Transform::Inlay(inlay) => {
189                let inlay_chunks = self.inlay_chunks.get_or_insert_with(|| {
190                    let start = self.output_offset - self.transforms.start().0;
191                    let end = cmp::min(self.max_output_offset, self.transforms.end(&()).0)
192                        - self.transforms.start().0;
193                    inlay.properties.text.chunks_in_range(start.0..end.0)
194                });
195
196                let chunk = inlay_chunks.next().unwrap();
197                self.output_offset.0 += chunk.len();
198                Chunk {
199                    text: chunk,
200                    ..Default::default()
201                }
202            }
203        };
204
205        if self.output_offset == self.transforms.end(&()).0 {
206            self.transforms.next(&());
207        }
208
209        Some(chunk)
210    }
211}
212
213impl<'a> Iterator for InlayBufferRows<'a> {
214    type Item = Option<u32>;
215
216    fn next(&mut self) -> Option<Self::Item> {
217        self.suggestion_rows.next()
218    }
219}
220
221impl InlayPoint {
222    pub fn new(row: u32, column: u32) -> Self {
223        Self(Point::new(row, column))
224    }
225
226    pub fn row(self) -> u32 {
227        self.0.row
228    }
229
230    pub fn column(self) -> u32 {
231        self.0.column
232    }
233}
234
235impl InlayMap {
236    pub fn new(suggestion_snapshot: SuggestionSnapshot) -> (Self, InlaySnapshot) {
237        let snapshot = InlaySnapshot {
238            suggestion_snapshot: suggestion_snapshot.clone(),
239            version: 0,
240            transforms: SumTree::from_item(
241                Transform::Isomorphic(suggestion_snapshot.text_summary()),
242                &(),
243            ),
244        };
245
246        (
247            Self {
248                snapshot: Mutex::new(snapshot.clone()),
249                next_inlay_id: 0,
250                inlays: HashMap::default(),
251            },
252            snapshot,
253        )
254    }
255
256    pub fn sync(
257        &self,
258        suggestion_snapshot: SuggestionSnapshot,
259        suggestion_edits: Vec<SuggestionEdit>,
260    ) -> (InlaySnapshot, Vec<InlayEdit>) {
261        let mut snapshot = self.snapshot.lock();
262
263        if snapshot.suggestion_snapshot.version != suggestion_snapshot.version {
264            snapshot.version += 1;
265        }
266
267        let mut inlay_edits = Vec::new();
268
269        dbg!(self.inlays.len());
270
271        for suggestion_edit in suggestion_edits {
272            let old = suggestion_edit.old;
273            let new = suggestion_edit.new;
274            // TODO kb copied from suggestion_map
275            inlay_edits.push(InlayEdit {
276                old: InlayOffset(old.start.0)..InlayOffset(old.end.0),
277                new: InlayOffset(old.start.0)..InlayOffset(new.end.0),
278            })
279        }
280
281        snapshot.suggestion_snapshot = suggestion_snapshot;
282
283        (snapshot.clone(), inlay_edits)
284    }
285
286    pub fn splice(
287        &mut self,
288        to_remove: HashSet<InlayId>,
289        to_insert: Vec<(InlayHintLocation, InlayProperties)>,
290    ) -> (InlaySnapshot, Vec<InlayEdit>, Vec<InlayId>) {
291        let mut snapshot = self.snapshot.lock();
292
293        let mut inlays = BTreeMap::new();
294        let mut new_ids = Vec::new();
295        for (location, properties) in to_insert {
296            let inlay = Inlay {
297                id: InlayId(post_inc(&mut self.next_inlay_id)),
298                properties,
299            };
300            self.inlays.insert(inlay.id, (location, inlay.clone()));
301            new_ids.push(inlay.id);
302
303            let buffer_point = inlay
304                .properties
305                .position
306                .to_point(snapshot.buffer_snapshot());
307            let fold_point = snapshot
308                .suggestion_snapshot
309                .fold_snapshot
310                .to_fold_point(buffer_point, Bias::Left);
311            let suggestion_point = snapshot.suggestion_snapshot.to_suggestion_point(fold_point);
312            let inlay_point = snapshot.to_inlay_point(suggestion_point);
313
314            inlays.insert((inlay_point, Reverse(inlay.id)), Some(inlay));
315        }
316
317        for inlay_id in to_remove {
318            if let Some((_, inlay)) = self.inlays.remove(&inlay_id) {
319                let buffer_point = inlay
320                    .properties
321                    .position
322                    .to_point(snapshot.buffer_snapshot());
323                let fold_point = snapshot
324                    .suggestion_snapshot
325                    .fold_snapshot
326                    .to_fold_point(buffer_point, Bias::Left);
327                let suggestion_point = snapshot.suggestion_snapshot.to_suggestion_point(fold_point);
328                let inlay_point = snapshot.to_inlay_point(suggestion_point);
329                inlays.insert((inlay_point, Reverse(inlay.id)), None);
330            }
331        }
332
333        let mut new_transforms = SumTree::new();
334        let mut cursor = snapshot
335            .transforms
336            .cursor::<(InlayPoint, SuggestionPoint)>();
337        for ((inlay_point, inlay_id), inlay) in inlays {
338            new_transforms.push_tree(cursor.slice(&inlay_point, Bias::Right, &()), &());
339            while let Some(transform) = cursor.item() {
340                match transform {
341                    Transform::Isomorphic(_) => break,
342                    Transform::Inlay(inlay) => {
343                        if inlay.id > inlay_id.0 {
344                            new_transforms.push(transform.clone(), &());
345                            cursor.next(&());
346                        } else {
347                            if inlay.id == inlay_id.0 {
348                                cursor.next(&());
349                            }
350                            break;
351                        }
352                    }
353                }
354            }
355
356            if let Some(inlay) = inlay {
357                if let Some(Transform::Isomorphic(transform)) = cursor.item() {
358                    let prefix = inlay_point.0 - cursor.start().0 .0;
359                    if !prefix.is_zero() {
360                        let prefix_suggestion_start = cursor.start().1;
361                        let prefix_suggestion_end = SuggestionPoint(cursor.start().1 .0 + prefix);
362                        new_transforms.push(
363                            Transform::Isomorphic(
364                                snapshot.suggestion_snapshot.text_summary_for_range(
365                                    prefix_suggestion_start..prefix_suggestion_end,
366                                ),
367                            ),
368                            &(),
369                        );
370                    }
371
372                    new_transforms.push(Transform::Inlay(inlay), &());
373
374                    let suffix_suggestion_start = SuggestionPoint(cursor.start().1 .0 + prefix);
375                    let suffix_suggestion_end = cursor.end(&()).1;
376                    new_transforms.push(
377                        Transform::Isomorphic(snapshot.suggestion_snapshot.text_summary_for_range(
378                            suffix_suggestion_start..suffix_suggestion_end,
379                        )),
380                        &(),
381                    );
382
383                    cursor.next(&());
384                } else {
385                    new_transforms.push(Transform::Inlay(inlay), &());
386                }
387            }
388        }
389
390        new_transforms.push_tree(cursor.suffix(&()), &());
391        drop(cursor);
392        snapshot.transforms = new_transforms;
393        snapshot.version += 1;
394
395        (snapshot.clone(), Vec::new(), new_ids)
396    }
397}
398
399impl InlaySnapshot {
400    pub fn buffer_snapshot(&self) -> &MultiBufferSnapshot {
401        // TODO kb copied from suggestion_map
402        self.suggestion_snapshot.buffer_snapshot()
403    }
404
405    pub fn to_point(&self, offset: InlayOffset) -> InlayPoint {
406        // TODO kb copied from suggestion_map
407        self.to_inlay_point(
408            self.suggestion_snapshot
409                .to_point(super::suggestion_map::SuggestionOffset(offset.0)),
410        )
411    }
412
413    pub fn len(&self) -> InlayOffset {
414        InlayOffset(self.transforms.summary().output.len)
415    }
416
417    pub fn max_point(&self) -> InlayPoint {
418        InlayPoint(self.transforms.summary().output.lines)
419    }
420
421    pub fn to_offset(&self, point: InlayPoint) -> InlayOffset {
422        // TODO kb copied from suggestion_map
423        InlayOffset(
424            self.suggestion_snapshot
425                .to_offset(self.to_suggestion_point(point, Bias::Left))
426                .0,
427        )
428    }
429
430    pub fn chars_at(&self, start: InlayPoint) -> impl '_ + Iterator<Item = char> {
431        self.suggestion_snapshot
432            .chars_at(self.to_suggestion_point(start, Bias::Left))
433    }
434
435    // TODO kb what to do with bias?
436    pub fn to_suggestion_point(&self, point: InlayPoint, _: Bias) -> SuggestionPoint {
437        SuggestionPoint(point.0)
438    }
439
440    pub fn to_suggestion_offset(&self, offset: InlayOffset) -> SuggestionOffset {
441        let mut cursor = self.transforms.cursor::<(InlayOffset, SuggestionOffset)>();
442        cursor.seek(&offset, Bias::Right, &());
443        match cursor.item() {
444            Some(Transform::Isomorphic(transform)) => {
445                let overshoot = offset - cursor.start().0;
446                cursor.start().1 + SuggestionOffset(overshoot.0)
447            }
448            Some(Transform::Inlay(inlay)) => cursor.start().1,
449            None => self.suggestion_snapshot.len(),
450        }
451    }
452
453    pub fn to_inlay_point(&self, point: SuggestionPoint) -> InlayPoint {
454        InlayPoint(point.0)
455    }
456
457    pub fn clip_point(&self, point: InlayPoint, bias: Bias) -> InlayPoint {
458        // TODO kb copied from suggestion_map
459        self.to_inlay_point(
460            self.suggestion_snapshot
461                .clip_point(self.to_suggestion_point(point, bias), bias),
462        )
463    }
464
465    pub fn text_summary_for_range(&self, range: Range<InlayPoint>) -> TextSummary {
466        // TODO kb copied from suggestion_map
467        self.suggestion_snapshot.text_summary_for_range(
468            self.to_suggestion_point(range.start, Bias::Left)
469                ..self.to_suggestion_point(range.end, Bias::Left),
470        )
471    }
472
473    pub fn buffer_rows<'a>(&'a self, row: u32) -> InlayBufferRows<'a> {
474        InlayBufferRows {
475            suggestion_rows: self.suggestion_snapshot.buffer_rows(row),
476        }
477    }
478
479    pub fn line_len(&self, row: u32) -> u32 {
480        // TODO kb copied from suggestion_map
481        self.suggestion_snapshot.line_len(row)
482    }
483
484    pub fn chunks<'a>(
485        &'a self,
486        range: Range<InlayOffset>,
487        language_aware: bool,
488        text_highlights: Option<&'a TextHighlights>,
489        suggestion_highlight: Option<HighlightStyle>,
490    ) -> InlayChunks<'a> {
491        dbg!(self.transforms.items(&()));
492
493        let mut cursor = self.transforms.cursor::<(InlayOffset, SuggestionOffset)>();
494        cursor.seek(&range.start, Bias::Right, &());
495
496        let suggestion_range =
497            self.to_suggestion_offset(range.start)..self.to_suggestion_offset(range.end);
498        let suggestion_chunks = self.suggestion_snapshot.chunks(
499            suggestion_range,
500            language_aware,
501            text_highlights,
502            suggestion_highlight,
503        );
504
505        InlayChunks {
506            transforms: cursor,
507            suggestion_chunks,
508            inlay_chunks: None,
509            suggestion_chunk: None,
510            output_offset: range.start,
511            max_output_offset: range.end,
512        }
513    }
514
515    #[cfg(test)]
516    pub fn text(&self) -> String {
517        self.chunks(Default::default()..self.len(), false, None, None)
518            .map(|chunk| chunk.text)
519            .collect()
520    }
521}
522
523#[cfg(test)]
524mod tests {
525    use super::*;
526    use crate::{
527        display_map::{fold_map::FoldMap, suggestion_map::SuggestionMap},
528        MultiBuffer,
529    };
530    use gpui::AppContext;
531
532    #[gpui::test]
533    fn test_basic_inlays(cx: &mut AppContext) {
534        let buffer = MultiBuffer::build_simple("abcdefghi", cx);
535        let buffer_edits = buffer.update(cx, |buffer, _| buffer.subscribe());
536        let (mut fold_map, fold_snapshot) = FoldMap::new(buffer.read(cx).snapshot(cx));
537        let (suggestion_map, suggestion_snapshot) = SuggestionMap::new(fold_snapshot.clone());
538        let (mut inlay_map, inlay_snapshot) = InlayMap::new(suggestion_snapshot.clone());
539        assert_eq!(inlay_snapshot.text(), "abcdefghi");
540
541        let (inlay_snapshot, _, inlay_ids) = inlay_map.splice(
542            HashSet::default(),
543            vec![(
544                InlayHintLocation {
545                    buffer_id: 0,
546                    excerpt_id: ExcerptId::default(),
547                },
548                InlayProperties {
549                    position: buffer.read(cx).read(cx).anchor_before(3),
550                    text: "|123|".into(),
551                },
552            )],
553        );
554        assert_eq!(inlay_snapshot.text(), "abc|123|defghi");
555
556        buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "XYZ")], None, cx));
557        let (fold_snapshot, fold_edits) = fold_map.read(
558            buffer.read(cx).snapshot(cx),
559            buffer_edits.consume().into_inner(),
560        );
561        let (suggestion_snapshot, suggestion_edits) =
562            suggestion_map.sync(fold_snapshot.clone(), fold_edits);
563        let (inlay_snapshot, _) = inlay_map.sync(suggestion_snapshot.clone(), suggestion_edits);
564        assert_eq!(inlay_snapshot.text(), "XYZabc|123|defghi");
565
566        //////// case: folding and unfolding the text should hine and then return the hint back
567        let (mut fold_map_writer, _, _) = fold_map.write(
568            buffer.read(cx).snapshot(cx),
569            buffer_edits.consume().into_inner(),
570        );
571        let (fold_snapshot, fold_edits) = fold_map_writer.fold([4..8]);
572        let (suggestion_snapshot, suggestion_edits) =
573            suggestion_map.sync(fold_snapshot.clone(), fold_edits);
574        let (inlay_snapshot, _) = inlay_map.sync(suggestion_snapshot.clone(), suggestion_edits);
575        assert_eq!(inlay_snapshot.text(), "XYZa⋯fghi");
576
577        let (fold_snapshot, fold_edits) = fold_map_writer.unfold([4..8], false);
578        let (suggestion_snapshot, suggestion_edits) =
579            suggestion_map.sync(fold_snapshot.clone(), fold_edits);
580        let (inlay_snapshot, _) = inlay_map.sync(suggestion_snapshot.clone(), suggestion_edits);
581        assert_eq!(inlay_snapshot.text(), "XYZabc|123|defghi");
582
583        ////////// case: replacing the anchor that got the hint: it should disappear, then undo and it should reappear again
584        buffer.update(cx, |buffer, cx| buffer.edit([(2..3, "C")], None, cx));
585        let (fold_snapshot, fold_edits) = fold_map.read(
586            buffer.read(cx).snapshot(cx),
587            buffer_edits.consume().into_inner(),
588        );
589        let (suggestion_snapshot, suggestion_edits) =
590            suggestion_map.sync(fold_snapshot.clone(), fold_edits);
591        let (inlay_snapshot, _) = inlay_map.sync(suggestion_snapshot.clone(), suggestion_edits);
592        assert_eq!(inlay_snapshot.text(), "XYZabCdefghi");
593
594        buffer.update(cx, |buffer, cx| buffer.undo(cx));
595        let (fold_snapshot, fold_edits) = fold_map.read(
596            buffer.read(cx).snapshot(cx),
597            buffer_edits.consume().into_inner(),
598        );
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(), "XYZabc|123|defghi");
603    }
604}