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