Rework inlay hints system (zed-industries#40183)

Closes zed-industries#40047
Closes zed-industries#24798
Closes zed-industries#24788

Before, each editor, even if it's the same buffer split in 2, was
querying for inlay hints separately, and storing the whole inlay hint
twice, in `Editor`'s `display_map` and its `inlay_hint_cache` fields.

Now, instead of `inlay_hint_cache`, each editor maintains a minimal set
of metadata (which area was queried by what task) instead, and all LSP
inlay hint data had been moved into `LspStore`, both local and remote
flavors store the data.
This allows Zed, as long as a buffer is open, to reuse the inlay hint
data similar to how document colors and code lens are now stored and
reused.

Unlike other reused LSP data, inlay hints data is the first one that's
possible to query by document ranges and previous version had issue with
caching and invalidating such ranges already queried for.
The new version re-approaches this by chunking the file into row ranges,
which are queried based on the editors' visible area.

Among the corresponding refactoring, one notable difference in inlays
display are multi buffers: buffers in them are not
[registered](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_didOpen)
in the language server until a caret/selection is placed inside their
excerpts inside the multi buffer.

New inlays code does not query language servers for unregistered
buffers, as servers usually respond with empty responses or errors in
such cases.

Release Notes:

- Reworked inlay hints to be less error-prone

---------

Co-authored-by: Lukas Wirth <[email protected]>
Co-authored-by: dino <[email protected]>
Co-authored-by: Lukas Wirth <[email protected]>

Author: 3 people

Cargo.lock

@@ -11,10 +11,10 @@ use assistant_slash_commands::codeblock_fence_for_path;
 use collections::{HashMap, HashSet};
 use editor::{
     Addon, Anchor, AnchorRangeExt, ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement,
-    EditorEvent, EditorMode, EditorSnapshot, EditorStyle, ExcerptId, FoldPlaceholder, InlayId,
+    EditorEvent, EditorMode, EditorSnapshot, EditorStyle, ExcerptId, FoldPlaceholder, Inlay,
     MultiBuffer, ToOffset,
     actions::Paste,
-    display_map::{Crease, CreaseId, FoldId, Inlay},
+    display_map::{Crease, CreaseId, FoldId},
 };
 use futures::{
     FutureExt as _,
@@ -29,7 +29,8 @@ use language::{Buffer, Language, language_settings::InlayHintKind};
 use language_model::LanguageModelImage;
 use postage::stream::Stream as _;
 use project::{
-    CompletionIntent, InlayHint, InlayHintLabel, Project, ProjectItem, ProjectPath, Worktree,
+    CompletionIntent, InlayHint, InlayHintLabel, InlayId, Project, ProjectItem, ProjectPath,
+    Worktree,
 };
 use prompt_store::{PromptId, PromptStore};
 use rope::Point;
@@ -75,7 +76,7 @@ pub enum MessageEditorEvent {
 
 impl EventEmitter<MessageEditorEvent> for MessageEditor {}
 
-const COMMAND_HINT_INLAY_ID: u32 = 0;
+const COMMAND_HINT_INLAY_ID: InlayId = InlayId::Hint(0);
 
 impl MessageEditor {
     pub fn new(
@@ -151,7 +152,7 @@ impl MessageEditor {
                         let has_new_hint = !new_hints.is_empty();
                         editor.splice_inlays(
                             if has_hint {
-                                &[InlayId::Hint(COMMAND_HINT_INLAY_ID)]
+                                &[COMMAND_HINT_INLAY_ID]
                             } else {
                                 &[]
                             },

crates/agent_ui/src/acp/message_editor.rs

@@ -343,7 +343,6 @@ impl Server {
             .add_request_handler(forward_read_only_project_request::<proto::OpenBufferForSymbol>)
             .add_request_handler(forward_read_only_project_request::<proto::OpenBufferById>)
             .add_request_handler(forward_read_only_project_request::<proto::SynchronizeBuffers>)
-            .add_request_handler(forward_read_only_project_request::<proto::InlayHints>)
             .add_request_handler(forward_read_only_project_request::<proto::ResolveInlayHint>)
             .add_request_handler(forward_read_only_project_request::<proto::GetColorPresentation>)
             .add_request_handler(forward_read_only_project_request::<proto::OpenBufferByPath>)

crates/collab/src/rpc.rs

@@ -1849,10 +1849,40 @@ async fn test_mutual_editor_inlay_hint_cache_update(
         ..lsp::ServerCapabilities::default()
     };
     client_a.language_registry().add(rust_lang());
+
+    // Set up the language server to return an additional inlay hint on each request.
+    let edits_made = Arc::new(AtomicUsize::new(0));
+    let closure_edits_made = Arc::clone(&edits_made);
     let mut fake_language_servers = client_a.language_registry().register_fake_lsp(
         "Rust",
         FakeLspAdapter {
             capabilities: capabilities.clone(),
+            initializer: Some(Box::new(move |fake_language_server| {
+                let closure_edits_made = closure_edits_made.clone();
+                fake_language_server.set_request_handler::<lsp::request::InlayHintRequest, _, _>(
+                    move |params, _| {
+                        let edits_made_2 = Arc::clone(&closure_edits_made);
+                        async move {
+                            assert_eq!(
+                                params.text_document.uri,
+                                lsp::Uri::from_file_path(path!("/a/main.rs")).unwrap(),
+                            );
+                            let edits_made =
+                                AtomicUsize::load(&edits_made_2, atomic::Ordering::Acquire);
+                            Ok(Some(vec![lsp::InlayHint {
+                                position: lsp::Position::new(0, edits_made as u32),
+                                label: lsp::InlayHintLabel::String(edits_made.to_string()),
+                                kind: None,
+                                text_edits: None,
+                                tooltip: None,
+                                padding_left: None,
+                                padding_right: None,
+                                data: None,
+                            }]))
+                        }
+                    },
+                );
+            })),
             ..FakeLspAdapter::default()
         },
     );
@@ -1894,61 +1924,20 @@ async fn test_mutual_editor_inlay_hint_cache_update(
         .unwrap();
 
     let (workspace_a, cx_a) = client_a.build_workspace(&project_a, cx_a);
-    executor.start_waiting();
 
     // The host opens a rust file.
-    let _buffer_a = project_a
-        .update(cx_a, |project, cx| {
-            project.open_local_buffer(path!("/a/main.rs"), cx)
-        })
-        .await
-        .unwrap();
-    let editor_a = workspace_a
-        .update_in(cx_a, |workspace, window, cx| {
-            workspace.open_path((worktree_id, rel_path("main.rs")), None, true, window, cx)
-        })
-        .await
-        .unwrap()
-        .downcast::<Editor>()
-        .unwrap();
-
+    let file_a = workspace_a.update_in(cx_a, |workspace, window, cx| {
+        workspace.open_path((worktree_id, rel_path("main.rs")), None, true, window, cx)
+    });
     let fake_language_server = fake_language_servers.next().await.unwrap();
-
-    // Set up the language server to return an additional inlay hint on each request.
-    let edits_made = Arc::new(AtomicUsize::new(0));
-    let closure_edits_made = Arc::clone(&edits_made);
-    fake_language_server
-        .set_request_handler::<lsp::request::InlayHintRequest, _, _>(move |params, _| {
-            let edits_made_2 = Arc::clone(&closure_edits_made);
-            async move {
-                assert_eq!(
-                    params.text_document.uri,
-                    lsp::Uri::from_file_path(path!("/a/main.rs")).unwrap(),
-                );
-                let edits_made = AtomicUsize::load(&edits_made_2, atomic::Ordering::Acquire);
-                Ok(Some(vec![lsp::InlayHint {
-                    position: lsp::Position::new(0, edits_made as u32),
-                    label: lsp::InlayHintLabel::String(edits_made.to_string()),
-                    kind: None,
-                    text_edits: None,
-                    tooltip: None,
-                    padding_left: None,
-                    padding_right: None,
-                    data: None,
-                }]))
-            }
-        })
-        .next()
-        .await
-        .unwrap();
-
+    let editor_a = file_a.await.unwrap().downcast::<Editor>().unwrap();
     executor.run_until_parked();
 
     let initial_edit = edits_made.load(atomic::Ordering::Acquire);
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert_eq!(
             vec![initial_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Host should get its first hints when opens an editor"
         );
     });
@@ -1963,10 +1952,10 @@ async fn test_mutual_editor_inlay_hint_cache_update(
         .unwrap();
 
     executor.run_until_parked();
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec![initial_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Client should get its first hints when opens an editor"
         );
     });
@@ -1981,16 +1970,16 @@ async fn test_mutual_editor_inlay_hint_cache_update(
     cx_b.focus(&editor_b);
 
     executor.run_until_parked();
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert_eq!(
             vec![after_client_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
         );
     });
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec![after_client_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
         );
     });
 
@@ -2004,16 +1993,16 @@ async fn test_mutual_editor_inlay_hint_cache_update(
     cx_a.focus(&editor_a);
 
     executor.run_until_parked();
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert_eq!(
             vec![after_host_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
         );
     });
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec![after_host_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
         );
     });
 
@@ -2025,26 +2014,22 @@ async fn test_mutual_editor_inlay_hint_cache_update(
         .expect("inlay refresh request failed");
 
     executor.run_until_parked();
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert_eq!(
             vec![after_special_edit_for_refresh.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Host should react to /refresh LSP request"
         );
     });
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec![after_special_edit_for_refresh.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Guest should get a /refresh LSP request propagated by host"
         );
     });
 }
 
