Make completion resolve async

7 files changed, 116 insertions, 53 deletions

diff --git a/crates/completion/src/item.rs b/crates/completion/src/item.rs
index 4e56f28f3..dd25ca75c 100644
--- a/crates/completion/src/item.rs
+++ b/crates/completion/src/item.rs
@@ -4,10 +4,10 @@ use std::fmt;
 
 use hir::{Documentation, ModPath, Mutability};
 use ide_db::helpers::{
-    insert_use::{self, ImportScope, MergeBehaviour},
+    insert_use::{self, ImportScope, ImportScopePtr, MergeBehaviour},
     mod_path_to_ast,
 };
-use syntax::{algo, TextRange};
+use syntax::{algo, SyntaxNode, TextRange};
 use text_edit::TextEdit;
 
 use crate::config::SnippetCap;
@@ -275,7 +275,32 @@ pub struct ImportEdit {
     pub merge_behaviour: Option<MergeBehaviour>,
 }
 
+#[derive(Debug, Clone)]
+pub struct ImportEditPtr {
+    pub import_path: ModPath,
+    pub import_scope: ImportScopePtr,
+    pub merge_behaviour: Option<MergeBehaviour>,
+}
+
+impl ImportEditPtr {
+    pub fn into_import_edit(self, root: &SyntaxNode) -> Option<ImportEdit> {
+        Some(ImportEdit {
+            import_path: self.import_path,
+            import_scope: self.import_scope.into_scope(root)?,
+            merge_behaviour: self.merge_behaviour,
+        })
+    }
+}
+
 impl ImportEdit {
+    pub fn get_edit_ptr(&self) -> ImportEditPtr {
+        ImportEditPtr {
+            import_path: self.import_path.clone(),
+            import_scope: self.import_scope.get_ptr(),
+            merge_behaviour: self.merge_behaviour,
+        }
+    }
+
     /// Attempts to insert the import to the given scope, producing a text edit.
     /// May return no edit in edge cases, such as scope already containing the import.
     pub fn to_text_edit(&self) -> Option<TextEdit> {
diff --git a/crates/completion/src/lib.rs b/crates/completion/src/lib.rs
index c57203c80..c277cd466 100644
--- a/crates/completion/src/lib.rs
+++ b/crates/completion/src/lib.rs
@@ -18,7 +18,10 @@ use crate::{completions::Completions, context::CompletionContext, item::Completi
 
 pub use crate::{
     config::{CompletionConfig, CompletionResolveCapability},
-    item::{CompletionItem, CompletionItemKind, CompletionScore, ImportEdit, InsertTextFormat},
+    item::{
+        CompletionItem, CompletionItemKind, CompletionScore, ImportEdit, ImportEditPtr,
+        InsertTextFormat,
+    },
 };
 
 //FIXME: split the following feature into fine-grained features.
diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs
index 9e38d6506..c52e53d27 100644
--- a/crates/ide/src/lib.rs
+++ b/crates/ide/src/lib.rs
@@ -81,7 +81,7 @@ pub use crate::{
 };
 pub use completion::{
     CompletionConfig, CompletionItem, CompletionItemKind, CompletionResolveCapability,
-    CompletionScore, ImportEdit, InsertTextFormat,
+    CompletionScore, ImportEdit, ImportEditPtr, InsertTextFormat,
 };
 pub use ide_db::{
     call_info::CallInfo,
diff --git a/crates/ide_db/src/helpers/insert_use.rs b/crates/ide_db/src/helpers/insert_use.rs
index 040843990..0dae9a541 100644
--- a/crates/ide_db/src/helpers/insert_use.rs
+++ b/crates/ide_db/src/helpers/insert_use.rs
@@ -11,7 +11,7 @@ use syntax::{
         edit::{AstNodeEdit, IndentLevel},
         make, AstNode, PathSegmentKind, VisibilityOwner,
     },
-    AstToken, InsertPosition, NodeOrToken, SyntaxElement, SyntaxNode, SyntaxToken,
+    AstToken, InsertPosition, NodeOrToken, SyntaxElement, SyntaxNode, SyntaxNodePtr, SyntaxToken,
 };
 use test_utils::mark;
 
@@ -22,6 +22,36 @@ pub enum ImportScope {
 }
 
 impl ImportScope {
+    pub fn get_ptr(&self) -> ImportScopePtr {
+        match self {
+            ImportScope::File(file) => ImportScopePtr::File(SyntaxNodePtr::new(file.syntax())),
+            ImportScope::Module(module) => {
+                ImportScopePtr::Module(SyntaxNodePtr::new(module.syntax()))
+            }
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub enum ImportScopePtr {
+    File(SyntaxNodePtr),
+    Module(SyntaxNodePtr),
+}
+
+impl ImportScopePtr {
+    pub fn into_scope(self, root: &SyntaxNode) -> Option<ImportScope> {
+        Some(match self {
+            ImportScopePtr::File(file_ptr) => {
+                ImportScope::File(ast::SourceFile::cast(file_ptr.to_node(root))?)
+            }
+            ImportScopePtr::Module(module_ptr) => {
+                ImportScope::File(ast::SourceFile::cast(module_ptr.to_node(root))?)
+            }
+        })
+    }
+}
+
+impl ImportScope {
     pub fn from(syntax: SyntaxNode) -> Option<Self> {
         if let Some(module) = ast::Module::cast(syntax.clone()) {
             module.item_list().map(ImportScope::Module)
diff --git a/crates/rust-analyzer/src/global_state.rs b/crates/rust-analyzer/src/global_state.rs
index a35abe86d..0fe69b996 100644
--- a/crates/rust-analyzer/src/global_state.rs
+++ b/crates/rust-analyzer/src/global_state.rs
@@ -7,7 +7,7 @@ use std::{sync::Arc, time::Instant};
 
 use crossbeam_channel::{unbounded, Receiver, Sender};
 use flycheck::FlycheckHandle;
-use ide::{Analysis, AnalysisHost, Change, FileId, ImportEdit};
+use ide::{Analysis, AnalysisHost, Change, FileId, ImportEditPtr};
 use ide_db::base_db::{CrateId, VfsPath};
 use lsp_types::{SemanticTokens, Url};
 use parking_lot::{Mutex, RwLock};
@@ -51,11 +51,6 @@ pub(crate) struct Handle<H, C> {
 pub(crate) type ReqHandler = fn(&mut GlobalState, lsp_server::Response);
 pub(crate) type ReqQueue = lsp_server::ReqQueue<(String, Instant), ReqHandler>;
 
-pub(crate) struct CompletionResolveData {
-    pub(crate) file_id: FileId,
-    pub(crate) import_edit: ImportEdit,
-}
-
 /// `GlobalState` is the primary mutable state of the language server
 ///
 /// The most interesting components are `vfs`, which stores a consistent
@@ -74,7 +69,7 @@ pub(crate) struct GlobalState {
     pub(crate) config: Config,
     pub(crate) analysis_host: AnalysisHost,
     pub(crate) diagnostics: DiagnosticCollection,
-    pub(crate) completion_resolve_data: FxHashMap<usize, CompletionResolveData>,
+    pub(crate) completion_resolve_data: Arc<FxHashMap<usize, ImportEditPtr>>,
     pub(crate) mem_docs: FxHashMap<VfsPath, DocumentData>,
     pub(crate) semantic_tokens_cache: Arc<Mutex<FxHashMap<Url, SemanticTokens>>>,
     pub(crate) vfs: Arc<RwLock<(vfs::Vfs, FxHashMap<FileId, LineEndings>)>>,
@@ -96,6 +91,7 @@ pub(crate) struct GlobalStateSnapshot {
     pub(crate) semantic_tokens_cache: Arc<Mutex<FxHashMap<Url, SemanticTokens>>>,
     vfs: Arc<RwLock<(vfs::Vfs, FxHashMap<FileId, LineEndings>)>>,
     pub(crate) workspaces: Arc<Vec<ProjectWorkspace>>,
+    pub(crate) completion_resolve_data: Arc<FxHashMap<usize, ImportEditPtr>>,
 }
 
 impl GlobalState {
@@ -127,7 +123,7 @@ impl GlobalState {
             config,
             analysis_host,
             diagnostics: Default::default(),
-            completion_resolve_data: FxHashMap::default(),
+            completion_resolve_data: Arc::new(FxHashMap::default()),
             mem_docs: FxHashMap::default(),
             semantic_tokens_cache: Arc::new(Default::default()),
             vfs: Arc::new(RwLock::new((vfs::Vfs::default(), FxHashMap::default()))),
@@ -198,6 +194,7 @@ impl GlobalState {
             check_fixes: Arc::clone(&self.diagnostics.check_fixes),
             mem_docs: self.mem_docs.clone(),
             semantic_tokens_cache: Arc::clone(&self.semantic_tokens_cache),
+            completion_resolve_data: Arc::clone(&self.completion_resolve_data),
         }
     }
 
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index 3eb5f26bc..1ea1e1f43 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -5,11 +5,13 @@
 use std::{
     io::Write as _,
     process::{self, Stdio},
+    sync::Arc,
 };
 
 use ide::{
-    FileId, FilePosition, FileRange, HoverAction, HoverGotoTypeData, ImportEdit, LineIndex,
-    NavigationTarget, Query, RangeInfo, Runnable, RunnableKind, SearchScope, TextEdit,
+    CompletionResolveCapability, FileId, FilePosition, FileRange, HoverAction, HoverGotoTypeData,
+    ImportEdit, LineIndex, NavigationTarget, Query, RangeInfo, Runnable, RunnableKind, SearchScope,
+    TextEdit,
 };
 use itertools::Itertools;
 use lsp_server::ErrorCode;
@@ -34,7 +36,7 @@ use crate::{
     cargo_target_spec::CargoTargetSpec,
     config::RustfmtConfig,
     from_json, from_proto,
-    global_state::{CompletionResolveData, GlobalState, GlobalStateSnapshot},
+    global_state::{GlobalState, GlobalStateSnapshot},
     line_endings::LineEndings,
     lsp_ext::{self, InlayHint, InlayHintsParams},
     to_proto, LspError, Result,
@@ -542,6 +544,7 @@ pub(crate) fn handle_completion(
 ) -> Result<Option<lsp_types::CompletionResponse>> {
     let _p = profile::span("handle_completion");
     let snap = global_state.snapshot();
+    let text_document_url = params.text_document_position.text_document.uri.clone();
     let position = from_proto::file_position(&snap, params.text_document_position)?;
     let completion_triggered_after_single_colon = {
         let mut res = false;
@@ -582,18 +585,15 @@ pub(crate) fn handle_completion(
 
             if snap.config.completion.resolve_additional_edits_lazily() {
                 if let Some(import_edit) = item.import_to_add() {
-                    completion_resolve_data.insert(
-                        item_index,
-                        CompletionResolveData {
-                            file_id: position.file_id,
-                            import_edit: import_edit.clone(),
-                        },
-                    );
-
-                    let item_id = serde_json::to_value(&item_index)
-                        .expect(&format!("Should be able to serialize usize value {}", item_index));
+                    completion_resolve_data.insert(item_index, import_edit.get_edit_ptr());
+
+                    let data = serde_json::to_value(&CompletionData {
+                        document_url: text_document_url.clone(),
+                        import_id: item_index,
+                    })
+                    .expect(&format!("Should be able to serialize usize value {}", item_index));
                     for new_item in &mut new_completion_items {
-                        new_item.data = Some(item_id.clone());
+                        new_item.data = Some(data.clone());
                     }
                 }
             }
@@ -602,50 +602,54 @@ pub(crate) fn handle_completion(
         })
         .collect();
 
-    global_state.completion_resolve_data = completion_resolve_data;
+    global_state.completion_resolve_data = Arc::new(completion_resolve_data);
 
     let completion_list = lsp_types::CompletionList { is_incomplete: true, items };
     Ok(Some(completion_list.into()))
 }
 
 pub(crate) fn handle_completion_resolve(
-    global_state: &mut GlobalState,
+    snap: GlobalStateSnapshot,
     mut original_completion: lsp_types::CompletionItem,
 ) -> Result<lsp_types::CompletionItem> {
     let _p = profile::span("handle_resolve_completion");
 
-    let active_resolve_caps = &global_state.config.completion.active_resolve_capabilities;
-    if active_resolve_caps.is_empty() {
+    // FIXME resolve the other capabilities also?
+    if !snap
+        .config
+        .completion
+        .active_resolve_capabilities
+        .contains(&CompletionResolveCapability::AdditionalTextEdits)
+    {
         return Ok(original_completion);
     }
 
-    let server_completion_data = match original_completion
+    let (import_edit_ptr, document_url) = match original_completion
         .data
         .as_ref()
-        .map(|data| serde_json::from_value::<usize>(data.clone()))
+        .map(|data| serde_json::from_value::<CompletionData>(data.clone()))
         .transpose()?
-        .and_then(|server_completion_id| {
-            global_state.completion_resolve_data.get(&server_completion_id)
+        .and_then(|data| {
+            let import_edit_ptr = snap.completion_resolve_data.get(&data.import_id).cloned();
+            Some((import_edit_ptr, data.document_url))
         }) {
         Some(data) => data,
         None => return Ok(original_completion),
     };
 
-    let snap = &global_state.snapshot();
-    for supported_completion_resolve_cap in active_resolve_caps {
-        match supported_completion_resolve_cap {
-            // FIXME actually add all additional edits here? see `to_proto::completion_item` for more
-            ide::CompletionResolveCapability::AdditionalTextEdits => {
-                append_import_edits(
-                    &mut original_completion,
-                    &server_completion_data.import_edit,
-                    snap.analysis.file_line_index(server_completion_data.file_id)?.as_ref(),
-                    snap.file_line_endings(server_completion_data.file_id),
-                );
-            }
-            // FIXME resolve the other capabilities also?
-            _ => {}
-        }
+    let file_id = from_proto::file_id(&snap, &document_url)?;
+    let root = snap.analysis.parse(file_id)?;
+
+    if let Some(import_to_add) =
+        import_edit_ptr.and_then(|import_edit| import_edit.into_import_edit(root.syntax()))
+    {
+        // FIXME actually add all additional edits here? see `to_proto::completion_item` for more
+        append_import_edits(
+            &mut original_completion,
+            &import_to_add,
+            snap.analysis.file_line_index(file_id)?.as_ref(),
+            snap.file_line_endings(file_id),
+        );
     }
 
     Ok(original_completion)
@@ -1609,6 +1613,12 @@ fn should_skip_target(runnable: &Runnable, cargo_spec: Option<&CargoTargetSpec>)
     }
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+struct CompletionData {
+    document_url: Url,
+    import_id: usize,
+}
+
 fn append_import_edits(
     completion: &mut lsp_types::CompletionItem,
     import_to_add: &ImportEdit,
diff --git a/crates/rust-analyzer/src/main_loop.rs b/crates/rust-analyzer/src/main_loop.rs
index db30b3dce..aad37fde1 100644
--- a/crates/rust-analyzer/src/main_loop.rs
+++ b/crates/rust-analyzer/src/main_loop.rs
@@ -437,9 +437,7 @@ impl GlobalState {
             })?
             .on_sync::<lsp_ext::MemoryUsage>(|s, p| handlers::handle_memory_usage(s, p))?
             .on_sync::<lsp_types::request::Completion>(handlers::handle_completion)?
-            .on_sync::<lsp_types::request::ResolveCompletionItem>(
-                handlers::handle_completion_resolve,
-            )?
+            .on::<lsp_types::request::ResolveCompletionItem>(handlers::handle_completion_resolve)
             .on::<lsp_ext::AnalyzerStatus>(handlers::handle_analyzer_status)
             .on::<lsp_ext::SyntaxTree>(handlers::handle_syntax_tree)
             .on::<lsp_ext::ExpandMacro>(handlers::handle_expand_macro)