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