inlay_cache.rs

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