inlay_hint_cache.rs

  1use std::{cmp, sync::Arc};
  2
  3use crate::{
  4    display_map::Inlay, editor_settings, Anchor, Editor, ExcerptId, InlayId, MultiBuffer,
  5    MultiBufferSnapshot,
  6};
  7use anyhow::Context;
  8use gpui::{ModelHandle, Task, ViewContext};
  9use language::{Buffer, BufferSnapshot};
 10use log::error;
 11use project::{InlayHint, InlayHintKind};
 12
 13use collections::{hash_map, HashMap, HashSet};
 14use util::post_inc;
 15
 16pub struct InlayHintCache {
 17    snapshot: CacheSnapshot,
 18    update_tasks: HashMap<ExcerptId, InlayHintUpdateTask>,
 19}
 20
 21struct InlayHintUpdateTask {
 22    version: usize,
 23    _task: Task<()>,
 24}
 25
 26struct CacheSnapshot {
 27    hints: HashMap<ExcerptId, Arc<CachedExcerptHints>>,
 28    allowed_hint_kinds: HashSet<Option<InlayHintKind>>,
 29    version: usize,
 30}
 31
 32struct CachedExcerptHints {
 33    version: usize,
 34    hints: Vec<(InlayId, InlayHint)>,
 35}
 36
 37#[derive(Debug, Clone, Copy)]
 38struct ExcerptQuery {
 39    buffer_id: u64,
 40    excerpt_id: ExcerptId,
 41    excerpt_range_start: language::Anchor,
 42    excerpt_range_end: language::Anchor,
 43    cache_version: usize,
 44    invalidate_cache: bool,
 45}
 46
 47#[derive(Debug, Default)]
 48pub struct InlaySplice {
 49    pub to_remove: Vec<InlayId>,
 50    pub to_insert: Vec<(Anchor, InlayId, InlayHint)>,
 51}
 52
 53#[derive(Debug)]
 54struct ExcerptHintsUpdate {
 55    excerpt_id: ExcerptId,
 56    cache_version: usize,
 57    remove_from_visible: Vec<InlayId>,
 58    remove_from_cache: HashSet<InlayId>,
 59    add_to_cache: Vec<InlayHint>,
 60}
 61
 62impl InlayHintCache {
 63    pub fn new(inlay_hint_settings: editor_settings::InlayHints) -> Self {
 64        Self {
 65            snapshot: CacheSnapshot {
 66                allowed_hint_kinds: allowed_hint_types(inlay_hint_settings),
 67                hints: HashMap::default(),
 68                version: 0,
 69            },
 70            update_tasks: HashMap::default(),
 71        }
 72    }
 73
 74    pub fn update_settings(
 75        &mut self,
 76        multi_buffer: &ModelHandle<MultiBuffer>,
 77        inlay_hint_settings: editor_settings::InlayHints,
 78        visible_hints: Vec<Inlay>,
 79        cx: &mut ViewContext<Editor>,
 80    ) -> Option<InlaySplice> {
 81        let new_allowed_hint_kinds = allowed_hint_types(inlay_hint_settings);
 82        if !inlay_hint_settings.enabled {
 83            if self.snapshot.hints.is_empty() {
 84                self.snapshot.allowed_hint_kinds = new_allowed_hint_kinds;
 85            } else {
 86                self.clear();
 87                self.snapshot.allowed_hint_kinds = new_allowed_hint_kinds;
 88                return Some(InlaySplice {
 89                    to_remove: visible_hints.iter().map(|inlay| inlay.id).collect(),
 90                    to_insert: Vec::new(),
 91                });
 92            }
 93
 94            return None;
 95        }
 96
 97        if new_allowed_hint_kinds == self.snapshot.allowed_hint_kinds {
 98            return None;
 99        }
100
101        let new_splice = new_allowed_hint_kinds_splice(
102            &self.snapshot,
103            multi_buffer,
104            &visible_hints,
105            &new_allowed_hint_kinds,
106            cx,
107        );
108        if new_splice.is_some() {
109            self.snapshot.version += 1;
110            self.update_tasks.clear();
111            self.snapshot.allowed_hint_kinds = new_allowed_hint_kinds;
112        }
113        new_splice
114    }
115
116    pub fn spawn_hints_update(
117        &mut self,
118        mut excerpts_to_query: HashMap<ExcerptId, ModelHandle<Buffer>>,
119        invalidate_cache: bool,
120        cx: &mut ViewContext<Editor>,
121    ) {
122        let update_tasks = &mut self.update_tasks;
123        if invalidate_cache {
124            update_tasks
125                .retain(|task_excerpt_id, _| excerpts_to_query.contains_key(task_excerpt_id));
126        }
127        excerpts_to_query.retain(|visible_excerpt_id, _| {
128            match update_tasks.entry(*visible_excerpt_id) {
129                hash_map::Entry::Occupied(o) => match o.get().version.cmp(&self.snapshot.version) {
130                    cmp::Ordering::Less => true,
131                    cmp::Ordering::Equal => invalidate_cache,
132                    cmp::Ordering::Greater => false,
133                },
134                hash_map::Entry::Vacant(_) => true,
135            }
136        });
137
138        if invalidate_cache {
139            update_tasks
140                .retain(|task_excerpt_id, _| excerpts_to_query.contains_key(task_excerpt_id));
141        }
142        let cache_version = self.snapshot.version;
143        excerpts_to_query.retain(|visible_excerpt_id, _| {
144            match update_tasks.entry(*visible_excerpt_id) {
145                hash_map::Entry::Occupied(o) => match o.get().version.cmp(&cache_version) {
146                    cmp::Ordering::Less => true,
147                    cmp::Ordering::Equal => invalidate_cache,
148                    cmp::Ordering::Greater => false,
149                },
150                hash_map::Entry::Vacant(_) => true,
151            }
152        });
153
154        cx.spawn(|editor, mut cx| async move {
155            editor
156                .update(&mut cx, |editor, cx| {
157                    let visible_hints =
158                        Arc::new(visible_inlay_hints(editor, cx).cloned().collect::<Vec<_>>());
159                    for (excerpt_id, buffer_handle) in excerpts_to_query {
160                        let (multi_buffer_snapshot, excerpt_range) =
161                            editor.buffer.update(cx, |multi_buffer, cx| {
162                                let multi_buffer_snapshot = multi_buffer.snapshot(cx);
163                                (
164                                    multi_buffer_snapshot,
165                                    multi_buffer
166                                        .excerpts_for_buffer(&buffer_handle, cx)
167                                        .into_iter()
168                                        .find(|(id, _)| id == &excerpt_id)
169                                        .map(|(_, range)| range.context),
170                                )
171                            });
172
173                        if let Some(excerpt_range) = excerpt_range {
174                            let buffer = buffer_handle.read(cx);
175                            let buffer_snapshot = buffer.snapshot();
176                            let query = ExcerptQuery {
177                                buffer_id: buffer.remote_id(),
178                                excerpt_id,
179                                excerpt_range_start: excerpt_range.start,
180                                excerpt_range_end: excerpt_range.end,
181                                cache_version,
182                                invalidate_cache,
183                            };
184                            let cached_excxerpt_hints = editor
185                                .inlay_hint_cache
186                                .snapshot
187                                .hints
188                                .get(&excerpt_id)
189                                .cloned();
190                            editor.inlay_hint_cache.update_tasks.insert(
191                                excerpt_id,
192                                new_update_task(
193                                    query,
194                                    multi_buffer_snapshot,
195                                    buffer_snapshot,
196                                    Arc::clone(&visible_hints),
197                                    cached_excxerpt_hints,
198                                    cx,
199                                ),
200                            );
201                        }
202                    }
203                })
204                .ok();
205        })
206        .detach();
207    }
208
209    fn clear(&mut self) {
210        self.snapshot.version += 1;
211        self.update_tasks.clear();
212        self.snapshot.hints.clear();
213        self.snapshot.allowed_hint_kinds.clear();
214    }
215}
216
217fn new_update_task(
218    query: ExcerptQuery,
219    multi_buffer_snapshot: MultiBufferSnapshot,
220    buffer_snapshot: BufferSnapshot,
221    visible_hints: Arc<Vec<Inlay>>,
222    cached_excerpt_hints: Option<Arc<CachedExcerptHints>>,
223    cx: &mut ViewContext<'_, '_, Editor>,
224) -> InlayHintUpdateTask {
225    let hints_fetch_task = hints_fetch_task(query, cx);
226    InlayHintUpdateTask {
227        version: query.cache_version,
228        _task: cx.spawn(|editor, mut cx| async move {
229            match hints_fetch_task.await {
230                Ok(Some(new_hints)) => {
231                    let task_buffer_snapshot = buffer_snapshot.clone();
232                    if let Some(new_update) = cx
233                        .background()
234                        .spawn(async move {
235                            new_excerpt_hints_update_result(
236                                query,
237                                new_hints,
238                                &task_buffer_snapshot,
239                                cached_excerpt_hints,
240                                &visible_hints,
241                            )
242                        })
243                        .await
244                    {
245                        editor
246                            .update(&mut cx, |editor, cx| {
247                                let cached_excerpt_hints = editor
248                                    .inlay_hint_cache
249                                    .snapshot
250                                    .hints
251                                    .entry(new_update.excerpt_id)
252                                    .or_insert_with(|| {
253                                        Arc::new(CachedExcerptHints {
254                                            version: new_update.cache_version,
255                                            hints: Vec::new(),
256                                        })
257                                    });
258                                let cached_excerpt_hints = Arc::get_mut(cached_excerpt_hints)
259                                    .expect("Cached excerot hints were dropped with the task");
260
261                                match new_update.cache_version.cmp(&cached_excerpt_hints.version) {
262                                    cmp::Ordering::Less => return,
263                                    cmp::Ordering::Greater | cmp::Ordering::Equal => {
264                                        cached_excerpt_hints.version = new_update.cache_version;
265                                    }
266                                }
267                                cached_excerpt_hints.hints.retain(|(hint_id, _)| {
268                                    !new_update.remove_from_cache.contains(hint_id)
269                                });
270
271                                editor.inlay_hint_cache.snapshot.version += 1;
272
273                                let mut splice = InlaySplice {
274                                    to_remove: new_update.remove_from_visible,
275                                    to_insert: Vec::new(),
276                                };
277
278                                for new_hint in new_update.add_to_cache {
279                                    let new_hint_position = multi_buffer_snapshot
280                                        .anchor_in_excerpt(query.excerpt_id, new_hint.position);
281                                    let new_inlay_id = InlayId(post_inc(&mut editor.next_inlay_id));
282                                    if editor
283                                        .inlay_hint_cache
284                                        .snapshot
285                                        .allowed_hint_kinds
286                                        .contains(&new_hint.kind)
287                                    {
288                                        splice.to_insert.push((
289                                            new_hint_position,
290                                            new_inlay_id,
291                                            new_hint.clone(),
292                                        ));
293                                    }
294
295                                    cached_excerpt_hints.hints.push((new_inlay_id, new_hint));
296                                }
297
298                                cached_excerpt_hints
299                                    .hints
300                                    .sort_by(|(_, hint_a), (_, hint_b)| {
301                                        hint_a.position.cmp(&hint_b.position, &buffer_snapshot)
302                                    });
303
304                                let InlaySplice {
305                                    to_remove,
306                                    to_insert,
307                                } = splice;
308                                if !to_remove.is_empty() || !to_insert.is_empty() {
309                                    editor.splice_inlay_hints(to_remove, to_insert, cx)
310                                }
311                            })
312                            .ok();
313                    }
314                }
315                Ok(None) => {}
316                Err(e) => error!(
317                    "Failed to fecth hints for excerpt {:?} in buffer {} : {}",
318                    query.excerpt_id, query.buffer_id, e
319                ),
320            }
321        }),
322    }
323}
324
325fn new_allowed_hint_kinds_splice(
326    cache: &CacheSnapshot,
327    multi_buffer: &ModelHandle<MultiBuffer>,
328    visible_hints: &[Inlay],
329    new_kinds: &HashSet<Option<InlayHintKind>>,
330    cx: &mut ViewContext<Editor>,
331) -> Option<InlaySplice> {
332    let old_kinds = &cache.allowed_hint_kinds;
333    if new_kinds == old_kinds {
334        return None;
335    }
336
337    let mut to_remove = Vec::new();
338    let mut to_insert = Vec::new();
339    let mut shown_hints_to_remove = visible_hints.iter().fold(
340        HashMap::<ExcerptId, Vec<(Anchor, InlayId)>>::default(),
341        |mut current_hints, inlay| {
342            current_hints
343                .entry(inlay.position.excerpt_id)
344                .or_default()
345                .push((inlay.position, inlay.id));
346            current_hints
347        },
348    );
349
350    let multi_buffer = multi_buffer.read(cx);
351    let multi_buffer_snapshot = multi_buffer.snapshot(cx);
352
353    for (excerpt_id, excerpt_cached_hints) in &cache.hints {
354        let shown_excerpt_hints_to_remove = shown_hints_to_remove.entry(*excerpt_id).or_default();
355        let mut excerpt_cache = excerpt_cached_hints.hints.iter().fuse().peekable();
356        shown_excerpt_hints_to_remove.retain(|(shown_anchor, shown_hint_id)| {
357            let Some(buffer) = shown_anchor
358                .buffer_id
359                .and_then(|buffer_id| multi_buffer.buffer(buffer_id)) else { return false };
360            let buffer_snapshot = buffer.read(cx).snapshot();
361            loop {
362                match excerpt_cache.peek() {
363                    Some((cached_hint_id, cached_hint)) => {
364                        if cached_hint_id == shown_hint_id {
365                            excerpt_cache.next();
366                            return !new_kinds.contains(&cached_hint.kind);
367                        }
368
369                        match cached_hint
370                            .position
371                            .cmp(&shown_anchor.text_anchor, &buffer_snapshot)
372                        {
373                            cmp::Ordering::Less | cmp::Ordering::Equal => {
374                                if !old_kinds.contains(&cached_hint.kind)
375                                    && new_kinds.contains(&cached_hint.kind)
376                                {
377                                    to_insert.push((
378                                        multi_buffer_snapshot
379                                            .anchor_in_excerpt(*excerpt_id, cached_hint.position),
380                                        *cached_hint_id,
381                                        cached_hint.clone(),
382                                    ));
383                                }
384                                excerpt_cache.next();
385                            }
386                            cmp::Ordering::Greater => return true,
387                        }
388                    }
389                    None => return true,
390                }
391            }
392        });
393
394        for (cached_hint_id, maybe_missed_cached_hint) in excerpt_cache {
395            let cached_hint_kind = maybe_missed_cached_hint.kind;
396            if !old_kinds.contains(&cached_hint_kind) && new_kinds.contains(&cached_hint_kind) {
397                to_insert.push((
398                    multi_buffer_snapshot
399                        .anchor_in_excerpt(*excerpt_id, maybe_missed_cached_hint.position),
400                    *cached_hint_id,
401                    maybe_missed_cached_hint.clone(),
402                ));
403            }
404        }
405    }
406
407    to_remove.extend(
408        shown_hints_to_remove
409            .into_values()
410            .flatten()
411            .map(|(_, hint_id)| hint_id),
412    );
413    if to_remove.is_empty() && to_insert.is_empty() {
414        None
415    } else {
416        Some(InlaySplice {
417            to_remove,
418            to_insert,
419        })
420    }
421}
422
423fn new_excerpt_hints_update_result(
424    query: ExcerptQuery,
425    new_excerpt_hints: Vec<InlayHint>,
426    buffer_snapshot: &BufferSnapshot,
427    cached_excerpt_hints: Option<Arc<CachedExcerptHints>>,
428    visible_hints: &[Inlay],
429) -> Option<ExcerptHintsUpdate> {
430    let mut add_to_cache: Vec<InlayHint> = Vec::new();
431
432    let mut excerpt_hints_to_persist = HashMap::default();
433    for new_hint in new_excerpt_hints {
434        let missing_from_cache = match &cached_excerpt_hints {
435            Some(cached_excerpt_hints) => {
436                match cached_excerpt_hints.hints.binary_search_by(|probe| {
437                    probe.1.position.cmp(&new_hint.position, buffer_snapshot)
438                }) {
439                    Ok(ix) => {
440                        let (cached_inlay_id, cached_hint) = &cached_excerpt_hints.hints[ix];
441                        if cached_hint == &new_hint {
442                            excerpt_hints_to_persist.insert(*cached_inlay_id, cached_hint.kind);
443                            false
444                        } else {
445                            true
446                        }
447                    }
448                    Err(_) => true,
449                }
450            }
451            None => true,
452        };
453        if missing_from_cache {
454            add_to_cache.push(new_hint);
455        }
456    }
457
458    let mut remove_from_visible = Vec::new();
459    let mut remove_from_cache = HashSet::default();
460    if query.invalidate_cache {
461        remove_from_visible.extend(
462            visible_hints
463                .iter()
464                .filter(|hint| hint.position.excerpt_id == query.excerpt_id)
465                .filter(|hint| {
466                    query
467                        .excerpt_range_start
468                        .cmp(&hint.position.text_anchor, buffer_snapshot)
469                        .is_le()
470                })
471                .filter(|hint| {
472                    query
473                        .excerpt_range_end
474                        .cmp(&hint.position.text_anchor, buffer_snapshot)
475                        .is_ge()
476                })
477                .map(|inlay_hint| inlay_hint.id)
478                .filter(|hint_id| !excerpt_hints_to_persist.contains_key(hint_id)),
479        );
480        remove_from_cache.extend(
481            cached_excerpt_hints
482                .iter()
483                .flat_map(|excerpt_hints| excerpt_hints.hints.iter())
484                .filter(|(cached_inlay_id, _)| {
485                    !excerpt_hints_to_persist.contains_key(cached_inlay_id)
486                })
487                .map(|(cached_inlay_id, _)| *cached_inlay_id),
488        );
489    }
490
491    if remove_from_visible.is_empty() && remove_from_cache.is_empty() && add_to_cache.is_empty() {
492        None
493    } else {
494        Some(ExcerptHintsUpdate {
495            cache_version: query.cache_version,
496            excerpt_id: query.excerpt_id,
497            remove_from_visible,
498            remove_from_cache,
499            add_to_cache,
500        })
501    }
502}
503
504fn allowed_hint_types(
505    inlay_hint_settings: editor_settings::InlayHints,
506) -> HashSet<Option<InlayHintKind>> {
507    let mut new_allowed_hint_types = HashSet::default();
508    if inlay_hint_settings.show_type_hints {
509        new_allowed_hint_types.insert(Some(InlayHintKind::Type));
510    }
511    if inlay_hint_settings.show_parameter_hints {
512        new_allowed_hint_types.insert(Some(InlayHintKind::Parameter));
513    }
514    if inlay_hint_settings.show_other_hints {
515        new_allowed_hint_types.insert(None);
516    }
517    new_allowed_hint_types
518}
519
520fn hints_fetch_task(
521    query: ExcerptQuery,
522    cx: &mut ViewContext<'_, '_, Editor>,
523) -> Task<anyhow::Result<Option<Vec<InlayHint>>>> {
524    cx.spawn(|editor, mut cx| async move {
525        let task = editor
526            .update(&mut cx, |editor, cx| {
527                editor
528                    .buffer()
529                    .read(cx)
530                    .buffer(query.buffer_id)
531                    .and_then(|buffer| {
532                        let project = editor.project.as_ref()?;
533                        Some(project.update(cx, |project, cx| {
534                            project.inlay_hints(
535                                buffer,
536                                query.excerpt_range_start..query.excerpt_range_end,
537                                cx,
538                            )
539                        }))
540                    })
541            })
542            .ok()
543            .flatten();
544        Ok(match task {
545            Some(task) => Some(task.await.context("inlays for buffer task")?),
546            None => None,
547        })
548    })
549}
550
551pub fn visible_inlay_hints<'a, 'b: 'a, 'c, 'd: 'a>(
552    editor: &'a Editor,
553    cx: &'b ViewContext<'c, 'd, Editor>,
554) -> impl Iterator<Item = &'b Inlay> + 'a {
555    editor
556        .display_map
557        .read(cx)
558        .current_inlays()
559        .filter(|inlay| Some(inlay.id) != editor.copilot_state.suggestion.as_ref().map(|h| h.id))
560}