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