Draft a shared inlay cache prototype

co-authored-by: Antonio Scandurra <antonio@zed.dev>

crates/project/src/lsp_store.rs

@@ -1,4 +1,5 @@
 pub mod clangd_ext;
+mod inlay_hint_cache;
 pub mod lsp_ext_command;
 pub mod rust_analyzer_ext;
 
@@ -10,7 +11,7 @@ use crate::{
     buffer_store::{BufferStore, BufferStoreEvent},
     environment::ProjectEnvironment,
     lsp_command::{self, *},
-    lsp_store,
+    lsp_store::{self, inlay_hint_cache::InlayHintCache},
     manifest_tree::{
         AdapterQuery, LanguageServerTree, LanguageServerTreeNode, LaunchDisposition,
         ManifestQueryDelegate, ManifestTree,
@@ -3551,7 +3552,8 @@ pub struct LspStore {
     _maintain_buffer_languages: Task<()>,
     diagnostic_summaries:
         HashMap<WorktreeId, HashMap<Arc<Path>, HashMap<LanguageServerId, DiagnosticSummary>>>,
-    lsp_data: HashMap<BufferId, DocumentColorData>,
+    color_data: HashMap<BufferId, LspDocumentColorData>,
+    inlay_hints: HashMap<BufferId, LspInlayHintData>,
 }
 
 #[derive(Debug, Default, Clone)]
@@ -3563,13 +3565,19 @@ pub struct DocumentColors {
 type DocumentColorTask = Shared<Task<std::result::Result<DocumentColors, Arc<anyhow::Error>>>>;
 
 #[derive(Debug, Default)]
-struct DocumentColorData {
+struct LspDocumentColorData {
     colors_for_version: Global,
     colors: HashMap<LanguageServerId, HashSet<DocumentColor>>,
     cache_version: usize,
     colors_update: Option<(Global, DocumentColorTask)>,
 }
 
+#[derive(Debug, Default)]
+struct LspInlayHintData {
+    hints_for_version: Global,
+    hints: HashMap<LanguageServerId, InlayHintCache>,
+}
+
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub enum ColorFetchStrategy {
     IgnoreCache,
@@ -3804,7 +3812,8 @@ impl LspStore {
             language_server_statuses: Default::default(),
             nonce: StdRng::from_entropy().r#gen(),
             diagnostic_summaries: HashMap::default(),
-            lsp_data: HashMap::default(),
+            color_data: HashMap::default(),
+            inlay_hints: HashMap::default(),
             active_entry: None,
             _maintain_workspace_config,
             _maintain_buffer_languages: Self::maintain_buffer_languages(languages, cx),
@@ -3861,7 +3870,8 @@ impl LspStore {
             language_server_statuses: Default::default(),
             nonce: StdRng::from_entropy().r#gen(),
             diagnostic_summaries: HashMap::default(),
-            lsp_data: HashMap::default(),
+            color_data: HashMap::default(),
+            inlay_hints: HashMap::default(),
             active_entry: None,
             toolchain_store,
             _maintain_workspace_config,
@@ -4162,7 +4172,8 @@ impl LspStore {
                         *refcount
                     };
                     if refcount == 0 {
-                        lsp_store.lsp_data.remove(&buffer_id);
+                        lsp_store.color_data.remove(&buffer_id);
+                        lsp_store.inlay_hints.remove(&buffer_id);
                         let local = lsp_store.as_local_mut().unwrap();
                         local.registered_buffers.remove(&buffer_id);
                         local.buffers_opened_in_servers.remove(&buffer_id);
@@ -6609,7 +6620,7 @@ impl LspStore {
             ColorFetchStrategy::UseCache {
                 known_cache_version,
             } => {
-                if let Some(cached_data) = self.lsp_data.get(&buffer_id) {
+                if let Some(cached_data) = self.color_data.get(&buffer_id) {
                     if !version_queried_for.changed_since(&cached_data.colors_for_version) {
                         let has_different_servers = self.as_local().is_some_and(|local| {
                             local
@@ -6642,7 +6653,7 @@ impl LspStore {
             }
         }
 
-        let lsp_data = self.lsp_data.entry(buffer_id).or_default();
+        let lsp_data = self.color_data.entry(buffer_id).or_default();
         if let Some((updating_for, running_update)) = &lsp_data.colors_update {
             if !version_queried_for.changed_since(&updating_for) {
                 return Some(running_update.clone());
@@ -6679,7 +6690,7 @@ impl LspStore {
                         lsp_store
                             .update(cx, |lsp_store, _| {
                                 lsp_store
-                                    .lsp_data
+                                    .color_data
                                     .entry(buffer_id)
                                     .or_default()
                                     .colors_update = None;
@@ -6691,7 +6702,7 @@ impl LspStore {
 
                 lsp_store
                     .update(cx, |lsp_store, _| {
-                        let lsp_data = lsp_store.lsp_data.entry(buffer_id).or_default();
+                        let lsp_data = lsp_store.color_data.entry(buffer_id).or_default();
 
                         if lsp_data.colors_for_version == query_version_queried_for {
                             lsp_data.colors.extend(fetched_colors.clone());
@@ -11245,9 +11256,12 @@ impl LspStore {
     }
 
     fn cleanup_lsp_data(&mut self, for_server: LanguageServerId) {
-        for buffer_lsp_data in self.lsp_data.values_mut() {
+        for color_data in self.color_data.values_mut() {
-            buffer_lsp_data.colors.remove(&for_server);
+            color_data.colors.remove(&for_server);
-            buffer_lsp_data.cache_version += 1;
+            color_data.cache_version += 1;
+        }
+        for hint_data in self.inlay_hints.values_mut() {
+            hint_data.hints.remove(&for_server);
         }
         if let Some(local) = self.as_local_mut() {
             local.buffer_pull_diagnostics_result_ids.remove(&for_server);

crates/project/src/lsp_store/inlay_hint_cache.rs

@@ -0,0 +1,107 @@
+// Every initial and invalidated state:
+// [unresolved(0..len)]
+//
+// [pending(0..50), unresolved(51..len)]
+//
+// [hint(43), unresolved(51..len)]
+// OR
+// [unresolved(51..len)]
+
+// [pending(0..50), unresolved(60..65), hint(70), hint(70), hint(90)]
+//
+// [pending(0..50), hint(70), hint(70), hint(90)]
+
+use std::{ops::Range, sync::Arc};
+
+use clock::Global;
+use futures::future::Shared;
+use gpui::Task;
+use sum_tree::SumTree;
+use text::{Anchor, Rope};
+
+#[derive(Debug, Clone)]
+enum LspInlayHintCacheItem {
+    Unresolved {
+        range: Range<Anchor>,
+        attempts: usize,
+    },
+    Pending {
+        range: Range<Anchor>,
+        version: usize,
+    },
+    InlayHint {
+        hint: InlayHint,
+        version: usize,
+    },
+}
+
+#[derive(Debug, Clone)]
+struct Summary {
+    // TODO kb
+}
+
+impl sum_tree::Summary for Summary {
+    type Context = ();
+
+    fn zero(cx: &Self::Context) -> Self {
+        todo!()
+    }
+
+    fn add_summary(&mut self, summary: &Self, cx: &Self::Context) {
+        todo!()
+    }
+}
+
+impl sum_tree::Item for LspInlayHintCacheItem {
+    type Summary = Summary;
+
+    fn summary(&self, cx: &<Self::Summary as sum_tree::Summary>::Context) -> Self::Summary {
+        todo!("TODO kb")
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+struct InlayHintId(usize);
+
+#[derive(Debug, Clone)]
+struct InlayHint {
+    pub id: InlayHintId,
+    pub position: Anchor,
+    pub text: Rope,
+}
+
+// TODO kb wrong: we have to pull by ranges
+type InlayHintsTask = Shared<Task<std::result::Result<Vec<InlayHint>, Arc<anyhow::Error>>>>;
+
+#[derive(Debug)]
+pub struct InlayHintCache {
+    // TODO kb is it needed? What about the inlay hint data, should there be a version too?
+    cache_version: usize,
+    hints_update: Option<(Global, InlayHintsTask)>,
+    items: SumTree<LspInlayHintCacheItem>,
+}
+
+impl InlayHintCache {
+    /// Invalidate this cache. This will keep previously cached results until a
+    /// call to `refresh` is made.
+    pub fn invalidate(&mut self) {}
+
+    /// Editor calls this every time when a viewport changes.
+    pub fn refresh(&mut self, range: Range<usize>) {
+        todo!()
+    }
+
+    /// Editor has to use this to keep its inlay may up-to-date,
+    /// this is done once on editor instantiation for the initial inlay splice.
+    ///
+    /// The rest is retrieved via the updates.
+    pub fn query(&self, range: Range<usize>) -> impl Iterator<Item = InlayHint> {
+        let output: Vec<InlayHint> = todo!();
+        output.into_iter()
+    }
+}
+
+enum InlayHintsChanged {
+    Added(Vec<InlayHint>),
+    Removed(Vec<InlayHintId>),
+}