inlay_hint_cache.rs

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