1use std::{collections::hash_map, ops::Range, sync::Arc};
2
3use collections::HashMap;
4use futures::future::Shared;
5use gpui::{App, Entity, Task};
6use language::{
7 Buffer,
8 row_chunk::{RowChunk, RowChunks},
9};
10use lsp::LanguageServerId;
11
12use crate::{InlayHint, InlayId};
13
14pub type CacheInlayHints = HashMap<LanguageServerId, Vec<(InlayId, InlayHint)>>;
15pub type CacheInlayHintsTask = Shared<Task<Result<CacheInlayHints, Arc<anyhow::Error>>>>;
16
17/// A logic to apply when querying for new inlay hints and deciding what to do with the old entries in the cache in case of conflicts.
18#[derive(Debug, Clone, Copy)]
19pub enum InvalidationStrategy {
20 /// Language servers reset hints via <a href="https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#workspace_inlayHint_refresh">request</a>.
21 /// Demands to re-query all inlay hints needed and invalidate all cached entries, but does not require instant update with invalidation.
22 ///
23 /// Despite nothing forbids language server from sending this request on every edit, it is expected to be sent only when certain internal server state update, invisible for the editor otherwise.
24 RefreshRequested {
25 server_id: LanguageServerId,
26 request_id: Option<usize>,
27 },
28 /// Multibuffer excerpt(s) and/or singleton buffer(s) were edited at least on one place.
29 /// Neither editor nor LSP is able to tell which open file hints' are not affected, so all of them have to be invalidated, re-queried and do that fast enough to avoid being slow, but also debounce to avoid loading hints on every fast keystroke sequence.
30 BufferEdited,
31 /// A new file got opened/new excerpt was added to a multibuffer/a [multi]buffer was scrolled to a new position.
32 /// No invalidation should be done at all, all new hints are added to the cache.
33 ///
34 /// A special case is the editor toggles and settings change:
35 /// in addition to LSP capabilities, Zed allows omitting certain hint kinds (defined by the corresponding LSP part: type/parameter/other) and toggling hints.
36 /// This does not lead to cache invalidation, but would require cache usage for determining which hints are not displayed and issuing an update to inlays on the screen.
37 None,
38}
39
40impl InvalidationStrategy {
41 pub fn should_invalidate(&self) -> bool {
42 matches!(
43 self,
44 InvalidationStrategy::RefreshRequested { .. } | InvalidationStrategy::BufferEdited
45 )
46 }
47}
48
49pub struct BufferInlayHints {
50 chunks: RowChunks,
51 hints_by_chunks: Vec<Option<CacheInlayHints>>,
52 fetches_by_chunks: Vec<Option<CacheInlayHintsTask>>,
53 hints_by_id: HashMap<InlayId, HintForId>,
54 latest_invalidation_requests: HashMap<LanguageServerId, Option<usize>>,
55 pub(super) hint_resolves: HashMap<InlayId, Shared<Task<()>>>,
56}
57
58#[derive(Debug, Clone, Copy)]
59struct HintForId {
60 chunk_id: usize,
61 server_id: LanguageServerId,
62 position: usize,
63}
64
65impl std::fmt::Debug for BufferInlayHints {
66 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67 f.debug_struct("BufferInlayHints")
68 .field("buffer_chunks", &self.chunks)
69 .field("hints_by_chunks", &self.hints_by_chunks)
70 .field("fetches_by_chunks", &self.fetches_by_chunks)
71 .field("hints_by_id", &self.hints_by_id)
72 .finish_non_exhaustive()
73 }
74}
75
76const MAX_ROWS_IN_A_CHUNK: u32 = 50;
77
78impl BufferInlayHints {
79 pub fn new(buffer: &Entity<Buffer>, cx: &mut App) -> Self {
80 let chunks = RowChunks::new(buffer.read(cx).text_snapshot(), MAX_ROWS_IN_A_CHUNK);
81
82 Self {
83 hints_by_chunks: vec![None; chunks.len()],
84 fetches_by_chunks: vec![None; chunks.len()],
85 latest_invalidation_requests: HashMap::default(),
86 hints_by_id: HashMap::default(),
87 hint_resolves: HashMap::default(),
88 chunks,
89 }
90 }
91
92 pub fn applicable_chunks(
93 &self,
94 ranges: &[Range<text::Anchor>],
95 ) -> impl Iterator<Item = RowChunk> {
96 self.chunks.applicable_chunks(ranges)
97 }
98
99 pub fn cached_hints(&mut self, chunk: &RowChunk) -> Option<&CacheInlayHints> {
100 self.hints_by_chunks[chunk.id].as_ref()
101 }
102
103 pub fn fetched_hints(&mut self, chunk: &RowChunk) -> &mut Option<CacheInlayHintsTask> {
104 &mut self.fetches_by_chunks[chunk.id]
105 }
106
107 #[cfg(any(test, feature = "test-support"))]
108 pub fn all_cached_hints(&self) -> Vec<InlayHint> {
109 self.hints_by_chunks
110 .iter()
111 .filter_map(|hints| hints.as_ref())
112 .flat_map(|hints| hints.values().cloned())
113 .flatten()
114 .map(|(_, hint)| hint)
115 .collect()
116 }
117
118 #[cfg(any(test, feature = "test-support"))]
119 pub fn all_fetched_hints(&self) -> Vec<CacheInlayHintsTask> {
120 self.fetches_by_chunks
121 .iter()
122 .filter_map(|fetches| fetches.clone())
123 .collect()
124 }
125
126 pub fn remove_server_data(&mut self, for_server: LanguageServerId) {
127 for (chunk_index, hints) in self.hints_by_chunks.iter_mut().enumerate() {
128 if let Some(hints) = hints {
129 if hints.remove(&for_server).is_some() {
130 self.fetches_by_chunks[chunk_index] = None;
131 }
132 }
133 }
134 }
135
136 pub fn clear(&mut self) {
137 self.hints_by_chunks = vec![None; self.chunks.len()];
138 self.fetches_by_chunks = vec![None; self.chunks.len()];
139 self.hints_by_id.clear();
140 self.hint_resolves.clear();
141 self.latest_invalidation_requests.clear();
142 }
143
144 pub fn insert_new_hints(
145 &mut self,
146 chunk: RowChunk,
147 server_id: LanguageServerId,
148 new_hints: Vec<(InlayId, InlayHint)>,
149 ) {
150 let existing_hints = self.hints_by_chunks[chunk.id]
151 .get_or_insert_default()
152 .entry(server_id)
153 .or_insert_with(Vec::new);
154 let existing_count = existing_hints.len();
155 existing_hints.extend(new_hints.into_iter().enumerate().filter_map(
156 |(i, (id, new_hint))| {
157 let new_hint_for_id = HintForId {
158 chunk_id: chunk.id,
159 server_id,
160 position: existing_count + i,
161 };
162 if let hash_map::Entry::Vacant(vacant_entry) = self.hints_by_id.entry(id) {
163 vacant_entry.insert(new_hint_for_id);
164 Some((id, new_hint))
165 } else {
166 None
167 }
168 },
169 ));
170 *self.fetched_hints(&chunk) = None;
171 }
172
173 pub fn hint_for_id(&mut self, id: InlayId) -> Option<&mut InlayHint> {
174 let hint_for_id = self.hints_by_id.get(&id)?;
175 let (hint_id, hint) = self
176 .hints_by_chunks
177 .get_mut(hint_for_id.chunk_id)?
178 .as_mut()?
179 .get_mut(&hint_for_id.server_id)?
180 .get_mut(hint_for_id.position)?;
181 debug_assert_eq!(*hint_id, id, "Invalid pointer {hint_for_id:?}");
182 Some(hint)
183 }
184
185 pub fn buffer_chunks_len(&self) -> usize {
186 self.chunks.len()
187 }
188
189 pub(crate) fn invalidate_for_server_refresh(
190 &mut self,
191 for_server: LanguageServerId,
192 request_id: Option<usize>,
193 ) -> bool {
194 match self.latest_invalidation_requests.entry(for_server) {
195 hash_map::Entry::Occupied(mut o) => {
196 if request_id > *o.get() {
197 o.insert(request_id);
198 } else {
199 return false;
200 }
201 }
202 hash_map::Entry::Vacant(v) => {
203 v.insert(request_id);
204 }
205 }
206
207 for (chunk_id, chunk_data) in self.hints_by_chunks.iter_mut().enumerate() {
208 if let Some(removed_hints) = chunk_data
209 .as_mut()
210 .and_then(|chunk_data| chunk_data.remove(&for_server))
211 {
212 for (id, _) in removed_hints {
213 self.hints_by_id.remove(&id);
214 self.hint_resolves.remove(&id);
215 }
216 self.fetches_by_chunks[chunk_id] = None;
217 }
218 }
219
220 true
221 }
222
223 pub(crate) fn invalidate_for_chunk(&mut self, chunk: RowChunk) {
224 self.fetches_by_chunks[chunk.id] = None;
225 if let Some(hints_by_server) = self.hints_by_chunks[chunk.id].take() {
226 for (hint_id, _) in hints_by_server.into_values().flatten() {
227 self.hints_by_id.remove(&hint_id);
228 self.hint_resolves.remove(&hint_id);
229 }
230 }
231 }
232}