inlay_hint_cache.rs

  1use std::ops::Range;
  2
  3use crate::{
  4    editor_settings, scroll::ScrollAnchor, Anchor, Editor, ExcerptId, InlayId, MultiBuffer,
  5};
  6use anyhow::Context;
  7use clock::Global;
  8use gpui::{ModelHandle, Task, ViewContext};
  9use language::Buffer;
 10use log::error;
 11use project::{InlayHint, InlayHintKind};
 12
 13use collections::{hash_map, HashMap, HashSet};
 14use util::post_inc;
 15
 16#[derive(Debug, Copy, Clone)]
 17pub enum InlayRefreshReason {
 18    SettingsChange(editor_settings::InlayHints),
 19    Scroll(ScrollAnchor),
 20    VisibleExcerptsChange,
 21}
 22
 23#[derive(Debug, Default)]
 24pub struct InlayHintCache {
 25    inlay_hints: HashMap<InlayId, InlayHint>,
 26    hints_in_buffers: HashMap<u64, BufferHints<(Anchor, InlayId)>>,
 27    allowed_hint_kinds: HashSet<Option<InlayHintKind>>,
 28}
 29
 30#[derive(Clone, Debug)]
 31struct BufferHints<H> {
 32    buffer_version: Global,
 33    hints_per_excerpt: HashMap<ExcerptId, ExcerptHints<H>>,
 34}
 35
 36#[derive(Clone, Debug)]
 37struct ExcerptHints<H> {
 38    cached_excerpt_offsets: Vec<Range<usize>>,
 39    hints: Vec<H>,
 40}
 41
 42impl<H> Default for ExcerptHints<H> {
 43    fn default() -> Self {
 44        Self {
 45            cached_excerpt_offsets: Vec::new(),
 46            hints: Vec::new(),
 47        }
 48    }
 49}
 50
 51impl<H> BufferHints<H> {
 52    fn new(buffer_version: Global) -> Self {
 53        Self {
 54            buffer_version,
 55            hints_per_excerpt: HashMap::default(),
 56        }
 57    }
 58}
 59
 60#[derive(Debug, Default)]
 61pub struct InlaySplice {
 62    pub to_remove: Vec<InlayId>,
 63    pub to_insert: Vec<(InlayId, Anchor, InlayHint)>,
 64}
 65
 66pub struct InlayHintQuery {
 67    pub buffer_id: u64,
 68    pub buffer_version: Global,
 69    pub excerpt_id: ExcerptId,
 70    pub excerpt_offset_query_range: Range<usize>,
 71}
 72
 73impl InlayHintCache {
 74    pub fn new(inlay_hint_settings: editor_settings::InlayHints) -> Self {
 75        Self {
 76            allowed_hint_kinds: allowed_hint_types(inlay_hint_settings),
 77            hints_in_buffers: HashMap::default(),
 78            inlay_hints: HashMap::default(),
 79        }
 80    }
 81
 82    pub fn apply_settings(
 83        &mut self,
 84        inlay_hint_settings: editor_settings::InlayHints,
 85        currently_visible_ranges: Vec<(ModelHandle<Buffer>, Range<usize>, ExcerptId)>,
 86        mut currently_shown_hints: HashMap<u64, HashMap<ExcerptId, Vec<(Anchor, InlayId)>>>,
 87        cx: &mut ViewContext<Editor>,
 88    ) -> Option<InlaySplice> {
 89        let new_allowed_hint_kinds = allowed_hint_types(inlay_hint_settings);
 90        if new_allowed_hint_kinds == self.allowed_hint_kinds {
 91            None
 92        } else {
 93            self.allowed_hint_kinds = new_allowed_hint_kinds;
 94            let mut to_remove = Vec::new();
 95            let mut to_insert = Vec::new();
 96            let mut considered_hints =
 97                HashMap::<u64, HashMap<ExcerptId, HashSet<InlayId>>>::default();
 98            for (visible_buffer, _, visible_excerpt_id) in currently_visible_ranges {
 99                let visible_buffer = visible_buffer.read(cx);
100                let visible_buffer_id = visible_buffer.remote_id();
101                match currently_shown_hints.entry(visible_buffer_id) {
102                    hash_map::Entry::Occupied(mut o) => {
103                        let shown_hints_per_excerpt = o.get_mut();
104                        for (_, shown_hint_id) in shown_hints_per_excerpt
105                            .remove(&visible_excerpt_id)
106                            .unwrap_or_default()
107                        {
108                            considered_hints
109                                .entry(visible_buffer_id)
110                                .or_default()
111                                .entry(visible_excerpt_id)
112                                .or_default()
113                                .insert(shown_hint_id);
114                            match self.inlay_hints.get(&shown_hint_id) {
115                                Some(shown_hint) => {
116                                    if !self.allowed_hint_kinds.contains(&shown_hint.kind) {
117                                        to_remove.push(shown_hint_id);
118                                    }
119                                }
120                                None => to_remove.push(shown_hint_id),
121                            }
122                        }
123                        if shown_hints_per_excerpt.is_empty() {
124                            o.remove();
125                        }
126                    }
127                    hash_map::Entry::Vacant(_) => {}
128                }
129            }
130
131            let reenabled_hints = self
132                .hints_in_buffers
133                .iter()
134                .filter_map(|(cached_buffer_id, cached_hints_per_excerpt)| {
135                    let considered_hints_in_excerpts = considered_hints.get(cached_buffer_id)?;
136                    let not_considered_cached_hints = cached_hints_per_excerpt
137                        .hints_per_excerpt
138                        .iter()
139                        .filter_map(|(cached_excerpt_id, cached_excerpt_hints)| {
140                            let considered_excerpt_hints =
141                                considered_hints_in_excerpts.get(&cached_excerpt_id)?;
142                            let not_considered_cached_hints = cached_excerpt_hints
143                                .hints
144                                .iter()
145                                .filter(|(_, cached_hint_id)| {
146                                    !considered_excerpt_hints.contains(cached_hint_id)
147                                })
148                                .copied();
149                            Some(not_considered_cached_hints)
150                        })
151                        .flatten();
152                    Some(not_considered_cached_hints)
153                })
154                .flatten()
155                .filter_map(|(cached_anchor, cached_hint_id)| {
156                    Some((
157                        cached_anchor,
158                        cached_hint_id,
159                        self.inlay_hints.get(&cached_hint_id)?,
160                    ))
161                })
162                .filter(|(_, _, cached_hint)| self.allowed_hint_kinds.contains(&cached_hint.kind))
163                .map(|(cached_anchor, cached_hint_id, reenabled_hint)| {
164                    (cached_hint_id, cached_anchor, reenabled_hint.clone())
165                });
166            to_insert.extend(reenabled_hints);
167
168            to_remove.extend(
169                currently_shown_hints
170                    .into_iter()
171                    .flat_map(|(_, hints_by_excerpt)| hints_by_excerpt)
172                    .flat_map(|(_, excerpt_hints)| excerpt_hints)
173                    .map(|(_, hint_id)| hint_id),
174            );
175
176            Some(InlaySplice {
177                to_remove,
178                to_insert,
179            })
180        }
181    }
182
183    pub fn clear(&mut self) -> Vec<InlayId> {
184        let ids_to_remove = self.inlay_hints.drain().map(|(id, _)| id).collect();
185        self.hints_in_buffers.clear();
186        ids_to_remove
187    }
188
189    pub fn append_hints(
190        &mut self,
191        multi_buffer: ModelHandle<MultiBuffer>,
192        ranges_to_add: impl Iterator<Item = InlayHintQuery>,
193        cx: &mut ViewContext<Editor>,
194    ) -> Task<anyhow::Result<InlaySplice>> {
195        let queries = filter_queries(ranges_to_add, &self.hints_in_buffers, false);
196
197        let task_multi_buffer = multi_buffer.clone();
198        let fetch_queries_task = fetch_queries(multi_buffer, queries.into_iter(), cx);
199        cx.spawn(|editor, mut cx| async move {
200            let new_hints = fetch_queries_task.await?;
201            editor.update(&mut cx, |editor, cx| {
202                let multi_buffer_snapshot = task_multi_buffer.read(cx).snapshot(cx);
203                let mut to_insert = Vec::new();
204                for (new_buffer_id, new_hints_per_buffer) in new_hints {
205                    let cached_buffer_hints = editor
206                        .inlay_hint_cache
207                        .hints_in_buffers
208                        .entry(new_buffer_id)
209                        .or_insert_with(|| {
210                            BufferHints::new(new_hints_per_buffer.buffer_version.clone())
211                        });
212                    if cached_buffer_hints
213                        .buffer_version
214                        .changed_since(&new_hints_per_buffer.buffer_version)
215                    {
216                        continue;
217                    }
218
219                    for (new_excerpt_id, new_excerpt_hints) in
220                        new_hints_per_buffer.hints_per_excerpt
221                    {
222                        let cached_excerpt_hints = cached_buffer_hints
223                            .hints_per_excerpt
224                            .entry(new_excerpt_id)
225                            .or_insert_with(|| ExcerptHints::default());
226                        for new_range in new_excerpt_hints.cached_excerpt_offsets {
227                            insert_and_merge_ranges(
228                                &mut cached_excerpt_hints.cached_excerpt_offsets,
229                                &new_range,
230                            )
231                        }
232                        for new_hint in new_excerpt_hints.hints {
233                            let hint_anchor = multi_buffer_snapshot
234                                .anchor_in_excerpt(new_excerpt_id, new_hint.position);
235                            let insert_ix =
236                                match cached_excerpt_hints.hints.binary_search_by(|probe| {
237                                    hint_anchor.cmp(&probe.0, &multi_buffer_snapshot)
238                                }) {
239                                    Ok(ix) | Err(ix) => ix,
240                                };
241
242                            let new_hint_id = InlayId(post_inc(&mut editor.next_inlay_id));
243                            cached_excerpt_hints
244                                .hints
245                                .insert(insert_ix, (hint_anchor, new_hint_id));
246                            editor
247                                .inlay_hint_cache
248                                .inlay_hints
249                                .insert(new_hint_id, new_hint.clone());
250                            if editor
251                                .inlay_hint_cache
252                                .allowed_hint_kinds
253                                .contains(&new_hint.kind)
254                            {
255                                to_insert.push((new_hint_id, hint_anchor, new_hint));
256                            }
257                        }
258                    }
259                }
260
261                InlaySplice {
262                    to_remove: Vec::new(),
263                    to_insert,
264                }
265            })
266        })
267    }
268
269    pub fn replace_hints(
270        &mut self,
271        multi_buffer: ModelHandle<MultiBuffer>,
272        mut range_updates: impl Iterator<Item = InlayHintQuery>,
273        mut currently_shown_hints: HashMap<u64, HashMap<ExcerptId, Vec<(Anchor, InlayId)>>>,
274        cx: &mut ViewContext<Editor>,
275    ) -> Task<anyhow::Result<InlaySplice>> {
276        let conflicts_with_cache = range_updates.any(|update_query| {
277            let Some(cached_buffer_hints) = self.hints_in_buffers.get(&update_query.buffer_id)
278                else { return false };
279            if cached_buffer_hints
280                .buffer_version
281                .changed_since(&update_query.buffer_version)
282            {
283                false
284            } else if update_query
285                .buffer_version
286                .changed_since(&cached_buffer_hints.buffer_version)
287            {
288                true
289            } else {
290                cached_buffer_hints
291                    .hints_per_excerpt
292                    .contains_key(&update_query.excerpt_id)
293            }
294        });
295
296        let queries = filter_queries(range_updates, &self.hints_in_buffers, conflicts_with_cache);
297        let task_multi_buffer = multi_buffer.clone();
298        let fetch_queries_task = fetch_queries(multi_buffer, queries.into_iter(), cx);
299        let mut to_remove = Vec::new();
300        let mut to_insert = Vec::new();
301        cx.spawn(|editor, mut cx| async move {
302            let new_hints = fetch_queries_task.await?;
303            editor.update(&mut cx, |editor, cx| {
304                let multi_buffer_snapshot = task_multi_buffer.read(cx).snapshot(cx);
305                for (new_buffer_id, new_hints_per_buffer) in new_hints {
306                    let cached_buffer_hints = editor
307                        .inlay_hint_cache
308                        .hints_in_buffers
309                        .entry(new_buffer_id)
310                        .or_insert_with(|| {
311                            BufferHints::new(new_hints_per_buffer.buffer_version.clone())
312                        });
313                    let mut shown_buffer_hints = currently_shown_hints
314                        .remove(&new_buffer_id)
315                        .unwrap_or_default();
316                    if cached_buffer_hints
317                        .buffer_version
318                        .changed_since(&new_hints_per_buffer.buffer_version)
319                    {
320                        continue;
321                    } else {
322                        cached_buffer_hints.buffer_version = new_hints_per_buffer.buffer_version;
323                    }
324
325                    for (new_excerpt_id, new_hints_per_excerpt) in
326                        new_hints_per_buffer.hints_per_excerpt
327                    {
328                        let cached_excerpt_hints = cached_buffer_hints
329                            .hints_per_excerpt
330                            .entry(new_excerpt_id)
331                            .or_default();
332                        let shown_excerpt_hints = shown_buffer_hints
333                            .remove(&new_excerpt_id)
334                            .unwrap_or_default();
335
336                        if conflicts_with_cache {
337                            cached_excerpt_hints.cached_excerpt_offsets.clear();
338                            // TODO kb need to add such into to_delete and do not cause extra changes
339                            // cached_excerpt_hints.hints.clear();
340                            // editor.inlay_hint_cache.inlay_hints.clear();
341                            todo!("TODO kb")
342                        } else {
343                            for new_range in new_hints_per_excerpt.cached_excerpt_offsets {
344                                insert_and_merge_ranges(
345                                    &mut cached_excerpt_hints.cached_excerpt_offsets,
346                                    &new_range,
347                                )
348                            }
349                            for new_hint in new_hints_per_excerpt.hints {
350                                let hint_anchor = multi_buffer_snapshot
351                                    .anchor_in_excerpt(new_excerpt_id, new_hint.position);
352                                let insert_ix =
353                                    match cached_excerpt_hints.hints.binary_search_by(|probe| {
354                                        hint_anchor.cmp(&probe.0, &multi_buffer_snapshot)
355                                    }) {
356                                        Ok(ix) | Err(ix) => ix,
357                                    };
358
359                                let new_hint_id = InlayId(post_inc(&mut editor.next_inlay_id));
360                                cached_excerpt_hints
361                                    .hints
362                                    .insert(insert_ix, (hint_anchor, new_hint_id));
363                                editor
364                                    .inlay_hint_cache
365                                    .inlay_hints
366                                    .insert(new_hint_id, new_hint.clone());
367                                if editor
368                                    .inlay_hint_cache
369                                    .allowed_hint_kinds
370                                    .contains(&new_hint.kind)
371                                {
372                                    to_insert.push((new_hint_id, hint_anchor, new_hint));
373                                }
374                            }
375                        }
376
377                        if cached_excerpt_hints.hints.is_empty() {
378                            cached_buffer_hints
379                                .hints_per_excerpt
380                                .remove(&new_excerpt_id);
381                        }
382                    }
383
384                    if shown_buffer_hints.is_empty() {
385                        currently_shown_hints.remove(&new_buffer_id);
386                    }
387                }
388
389                to_remove.extend(
390                    currently_shown_hints
391                        .into_iter()
392                        .flat_map(|(_, hints_by_excerpt)| hints_by_excerpt)
393                        .flat_map(|(_, excerpt_hints)| excerpt_hints)
394                        .map(|(_, hint_id)| hint_id),
395                );
396                InlaySplice {
397                    to_remove,
398                    to_insert,
399                }
400            })
401        })
402    }
403}
404
405fn filter_queries(
406    queries: impl Iterator<Item = InlayHintQuery>,
407    cached_hints: &HashMap<u64, BufferHints<(Anchor, InlayId)>>,
408    invalidate_cache: bool,
409) -> Vec<InlayHintQuery> {
410    queries
411        .filter_map(|query| {
412            let Some(cached_buffer_hints) = cached_hints.get(&query.buffer_id)
413                else { return Some(vec![query]) };
414            if cached_buffer_hints
415                .buffer_version
416                .changed_since(&query.buffer_version)
417            {
418                return None;
419            }
420            let Some(excerpt_hints) = cached_buffer_hints.hints_per_excerpt.get(&query.excerpt_id)
421                else { return Some(vec![query]) };
422
423            if invalidate_cache {
424                Some(vec![query])
425            } else {
426                let non_cached_ranges = missing_subranges(
427                    &excerpt_hints.cached_excerpt_offsets,
428                    &query.excerpt_offset_query_range,
429                );
430                if non_cached_ranges.is_empty() {
431                    None
432                } else {
433                    Some(
434                        non_cached_ranges
435                            .into_iter()
436                            .map(|non_cached_range| InlayHintQuery {
437                                buffer_id: query.buffer_id,
438                                buffer_version: query.buffer_version.clone(),
439                                excerpt_id: query.excerpt_id,
440                                excerpt_offset_query_range: non_cached_range,
441                            })
442                            .collect(),
443                    )
444                }
445            }
446        })
447        .flatten()
448        .collect()
449}
450
451fn allowed_hint_types(
452    inlay_hint_settings: editor_settings::InlayHints,
453) -> HashSet<Option<InlayHintKind>> {
454    let mut new_allowed_hint_types = HashSet::default();
455    if inlay_hint_settings.show_type_hints {
456        new_allowed_hint_types.insert(Some(InlayHintKind::Type));
457    }
458    if inlay_hint_settings.show_parameter_hints {
459        new_allowed_hint_types.insert(Some(InlayHintKind::Parameter));
460    }
461    if inlay_hint_settings.show_other_hints {
462        new_allowed_hint_types.insert(None);
463    }
464    new_allowed_hint_types
465}
466
467fn missing_subranges(cache: &[Range<usize>], input: &Range<usize>) -> Vec<Range<usize>> {
468    let mut missing = Vec::new();
469
470    // Find where the input range would fit in the cache
471    let index = match cache.binary_search_by_key(&input.start, |probe| probe.start) {
472        Ok(pos) | Err(pos) => pos,
473    };
474
475    // Check for a gap from the start of the input range to the first range in the cache
476    if index == 0 {
477        if input.start < cache[index].start {
478            missing.push(input.start..cache[index].start);
479        }
480    } else {
481        let prev_end = cache[index - 1].end;
482        if input.start < prev_end {
483            missing.push(input.start..prev_end);
484        }
485    }
486
487    // Iterate through the cache ranges starting from index
488    for i in index..cache.len() {
489        let start = if i > 0 { cache[i - 1].end } else { input.start };
490        let end = cache[i].start;
491
492        if start < end {
493            missing.push(start..end);
494        }
495    }
496
497    // Check for a gap from the last range in the cache to the end of the input range
498    if let Some(last_range) = cache.last() {
499        if last_range.end < input.end {
500            missing.push(last_range.end..input.end);
501        }
502    } else {
503        // If cache is empty, the entire input range is missing
504        missing.push(input.start..input.end);
505    }
506
507    missing
508}
509
510fn insert_and_merge_ranges(cache: &mut Vec<Range<usize>>, new_range: &Range<usize>) {
511    if cache.is_empty() {
512        cache.push(new_range.clone());
513        return;
514    }
515
516    // Find the index to insert the new range
517    let index = match cache.binary_search_by_key(&new_range.start, |probe| probe.start) {
518        Ok(pos) | Err(pos) => pos,
519    };
520
521    // Check if the new range overlaps with the previous range in the cache
522    if index > 0 && cache[index - 1].end >= new_range.start {
523        // Merge with the previous range
524        cache[index - 1].end = std::cmp::max(cache[index - 1].end, new_range.end);
525    } else {
526        // Insert the new range, as it doesn't overlap with the previous range
527        cache.insert(index, new_range.clone());
528    }
529
530    // Merge overlaps with subsequent ranges
531    let mut i = index;
532    while i + 1 < cache.len() && cache[i].end >= cache[i + 1].start {
533        cache[i].end = std::cmp::max(cache[i].end, cache[i + 1].end);
534        cache.remove(i + 1);
535        i += 1;
536    }
537}
538
539fn fetch_queries<'a, 'b>(
540    multi_buffer: ModelHandle<MultiBuffer>,
541    queries: impl Iterator<Item = InlayHintQuery>,
542    cx: &mut ViewContext<'a, 'b, Editor>,
543) -> Task<anyhow::Result<HashMap<u64, BufferHints<InlayHint>>>> {
544    let mut inlay_fetch_tasks = Vec::new();
545    for query in queries {
546        let task_multi_buffer = multi_buffer.clone();
547        let task = cx.spawn(|editor, mut cx| async move {
548            let Some(buffer_handle) = cx.read(|cx| task_multi_buffer.read(cx).buffer(query.buffer_id))
549                else { return anyhow::Ok((query, Some(Vec::new()))) };
550            let task = editor
551                .update(&mut cx, |editor, cx| {
552                    editor.project.as_ref().map(|project| {
553                        project.update(cx, |project, cx| {
554                            project.query_inlay_hints_for_buffer(
555                                buffer_handle,
556                                query.excerpt_offset_query_range.clone(),
557                                cx,
558                            )
559                        })
560                    })
561                })
562                .context("inlays fecth task spawn")?;
563            Ok((
564                query,
565                match task {
566                    Some(task) => task.await.context("inlays for buffer task")?,
567                    None => Some(Vec::new()),
568                },
569            ))
570        });
571
572        inlay_fetch_tasks.push(task);
573    }
574
575    cx.spawn(|editor, cx| async move {
576        let mut inlay_updates: HashMap<u64, BufferHints<InlayHint>> = HashMap::default();
577        for task_result in futures::future::join_all(inlay_fetch_tasks).await {
578            match task_result {
579                Ok((query, Some(response_hints))) => {
580                    let Some(buffer_snapshot) = editor.read_with(&cx, |editor, cx| {
581                        editor.buffer().read(cx).buffer(query.buffer_id).map(|buffer| buffer.read(cx).snapshot())
582                    })? else { continue; };
583                    let buffer_hints = inlay_updates
584                        .entry(query.buffer_id)
585                        .or_insert_with(|| BufferHints::new(query.buffer_version.clone()));
586                    if buffer_snapshot.version().changed_since(&buffer_hints.buffer_version) {
587                        continue;
588                    }
589                    let cached_excerpt_hints = buffer_hints
590                        .hints_per_excerpt
591                        .entry(query.excerpt_id)
592                        .or_default();
593                    insert_and_merge_ranges(&mut cached_excerpt_hints.cached_excerpt_offsets, &query.excerpt_offset_query_range);
594                    let excerpt_hints = &mut cached_excerpt_hints.hints;
595                    for inlay in response_hints {
596                        match excerpt_hints.binary_search_by(|probe| {
597                            inlay.position.cmp(&probe.position, &buffer_snapshot)
598                        }) {
599                            Ok(ix) | Err(ix) => excerpt_hints.insert(ix, inlay),
600                        }
601                    }
602                }
603                Ok((_, None)) => {}
604                Err(e) => error!("Failed to update inlays for buffer: {e:#}"),
605            }
606        }
607        Ok(inlay_updates)
608    })
609}