-// This test started hanging on seed 2 after the theme settings
-// PR. The hypothesis is that it's been buggy for a while, but got lucky
-// on seeds.
-#[ignore]
 #[gpui::test(iterations = 10)]
 async fn test_inlay_hint_refresh_is_forwarded(
     cx_a: &mut TestAppContext,
@@ -2206,18 +2191,18 @@ async fn test_inlay_hint_refresh_is_forwarded(
     executor.finish_waiting();
 
     executor.run_until_parked();
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert!(
-            extract_hint_labels(editor).is_empty(),
+            extract_hint_labels(editor, cx).is_empty(),
             "Host should get no hints due to them turned off"
         );
     });
 
     executor.run_until_parked();
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec!["initial hint".to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Client should get its first hints when opens an editor"
         );
     });
@@ -2229,18 +2214,18 @@ async fn test_inlay_hint_refresh_is_forwarded(
         .into_response()
         .expect("inlay refresh request failed");
     executor.run_until_parked();
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert!(
-            extract_hint_labels(editor).is_empty(),
+            extract_hint_labels(editor, cx).is_empty(),
             "Host should get no hints due to them turned off, even after the /refresh"
         );
     });
 
     executor.run_until_parked();
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec!["other hint".to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Guest should get a /refresh LSP request propagated by host despite host hints are off"
         );
     });
