inlay_cache.rs

  1use std::{
  2    cmp,
  3    ops::Range,
  4    path::{Path, PathBuf},
  5};
  6
  7use crate::{editor_settings, scroll::ScrollAnchor, Anchor, Editor, ExcerptId, MultiBuffer};
  8use anyhow::Context;
  9use clock::{Global, Local};
 10use gpui::{ModelHandle, Task, ViewContext};
 11use log::error;
 12use project::{InlayHint, InlayHintKind};
 13use util::post_inc;
 14
 15use collections::{hash_map, BTreeMap, HashMap, HashSet};
 16
 17#[derive(Debug, Clone)]
 18pub struct Inlay {
 19    pub id: InlayId,
 20    pub position: Anchor,
 21    pub text: text::Rope,
 22}
 23
 24#[derive(Debug, Clone)]
 25pub struct InlayProperties<T> {
 26    pub position: Anchor,
 27    pub text: T,
 28}
 29
 30#[derive(Debug, Copy, Clone)]
 31pub enum InlayRefreshReason {
 32    SettingsChange(editor_settings::InlayHints),
 33    Scroll(ScrollAnchor),
 34    VisibleExcerptsChange,
 35}
 36
 37#[derive(Debug, Clone, Default)]
 38pub struct InlayCache {
 39    inlays_per_buffer: HashMap<PathBuf, BufferInlays>,
 40    allowed_hint_kinds: HashSet<Option<InlayHintKind>>,
 41    next_inlay_id: usize,
 42}
 43
 44#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
 45pub struct AnchorKey {
 46    offset: usize,
 47    version: Local,
 48}
 49
 50#[derive(Clone, Debug)]
 51pub struct OrderedByAnchorOffset<T>(pub BTreeMap<AnchorKey, (Anchor, T)>);
 52
 53impl<T> OrderedByAnchorOffset<T> {
 54    pub fn add(&mut self, anchor: Anchor, t: T) {
 55        let key = AnchorKey {
 56            offset: anchor.text_anchor.offset,
 57            version: anchor.text_anchor.timestamp,
 58        };
 59        self.0.insert(key, (anchor, t));
 60    }
 61
 62    fn into_ordered_elements(self) -> impl Iterator<Item = (Anchor, T)> {
 63        self.0.into_values()
 64    }
 65
 66    fn ordered_elements(&self) -> impl Iterator<Item = &(Anchor, T)> {
 67        self.0.values()
 68    }
 69}
 70
 71impl<T> Default for OrderedByAnchorOffset<T> {
 72    fn default() -> Self {
 73        Self(BTreeMap::default())
 74    }
 75}
 76
 77#[derive(Clone, Debug, Default)]
 78struct BufferInlays {
 79    buffer_version: Global,
 80    inlays_per_excerpts: HashMap<ExcerptId, OrderedByAnchorOffset<(InlayId, InlayHint)>>,
 81}
 82
 83#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
 84pub struct InlayId(pub usize);
 85
 86#[derive(Debug, Default)]
 87pub struct InlaySplice {
 88    pub to_remove: Vec<InlayId>,
 89    pub to_insert: Vec<(InlayId, Anchor, InlayHint)>,
 90}
 91
 92pub struct InlayHintQuery {
 93    pub buffer_id: u64,
 94    pub buffer_path: PathBuf,
 95    pub buffer_version: Global,
 96    pub excerpt_id: ExcerptId,
 97    pub excerpt_offset_query_range: Range<usize>,
 98}
 99
