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