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}