1use std::{collections::hash_map, ops::Range, sync::Arc};
2
3use anyhow::{Context as _, Result};
4use collections::HashMap;
5use futures::future::Shared;
6use gpui::{App, AppContext as _, AsyncApp, Context, Entity, Task};
7use language::{
8 Buffer,
9 row_chunk::{RowChunk, RowChunks},
10};
11use lsp::LanguageServerId;
12use rpc::{TypedEnvelope, proto};
13use settings::Settings as _;
14use text::{BufferId, Point};
15
16use crate::{
17 InlayHint, InlayId, LspStore, LspStoreEvent, ResolveState, lsp_command::InlayHints,
18 project_settings::ProjectSettings,
19};
20
21pub type CacheInlayHints = HashMap<LanguageServerId, Vec<(InlayId, InlayHint)>>;
22pub type CacheInlayHintsTask = Shared<Task<Result<CacheInlayHints, Arc<anyhow::Error>>>>;
23
24/// 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.
25#[derive(Debug, Clone, Copy)]
26pub enum InvalidationStrategy {
27 /// Language servers reset hints via <a href="https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#workspace_inlayHint_refresh">request</a>.
28 /// Demands to re-query all inlay hints needed and invalidate all cached entries, but does not require instant update with invalidation.
29 ///
30 /// 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.
31 RefreshRequested {
32 server_id: LanguageServerId,
33 request_id: Option<usize>,
34 },
35 /// Multibuffer excerpt(s) and/or singleton buffer(s) were edited at least on one place.
36 /// 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.
37 BufferEdited,
38 /// A new file got opened/new excerpt was added to a multibuffer/a [multi]buffer was scrolled to a new position.
39 /// No invalidation should be done at all, all new hints are added to the cache.
40 ///
41 /// A special case is the editor toggles and settings change:
42 /// in addition to LSP capabilities, Zed allows omitting certain hint kinds (defined by the corresponding LSP part: type/parameter/other) and toggling hints.
43 /// 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.
44 None,
45}
46
47impl InvalidationStrategy {
48 pub fn should_invalidate(&self) -> bool {
49 matches!(
50 self,
51 InvalidationStrategy::RefreshRequested { .. } | InvalidationStrategy::BufferEdited
52 )
53 }
54}
55
56pub struct BufferInlayHints {
57 chunks: RowChunks,
58 hints_by_chunks: Vec<Option<CacheInlayHints>>,
59 fetches_by_chunks: Vec<Option<CacheInlayHintsTask>>,
60 hints_by_id: HashMap<InlayId, HintForId>,
61 latest_invalidation_requests: HashMap<LanguageServerId, Option<usize>>,
62 pub(super) hint_resolves: HashMap<InlayId, Shared<Task<()>>>,
63}
64
65#[derive(Debug, Clone, Copy)]
66struct HintForId {
67 chunk_id: usize,
68 server_id: LanguageServerId,
69 position: usize,
70}
71
72impl std::fmt::Debug for BufferInlayHints {
73 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74 f.debug_struct("BufferInlayHints")
75 .field("buffer_chunks", &self.chunks)
76 .field("hints_by_chunks", &self.hints_by_chunks)
77 .field("fetches_by_chunks", &self.fetches_by_chunks)
78 .field("hints_by_id", &self.hints_by_id)
79 .finish_non_exhaustive()
80 }
81}
82
83const MAX_ROWS_IN_A_CHUNK: u32 = 50;
84
85impl BufferInlayHints {
86 pub fn new(buffer: &Entity<Buffer>, cx: &mut App) -> Self {
87 let chunks = RowChunks::new(buffer.read(cx).as_text_snapshot(), MAX_ROWS_IN_A_CHUNK);
88
89 Self {
90 hints_by_chunks: vec![None; chunks.len()],
91 fetches_by_chunks: vec![None; chunks.len()],
92 latest_invalidation_requests: HashMap::default(),
93 hints_by_id: HashMap::default(),
94 hint_resolves: HashMap::default(),
95 chunks,
96 }
97 }
98
99 pub fn applicable_chunks(&self, ranges: &[Range<Point>]) -> impl Iterator<Item = RowChunk> {
100 self.chunks.applicable_chunks(ranges)
101 }
102
103 pub fn cached_hints(&mut self, chunk: &RowChunk) -> Option<&CacheInlayHints> {
104 self.hints_by_chunks[chunk.id].as_ref()
105 }
106
107 pub fn fetched_hints(&mut self, chunk: &RowChunk) -> &mut Option<CacheInlayHintsTask> {
108 &mut self.fetches_by_chunks[chunk.id]
109 }
110
111 #[cfg(any(test, feature = "test-support"))]
112 pub fn all_cached_hints(&self) -> Vec<InlayHint> {
113 self.hints_by_chunks
114 .iter()
115 .filter_map(|hints| hints.as_ref())
116 .flat_map(|hints| hints.values().cloned())
117 .flatten()
118 .map(|(_, hint)| hint)
119 .collect()
120 }
121
122 #[cfg(any(test, feature = "test-support"))]
123 pub fn all_fetched_hints(&self) -> Vec<CacheInlayHintsTask> {
124 self.fetches_by_chunks
125 .iter()
126 .filter_map(|fetches| fetches.clone())
127 .collect()
128 }
129
130 pub fn remove_server_data(&mut self, for_server: LanguageServerId) {
131 for (chunk_index, hints) in self.hints_by_chunks.iter_mut().enumerate() {
132 if let Some(hints) = hints {
133 if hints.remove(&for_server).is_some() {
134 self.fetches_by_chunks[chunk_index] = None;
135 }
136 }
137 }
138 }
139
140 pub fn clear(&mut self) {
141 self.hints_by_chunks = vec![None; self.chunks.len()];
142 self.fetches_by_chunks = vec![None; self.chunks.len()];
143 self.hints_by_id.clear();
144 self.hint_resolves.clear();
145 self.latest_invalidation_requests.clear();
146 }
147
148 pub fn insert_new_hints(
149 &mut self,
150 chunk: RowChunk,
151 server_id: LanguageServerId,
152 new_hints: Vec<(InlayId, InlayHint)>,
153 ) {
154 let existing_hints = self.hints_by_chunks[chunk.id]
155 .get_or_insert_default()
156 .entry(server_id)
157 .or_insert_with(Vec::new);
158 let existing_count = existing_hints.len();
159 existing_hints.extend(new_hints.into_iter().enumerate().filter_map(
160 |(i, (id, new_hint))| {
161 let new_hint_for_id = HintForId {
162 chunk_id: chunk.id,
163 server_id,
164 position: existing_count + i,
165 };
166 if let hash_map::Entry::Vacant(vacant_entry) = self.hints_by_id.entry(id) {
167 vacant_entry.insert(new_hint_for_id);
168 Some((id, new_hint))
169 } else {
170 None
171 }
172 },
173 ));
174 *self.fetched_hints(&chunk) = None;
175 }
176
177 pub fn hint_for_id(&mut self, id: InlayId) -> Option<&mut InlayHint> {
178 let hint_for_id = self.hints_by_id.get(&id)?;
179 let (hint_id, hint) = self
180 .hints_by_chunks
181 .get_mut(hint_for_id.chunk_id)?
182 .as_mut()?
183 .get_mut(&hint_for_id.server_id)?
184 .get_mut(hint_for_id.position)?;
185 debug_assert_eq!(*hint_id, id, "Invalid pointer {hint_for_id:?}");
186 Some(hint)
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}
233
234impl LspStore {
235 pub(super) fn resolve_inlay_hint(
236 &self,
237 mut hint: InlayHint,
238 buffer: Entity<Buffer>,
239 server_id: LanguageServerId,
240 cx: &mut Context<Self>,
241 ) -> Task<anyhow::Result<InlayHint>> {
242 if let Some((upstream_client, project_id)) = self.upstream_client() {
243 if !self.check_if_capable_for_proto_request(&buffer, InlayHints::can_resolve_inlays, cx)
244 {
245 hint.resolve_state = ResolveState::Resolved;
246 return Task::ready(Ok(hint));
247 }
248 let request = proto::ResolveInlayHint {
249 project_id,
250 buffer_id: buffer.read(cx).remote_id().into(),
251 language_server_id: server_id.0 as u64,
252 hint: Some(InlayHints::project_to_proto_hint(hint.clone())),
253 };
254 cx.background_spawn(async move {
255 let response = upstream_client
256 .request(request)
257 .await
258 .context("inlay hints proto request")?;
259 match response.hint {
260 Some(resolved_hint) => InlayHints::proto_to_project_hint(resolved_hint)
261 .context("inlay hints proto resolve response conversion"),
262 None => Ok(hint),
263 }
264 })
265 } else {
266 let Some(lang_server) = buffer.update(cx, |buffer, cx| {
267 self.language_server_for_local_buffer(buffer, server_id, cx)
268 .map(|(_, server)| server.clone())
269 }) else {
270 return Task::ready(Ok(hint));
271 };
272 if !InlayHints::can_resolve_inlays(&lang_server.capabilities()) {
273 return Task::ready(Ok(hint));
274 }
275 let buffer_snapshot = buffer.read(cx).snapshot();
276 let request_timeout = ProjectSettings::get_global(cx)
277 .global_lsp_settings
278 .get_request_timeout();
279 cx.spawn(async move |_, cx| {
280 let resolve_task = lang_server.request::<lsp::request::InlayHintResolveRequest>(
281 InlayHints::project_to_lsp_hint(hint, &buffer_snapshot),
282 request_timeout,
283 );
284 let resolved_hint = resolve_task
285 .await
286 .into_response()
287 .context("inlay hint resolve LSP request")?;
288 let resolved_hint = InlayHints::lsp_to_project_hint(
289 resolved_hint,
290 &buffer,
291 server_id,
292 ResolveState::Resolved,
293 false,
294 cx,
295 )
296 .await?;
297 Ok(resolved_hint)
298 })
299 }
300 }
301
302 pub(super) async fn handle_refresh_inlay_hints(
303 lsp_store: Entity<Self>,
304 envelope: TypedEnvelope<proto::RefreshInlayHints>,
305 mut cx: AsyncApp,
306 ) -> Result<proto::Ack> {
307 lsp_store.update(&mut cx, |_, cx| {
308 cx.emit(LspStoreEvent::RefreshInlayHints {
309 server_id: LanguageServerId::from_proto(envelope.payload.server_id),
310 request_id: envelope.payload.request_id.map(|id| id as usize),
311 });
312 });
313 Ok(proto::Ack {})
314 }
315
316 pub(super) async fn handle_resolve_inlay_hint(
317 lsp_store: Entity<Self>,
318 envelope: TypedEnvelope<proto::ResolveInlayHint>,
319 mut cx: AsyncApp,
320 ) -> Result<proto::ResolveInlayHintResponse> {
321 let proto_hint = envelope
322 .payload
323 .hint
324 .expect("incorrect protobuf resolve inlay hint message: missing the inlay hint");
325 let hint = InlayHints::proto_to_project_hint(proto_hint)
326 .context("resolved proto inlay hint conversion")?;
327 let buffer = lsp_store.update(&mut cx, |lsp_store, cx| {
328 let buffer_id = BufferId::new(envelope.payload.buffer_id)?;
329 lsp_store.buffer_store.read(cx).get_existing(buffer_id)
330 })?;
331 let response_hint = lsp_store
332 .update(&mut cx, |lsp_store, cx| {
333 lsp_store.resolve_inlay_hint(
334 hint,
335 buffer,
336 LanguageServerId(envelope.payload.language_server_id as usize),
337 cx,
338 )
339 })
340 .await
341 .context("inlay hints fetch")?;
342 Ok(proto::ResolveInlayHintResponse {
343 hint: Some(InlayHints::project_to_proto_hint(response_hint)),
344 })
345 }
346}