@@ -4217,15 +4202,35 @@ fn tab_undo_assert(
     cx_b.assert_editor_state(expected_initial);
 }
 
-fn extract_hint_labels(editor: &Editor) -> Vec<String> {
-    let mut labels = Vec::new();
-    for hint in editor.inlay_hint_cache().hints() {
-        match hint.label {
-            project::InlayHintLabel::String(s) => labels.push(s),
-            _ => unreachable!(),
-        }
+fn extract_hint_labels(editor: &Editor, cx: &mut App) -> Vec<String> {
+    let lsp_store = editor.project().unwrap().read(cx).lsp_store();
+
+    let mut all_cached_labels = Vec::new();
+    let mut all_fetched_hints = Vec::new();
+    for buffer in editor.buffer().read(cx).all_buffers() {
+        lsp_store.update(cx, |lsp_store, cx| {
+            let hints = &lsp_store.latest_lsp_data(&buffer, cx).inlay_hints();
+            all_cached_labels.extend(hints.all_cached_hints().into_iter().map(|hint| {
+                let mut label = hint.text().to_string();
+                if hint.padding_left {
+                    label.insert(0, ' ');
+                }
+                if hint.padding_right {
+                    label.push_str(" ");
+                }
+                label
+            }));
+            all_fetched_hints.extend(hints.all_fetched_hints());
+        });
     }
-    labels
+
+    assert!(
+        all_fetched_hints.is_empty(),
+        "Did not expect background hints fetch tasks, but got {} of them",
+        all_fetched_hints.len()
+    );
+
+    all_cached_labels
 }
 
 #[track_caller]

crates/collab/src/tests/editor_tests.rs

@@ -1,9 +1,9 @@
 use super::*;
 use collections::{HashMap, HashSet};
 use editor::{
-    DisplayPoint, EditorSettings,
+    DisplayPoint, EditorSettings, Inlay,
     actions::{GoToDiagnostic, GoToPreviousDiagnostic, Hover, MoveToBeginning},
-    display_map::{DisplayRow, Inlay},
+    display_map::DisplayRow,
     test::{
         editor_content_with_blocks, editor_lsp_test_context::EditorLspTestContext,
         editor_test_context::EditorTestContext,