100impl InlayCache {
101    pub fn new(inlay_hint_settings: editor_settings::InlayHints) -> Self {
102        Self {
103            inlays_per_buffer: HashMap::default(),
104            allowed_hint_kinds: allowed_inlay_hint_types(inlay_hint_settings),
105            next_inlay_id: 0,
106        }
107    }
108
109    pub fn append_inlays(
110        &mut self,
111        multi_buffer: ModelHandle<MultiBuffer>,
112        ranges_to_add: impl Iterator<Item = InlayHintQuery>,
113        cx: &mut ViewContext<Editor>,
114    ) -> Task<anyhow::Result<InlaySplice>> {
115        self.fetch_inlays(multi_buffer, ranges_to_add, false, cx)
116    }
117
118    pub fn replace_inlays(
119        &mut self,
120        multi_buffer: ModelHandle<MultiBuffer>,
121        new_ranges: impl Iterator<Item = InlayHintQuery>,
122        cx: &mut ViewContext<Editor>,
123    ) -> Task<anyhow::Result<InlaySplice>> {
124        self.fetch_inlays(multi_buffer, new_ranges, true, cx)
125    }
126
127    fn fetch_inlays(
128        &mut self,
129        multi_buffer: ModelHandle<MultiBuffer>,
130        inlay_fetch_ranges: impl Iterator<Item = InlayHintQuery>,
131        replace_old: bool,
132        cx: &mut ViewContext<Editor>,
133    ) -> Task<anyhow::Result<InlaySplice>> {
134        let mut inlay_fetch_tasks = Vec::new();
135        for inlay_fetch_range in inlay_fetch_ranges {
136            let inlays_up_to_date = self.inlays_up_to_date(
137                &inlay_fetch_range.buffer_path,
138                &inlay_fetch_range.buffer_version,
139                inlay_fetch_range.excerpt_id,
140            );
141            let task_multi_buffer = multi_buffer.clone();
142            let task = cx.spawn(|editor, mut cx| async move {
143                if inlays_up_to_date {
144                    anyhow::Ok((inlay_fetch_range, None))
145                } else {
146                    let Some(buffer_handle) = cx.read(|cx| task_multi_buffer.read(cx).buffer(inlay_fetch_range.buffer_id))
147                        else { return Ok((inlay_fetch_range, Some(Vec::new()))) };
148                    let task = editor
149                        .update(&mut cx, |editor, cx| {
150                            editor.project.as_ref().map(|project| {
151                                project.update(cx, |project, cx| {
152                                    project.query_inlay_hints_for_buffer(
153                                        buffer_handle,
154                                        inlay_fetch_range.excerpt_offset_query_range.clone(),
155                                        cx,
156                                    )
157                                })
158                            })
159                        })
160                        .context("inlays fecth task spawn")?;
161
162                    Ok((inlay_fetch_range, match task {
163                        Some(task) => task.await.context("inlays for buffer task")?,
164                        None => Some(Vec::new()),
165                    }))
166                }
167            });
168            inlay_fetch_tasks.push(task);
169        }
170
171        let final_task = cx.spawn(|editor, mut cx| async move {
172            let mut inlay_updates: HashMap<
173                PathBuf,
174                (
175                    Global,
176                    HashMap<ExcerptId, Option<(Range<usize>, OrderedByAnchorOffset<InlayHint>)>>,
177                ),
178            > = HashMap::default();
179            let multi_buffer_snapshot =
180                editor.read_with(&cx, |editor, cx| editor.buffer().read(cx).snapshot(cx))?;
181
182            for task_result in futures::future::join_all(inlay_fetch_tasks).await {
183                match task_result {
184                    Ok((inlay_fetch_range, response_inlays)) => {
185                        // TODO kb different caching now
186                        let inlays_per_excerpt = HashMap::from_iter([(
187                            inlay_fetch_range.excerpt_id,
188                            response_inlays
189                                .map(|excerpt_inlays| {
190                                    excerpt_inlays.into_iter().fold(
191                                        OrderedByAnchorOffset::default(),
192                                        |mut ordered_inlays, inlay| {
193                                            let anchor = multi_buffer_snapshot.anchor_in_excerpt(
194                                                inlay_fetch_range.excerpt_id,
195                                                inlay.position,
196                                            );
197                                            ordered_inlays.add(anchor, inlay);
198                                            ordered_inlays
199                                        },
200                                    )
201                                })
202                                .map(|inlays| {
203                                    (inlay_fetch_range.excerpt_offset_query_range, inlays)
204                                }),
205                        )]);
206                        match inlay_updates.entry(inlay_fetch_range.buffer_path) {
207                            hash_map::Entry::Occupied(mut o) => {
208                                o.get_mut().1.extend(inlays_per_excerpt);
209                            }
210                            hash_map::Entry::Vacant(v) => {
211                                v.insert((inlay_fetch_range.buffer_version, inlays_per_excerpt));
212                            }
213                        }
214                    }
215                    Err(e) => error!("Failed to update inlays for buffer: {e:#}"),
216                }
217            }
218
219            let updates = if !inlay_updates.is_empty() {
220                let inlays_update = editor.update(&mut cx, |editor, _| {
221                    editor.inlay_cache.apply_fetch_inlays(inlay_updates)
222                })?;
223                inlays_update
224            } else {
225                InlaySplice::default()
226            };
227
228            anyhow::Ok(updates)
229        });
230
231        final_task
232    }
233
234    fn inlays_up_to_date(
235        &self,
236        buffer_path: &Path,
237        buffer_version: &Global,
238        excerpt_id: ExcerptId,
239    ) -> bool {
240        let Some(buffer_inlays) = self.inlays_per_buffer.get(buffer_path) else { return false };
241        let buffer_up_to_date = buffer_version == &buffer_inlays.buffer_version
242            || buffer_inlays.buffer_version.changed_since(&buffer_version);
243        buffer_up_to_date && buffer_inlays.inlays_per_excerpts.contains_key(&excerpt_id)
244    }
245
246    fn apply_fetch_inlays(
247        &mut self,
248        fetched_inlays: HashMap<
249            PathBuf,
250            (
251                Global,
252                HashMap<ExcerptId, Option<(Range<usize>, OrderedByAnchorOffset<InlayHint>)>>,
253            ),
254        >,
255    ) -> InlaySplice {
256        let mut old_inlays = self.inlays_per_buffer.clone();
257        let mut to_remove = Vec::new();
258        let mut to_insert = Vec::new();
259
260        for (buffer_path, (buffer_version, new_buffer_inlays)) in fetched_inlays {
261            match old_inlays.remove(&buffer_path) {
262                Some(mut old_buffer_inlays) => {
263                    for (excerpt_id, new_excerpt_inlays) in new_buffer_inlays {
264                        let (_, mut new_excerpt_inlays) = match new_excerpt_inlays {
265                            Some((excerpt_offset_range, new_inlays)) => (
266                                excerpt_offset_range,
267                                new_inlays.into_ordered_elements().fuse().peekable(),
268                            ),
269                            None => continue,
270                        };
271                        if self.inlays_up_to_date(&buffer_path, &buffer_version, excerpt_id) {
272                            continue;
273                        }
274
275                        let self_inlays_per_buffer = self
276                            .inlays_per_buffer
277                            .get_mut(&buffer_path)
278                            .expect("element expected: `old_inlays.remove` returned `Some`");
279
280                        if old_buffer_inlays
281                            .inlays_per_excerpts
282                            .remove(&excerpt_id)
283                            .is_some()
284                        {
285                            let self_excerpt_inlays = self_inlays_per_buffer
286                                .inlays_per_excerpts
287                                .get_mut(&excerpt_id)
288                                .expect("element expected: `old_excerpt_inlays` is `Some`");
289                            let mut hints_to_add = Vec::<(Anchor, (InlayId, InlayHint))>::new();
290                            // TODO kb update inner buffer_id and version with the new data?
291                            self_excerpt_inlays.0.retain(
292                                |_, (old_anchor, (old_inlay_id, old_inlay))| {
293                                    let mut retain = false;
294
295                                    while let Some(new_offset) = new_excerpt_inlays
296                                        .peek()
297                                        .map(|(new_anchor, _)| new_anchor.text_anchor.offset)
298                                    {
299                                        let old_offset = old_anchor.text_anchor.offset;
300                                        match new_offset.cmp(&old_offset) {
301                                            cmp::Ordering::Less => {
302                                                let (new_anchor, new_inlay) =
303                                                    new_excerpt_inlays.next().expect(
304                                                        "element expected: `peek` returned `Some`",
305                                                    );
306                                                hints_to_add.push((
307                                                    new_anchor,
308                                                    (
309                                                        InlayId(post_inc(&mut self.next_inlay_id)),
310                                                        new_inlay,
311                                                    ),
312                                                ));
313                                            }
314                                            cmp::Ordering::Equal => {
315                                                let (new_anchor, new_inlay) =
316                                                    new_excerpt_inlays.next().expect(
317                                                        "element expected: `peek` returned `Some`",
318                                                    );
319                                                if &new_inlay == old_inlay {
320                                                    retain = true;
321                                                } else {
322                                                    hints_to_add.push((
323                                                        new_anchor,
324                                                        (
325                                                            InlayId(post_inc(
326                                                                &mut self.next_inlay_id,
327                                                            )),
328                                                            new_inlay,
329                                                        ),
330                                                    ));
331                                                }
332                                            }
333                                            cmp::Ordering::Greater => break,
334                                        }
335                                    }
336
337                                    if !retain {
338                                        to_remove.push(*old_inlay_id);
339                                    }
340                                    retain
341                                },
342                            );
343
344                            for (new_anchor, (id, new_inlay)) in hints_to_add {
345                                self_excerpt_inlays.add(new_anchor, (id, new_inlay.clone()));
346                                to_insert.push((id, new_anchor, new_inlay));
347                            }
348                        }
349
350                        for (new_anchor, new_inlay) in new_excerpt_inlays {
351                            let id = InlayId(post_inc(&mut self.next_inlay_id));
352                            self_inlays_per_buffer
353                                .inlays_per_excerpts
354                                .entry(excerpt_id)
355                                .or_default()
356                                .add(new_anchor, (id, new_inlay.clone()));
357                            to_insert.push((id, new_anchor, new_inlay));
358                        }
359                    }
360                }
361                None => {
362                    let mut inlays_per_excerpts: HashMap<
363                        ExcerptId,
364                        OrderedByAnchorOffset<(InlayId, InlayHint)>,
365                    > = HashMap::default();
366                    for (new_excerpt_id, new_ordered_inlays) in new_buffer_inlays {
367                        if let Some((_, new_ordered_inlays)) = new_ordered_inlays {
368                            for (new_anchor, new_inlay) in
369                                new_ordered_inlays.into_ordered_elements()
370                            {
371                                let id = InlayId(post_inc(&mut self.next_inlay_id));
372                                inlays_per_excerpts
373                                    .entry(new_excerpt_id)
374                                    .or_default()
375                                    .add(new_anchor, (id, new_inlay.clone()));
376                                to_insert.push((id, new_anchor, new_inlay));
377                            }
378                        }
379                    }
380                    self.inlays_per_buffer.insert(
381                        buffer_path,
382                        BufferInlays {
383                            buffer_version,
384                            inlays_per_excerpts,
385                        },
386                    );
387                }
388            }
389        }
390
391        for (_, old_buffer_inlays) in old_inlays {
392            for (_, old_excerpt_inlays) in old_buffer_inlays.inlays_per_excerpts {
393                for (_, (id_to_remove, _)) in old_excerpt_inlays.into_ordered_elements() {
394                    to_remove.push(id_to_remove);
395                }
396            }
397        }
398
399        to_insert.retain(|(_, _, new_hint)| self.allowed_hint_kinds.contains(&new_hint.kind));
400
401        InlaySplice {
402            to_remove,
403            to_insert,
404        }
405    }
406
407    pub fn apply_settings(
408        &mut self,
409        inlay_hint_settings: editor_settings::InlayHints,
410    ) -> InlaySplice {
411        let new_allowed_inlay_hint_types = allowed_inlay_hint_types(inlay_hint_settings);
412
413        let new_allowed_hint_kinds = new_allowed_inlay_hint_types
414            .difference(&self.allowed_hint_kinds)
415            .copied()
416            .collect::<HashSet<_>>();
417        let removed_hint_kinds = self
418            .allowed_hint_kinds
419            .difference(&new_allowed_inlay_hint_types)
420            .collect::<HashSet<_>>();
421        let mut to_remove = Vec::new();
422        let mut to_insert = Vec::new();
423        for (anchor, (inlay_id, inlay_hint)) in self
424            .inlays_per_buffer
425            .iter()
426            .map(|(_, buffer_inlays)| {
427                buffer_inlays
428                    .inlays_per_excerpts
429                    .iter()
430                    .map(|(_, excerpt_inlays)| excerpt_inlays.ordered_elements())
431                    .flatten()
432            })
433            .flatten()
434        {
435            if removed_hint_kinds.contains(&inlay_hint.kind) {
436                to_remove.push(*inlay_id);
437            } else if new_allowed_hint_kinds.contains(&inlay_hint.kind) {
438                to_insert.push((*inlay_id, *anchor, inlay_hint.to_owned()));
439            }
440        }
441
442        self.allowed_hint_kinds = new_allowed_hint_kinds;
443
444        InlaySplice {
445            to_remove,
446            to_insert,
447        }
448    }
449
450    pub fn clear(&mut self) -> Vec<InlayId> {
451        self.inlays_per_buffer
452            .drain()
453            .map(|(_, buffer_inlays)| {
454                buffer_inlays
455                    .inlays_per_excerpts
456                    .into_iter()
457                    .map(|(_, excerpt_inlays)| {
458                        excerpt_inlays
459                            .into_ordered_elements()
460                            .map(|(_, (id, _))| id)
461                    })
462                    .flatten()
463            })
464            .flatten()
465            .collect()
466    }
467}
468
469fn allowed_inlay_hint_types(
470    inlay_hint_settings: editor_settings::InlayHints,
471) -> HashSet<Option<InlayHintKind>> {
472    let mut new_allowed_inlay_hint_types = HashSet::default();
473    if inlay_hint_settings.show_type_hints {
474        new_allowed_inlay_hint_types.insert(Some(InlayHintKind::Type));
475    }
476    if inlay_hint_settings.show_parameter_hints {
477        new_allowed_inlay_hint_types.insert(Some(InlayHintKind::Parameter));
478    }
479    if inlay_hint_settings.show_other_hints {
480        new_allowed_inlay_hint_types.insert(None);
481    }
482    new_allowed_inlay_hint_types
483}