inlay_hint_cache.rs

  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}