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::Point;
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(&self, ranges: &[Range<Point>]) -> impl Iterator<Item = RowChunk> {
94 self.chunks.applicable_chunks(ranges)
95 }
96
97 pub fn cached_hints(&mut self, chunk: &RowChunk) -> Option<&CacheInlayHints> {
98 self.hints_by_chunks[chunk.id].as_ref()
99 }
100
101 pub fn fetched_hints(&mut self, chunk: &RowChunk) -> &mut Option<CacheInlayHintsTask> {
102 &mut self.fetches_by_chunks[chunk.id]
103 }
104
105 #[cfg(any(test, feature = "test-support"))]
106 pub fn all_cached_hints(&self) -> Vec<InlayHint> {
107 self.hints_by_chunks
108 .iter()
109 .filter_map(|hints| hints.as_ref())
110 .flat_map(|hints| hints.values().cloned())
111 .flatten()
112 .map(|(_, hint)| hint)
113 .collect()
114 }
115
116 #[cfg(any(test, feature = "test-support"))]
117 pub fn all_fetched_hints(&self) -> Vec<CacheInlayHintsTask> {
118 self.fetches_by_chunks
119 .iter()
120 .filter_map(|fetches| fetches.clone())
121 .collect()
122 }
123
124 pub fn remove_server_data(&mut self, for_server: LanguageServerId) {
125 for (chunk_index, hints) in self.hints_by_chunks.iter_mut().enumerate() {
126 if let Some(hints) = hints {
127 if hints.remove(&for_server).is_some() {
128 self.fetches_by_chunks[chunk_index] = None;
129 }
130 }
131 }
132 }
133
134 pub fn clear(&mut self) {
135 self.hints_by_chunks = vec![None; self.chunks.len()];
136 self.fetches_by_chunks = vec![None; self.chunks.len()];
137 self.hints_by_id.clear();
138 self.hint_resolves.clear();
139 self.latest_invalidation_requests.clear();
140 }
141
142 pub fn insert_new_hints(
143 &mut self,
144 chunk: RowChunk,
145 server_id: LanguageServerId,
146 new_hints: Vec<(InlayId, InlayHint)>,
147 ) {
148 let existing_hints = self.hints_by_chunks[chunk.id]
149 .get_or_insert_default()
150 .entry(server_id)
151 .or_insert_with(Vec::new);
152 let existing_count = existing_hints.len();
153 existing_hints.extend(new_hints.into_iter().enumerate().filter_map(
154 |(i, (id, new_hint))| {
155 let new_hint_for_id = HintForId {
156 chunk_id: chunk.id,
157 server_id,
158 position: existing_count + i,
159 };
160 if let hash_map::Entry::Vacant(vacant_entry) = self.hints_by_id.entry(id) {
161 vacant_entry.insert(new_hint_for_id);
162 Some((id, new_hint))
163 } else {
164 None
165 }
166 },
167 ));
168 *self.fetched_hints(&chunk) = None;
169 }
170
171 pub fn hint_for_id(&mut self, id: InlayId) -> Option<&mut InlayHint> {
172 let hint_for_id = self.hints_by_id.get(&id)?;
173 let (hint_id, hint) = self
174 .hints_by_chunks
175 .get_mut(hint_for_id.chunk_id)?
176 .as_mut()?
177 .get_mut(&hint_for_id.server_id)?
178 .get_mut(hint_for_id.position)?;
179 debug_assert_eq!(*hint_id, id, "Invalid pointer {hint_for_id:?}");
180 Some(hint)
181 }
182
183 pub(crate) fn invalidate_for_server_refresh(
184 &mut self,
185 for_server: LanguageServerId,
186 request_id: Option<usize>,
187 ) -> bool {
188 match self.latest_invalidation_requests.entry(for_server) {
189 hash_map::Entry::Occupied(mut o) => {
190 if request_id > *o.get() {
191 o.insert(request_id);
192 } else {
193 return false;
194 }
195 }
196 hash_map::Entry::Vacant(v) => {
197 v.insert(request_id);
198 }
199 }
200
201 for (chunk_id, chunk_data) in self.hints_by_chunks.iter_mut().enumerate() {
202 if let Some(removed_hints) = chunk_data
203 .as_mut()
204 .and_then(|chunk_data| chunk_data.remove(&for_server))
205 {
206 for (id, _) in removed_hints {
207 self.hints_by_id.remove(&id);
208 self.hint_resolves.remove(&id);
209 }
210 self.fetches_by_chunks[chunk_id] = None;
211 }
212 }
213
214 true
215 }
216
217 pub(crate) fn invalidate_for_chunk(&mut self, chunk: RowChunk) {
218 self.fetches_by_chunks[chunk.id] = None;
219 if let Some(hints_by_server) = self.hints_by_chunks[chunk.id].take() {
220 for (hint_id, _) in hints_by_server.into_values().flatten() {
221 self.hints_by_id.remove(&hint_id);
222 self.hint_resolves.remove(&hint_id);
223 }
224 }
225 }
226}