9 files changed, 78 insertions, 46 deletions

diff --git a/crates/completion/src/item.rs b/crates/completion/src/item.rs
index ce6a44e57..0e59f73cb 100644
--- a/crates/completion/src/item.rs
+++ b/crates/completion/src/item.rs
@@ -257,14 +257,18 @@ impl CompletionItem {
     pub fn ref_match(&self) -> Option<(Mutability, CompletionScore)> {
         self.ref_match
     }
+
+    pub fn import_to_add(&self) -> Option<&ImportToAdd> {
+        self.import_to_add.as_ref()
+    }
 }
 
 /// An extra import to add after the completion is applied.
-#[derive(Clone)]
-pub(crate) struct ImportToAdd {
-    pub(crate) import_path: ModPath,
-    pub(crate) import_scope: ImportScope,
-    pub(crate) merge_behaviour: Option<MergeBehaviour>,
+#[derive(Debug, Clone)]
+pub struct ImportToAdd {
+    pub import_path: ModPath,
+    pub import_scope: ImportScope,
+    pub merge_behaviour: Option<MergeBehaviour>,
 }
 
 /// A helper to make `CompletionItem`s.
diff --git a/crates/completion/src/lib.rs b/crates/completion/src/lib.rs
index 1ec2e9be7..28209d4e0 100644
--- a/crates/completion/src/lib.rs
+++ b/crates/completion/src/lib.rs
@@ -18,7 +18,7 @@ use crate::{completions::Completions, context::CompletionContext, item::Completi
 
 pub use crate::{
     config::CompletionConfig,
-    item::{CompletionItem, CompletionItemKind, CompletionScore, InsertTextFormat},
+    item::{CompletionItem, CompletionItemKind, CompletionScore, ImportToAdd, 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 5244bdd61..7015a5126 100644
--- a/crates/ide/src/lib.rs
+++ b/crates/ide/src/lib.rs
@@ -80,7 +80,8 @@ pub use crate::{
     },
 };
 pub use completion::{
-    CompletionConfig, CompletionItem, CompletionItemKind, CompletionScore, InsertTextFormat,
+    CompletionConfig, CompletionItem, CompletionItemKind, CompletionScore, ImportToAdd,
+    InsertTextFormat,
 };
 pub use ide_db::{
     call_info::CallInfo,
diff --git a/crates/rust-analyzer/src/completions.rs b/crates/rust-analyzer/src/completions.rs
new file mode 100644
index 000000000..d4971f06e
--- /dev/null
+++ b/crates/rust-analyzer/src/completions.rs
@@ -0,0 +1,2 @@
+#[derive(Debug, Default)]
+pub struct CompletionResolveActions {}
diff --git a/crates/rust-analyzer/src/global_state.rs b/crates/rust-analyzer/src/global_state.rs
index a27495d0d..dc9e7113b 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};
+use ide::{Analysis, AnalysisHost, Change, FileId, ImportToAdd};
 use ide_db::base_db::{CrateId, VfsPath};
 use lsp_types::{SemanticTokens, Url};
 use parking_lot::{Mutex, RwLock};
@@ -69,6 +69,7 @@ pub(crate) struct GlobalState {
     pub(crate) config: Config,
     pub(crate) analysis_host: AnalysisHost,
     pub(crate) diagnostics: DiagnosticCollection,
+    pub(crate) additional_imports: FxHashMap<String, ImportToAdd>,
     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>)>>,
@@ -121,6 +122,7 @@ impl GlobalState {
             config,
             analysis_host,
             diagnostics: Default::default(),
+            additional_imports: FxHashMap::default(),
             mem_docs: FxHashMap::default(),
             semantic_tokens_cache: Arc::new(Default::default()),
             vfs: Arc::new(RwLock::new((vfs::Vfs::default(), FxHashMap::default()))),
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index 255a6e489..853f7fa84 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -11,6 +11,7 @@ use ide::{
     FileId, FilePosition, FileRange, HoverAction, HoverGotoTypeData, NavigationTarget, Query,
     RangeInfo, Runnable, RunnableKind, SearchScope, TextEdit,
 };
+use ide_db::helpers::{insert_use, mod_path_to_ast};
 use itertools::Itertools;
 use lsp_server::ErrorCode;
 use lsp_types::{
@@ -24,6 +25,7 @@ use lsp_types::{
     SymbolTag, TextDocumentIdentifier, Url, WorkspaceEdit,
 };
 use project_model::TargetKind;
+use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
 use serde_json::to_value;
 use stdx::{format_to, split_once};
@@ -535,10 +537,11 @@ pub(crate) fn handle_runnables(
 }
 
 pub(crate) fn handle_completion(
-    snap: GlobalStateSnapshot,
+    global_state: &mut GlobalState,
     params: lsp_types::CompletionParams,
 ) -> Result<Option<lsp_types::CompletionResponse>> {
     let _p = profile::span("handle_completion");
+    let snap = global_state.snapshot();
     let position = from_proto::file_position(&snap, params.text_document_position)?;
     let completion_triggered_after_single_colon = {
         let mut res = false;
@@ -568,22 +571,68 @@ pub(crate) fn handle_completion(
     };
     let line_index = snap.analysis.file_line_index(position.file_id)?;
     let line_endings = snap.file_line_endings(position.file_id);
+    let mut additional_imports = FxHashMap::default();
+
     let items: Vec<CompletionItem> = items
         .into_iter()
-        .flat_map(|item| to_proto::completion_item(&line_index, line_endings, item))
+        .flat_map(|item| {
+            let import_to_add = item.import_to_add().cloned();
+            let new_completion_items = to_proto::completion_item(&line_index, line_endings, item);
+            if let Some(import_to_add) = import_to_add {
+                for new_item in &new_completion_items {
+                    additional_imports.insert(new_item.label.clone(), import_to_add.clone());
+                }
+            }
+            new_completion_items
+        })
+        .map(|mut item| {
+            item.data = Some(position.file_id.0.into());
+            item
+        })
         .collect();
 
+    global_state.additional_imports = additional_imports;
+
     let completion_list = lsp_types::CompletionList { is_incomplete: true, items };
     Ok(Some(completion_list.into()))
 }
 
 pub(crate) fn handle_resolve_completion(
-    snap: GlobalStateSnapshot,
-    original_completion: CompletionItem,
-) -> Result<CompletionItem> {
+    global_state: &mut GlobalState,
+    mut original_completion: lsp_types::CompletionItem,
+) -> Result<lsp_types::CompletionItem> {
+    // TODO kb slow, takes over 130ms
     let _p = profile::span("handle_resolve_completion");
-    // TODO kb use the field to detect it's for autocompletion and do the insert logic
-    let _data = dbg!(original_completion).data;
+
+    if let Some(import_data) =
+        global_state.additional_imports.get(dbg!(original_completion.label.as_str()))
+    {
+        let rewriter = insert_use::insert_use(
+            &import_data.import_scope,
+            mod_path_to_ast(&import_data.import_path),
+            import_data.merge_behaviour,
+        );
+        if let Some((old_ast, file_id)) =
+            // TODO kb for file_id, better use &str and then cast to u32?
+            rewriter
+            .rewrite_root()
+            .zip(original_completion.data.as_ref().and_then(|value| Some(value.as_u64()? as u32)))
+        {
+            let snap = global_state.snapshot();
+            let mut import_insert = TextEdit::builder();
+            algo::diff(&old_ast, &rewriter.rewrite(&old_ast)).into_text_edit(&mut import_insert);
+            let line_index = snap.analysis.file_line_index(FileId(file_id))?;
+            let line_endings = snap.file_line_endings(FileId(file_id));
+            let text_edit = import_insert.finish();
+
+            let mut new_edits = original_completion.additional_text_edits.unwrap_or_default();
+            for indel in text_edit {
+                new_edits.push(to_proto::text_edit(&line_index, line_endings, indel));
+            }
+            original_completion.additional_text_edits = Some(new_edits);
+        }
+    }
+
     Ok(original_completion)
 }
 
diff --git a/crates/rust-analyzer/src/lib.rs b/crates/rust-analyzer/src/lib.rs
index ad08f1afb..13f14398f 100644
--- a/crates/rust-analyzer/src/lib.rs
+++ b/crates/rust-analyzer/src/lib.rs
@@ -36,6 +36,7 @@ mod thread_pool;
 mod document;
 pub mod lsp_ext;
 pub mod config;
+mod completions;
 
 use ide::AnalysisHost;
 use serde::de::DeserializeOwned;
diff --git a/crates/rust-analyzer/src/main_loop.rs b/crates/rust-analyzer/src/main_loop.rs
index 12b0946ac..21c58d959 100644
--- a/crates/rust-analyzer/src/main_loop.rs
+++ b/crates/rust-analyzer/src/main_loop.rs
@@ -436,6 +436,10 @@ impl GlobalState {
                 handlers::handle_matching_brace(s.snapshot(), p)
             })?
             .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_resolve_completion,
+            )?
             .on::<lsp_ext::AnalyzerStatus>(handlers::handle_analyzer_status)
             .on::<lsp_ext::SyntaxTree>(handlers::handle_syntax_tree)
             .on::<lsp_ext::ExpandMacro>(handlers::handle_expand_macro)
@@ -453,8 +457,6 @@ impl GlobalState {
             .on::<lsp_types::request::GotoDefinition>(handlers::handle_goto_definition)
             .on::<lsp_types::request::GotoImplementation>(handlers::handle_goto_implementation)
             .on::<lsp_types::request::GotoTypeDefinition>(handlers::handle_goto_type_definition)
-            .on::<lsp_types::request::Completion>(handlers::handle_completion)
-            .on::<lsp_types::request::ResolveCompletionItem>(handlers::handle_resolve_completion)
             .on::<lsp_types::request::CodeLensRequest>(handlers::handle_code_lens)
             .on::<lsp_types::request::CodeLensResolve>(handlers::handle_code_lens_resolve)
             .on::<lsp_types::request::FoldingRangeRequest>(handlers::handle_folding_range)
diff --git a/crates/rust-analyzer/src/to_proto.rs b/crates/rust-analyzer/src/to_proto.rs
index db9ed08f6..01eabe852 100644
--- a/crates/rust-analyzer/src/to_proto.rs
+++ b/crates/rust-analyzer/src/to_proto.rs
@@ -231,35 +231,6 @@ pub(crate) fn completion_item(
         None => vec![res],
     };
 
-    // TODO kb need to get this logic away and store for the later resolve request
-    /*
-    let mut label = self.label;
-    let mut lookup = self.lookup;
-    let mut insert_text = self.insert_text;
-    let mut text_edits = TextEdit::builder();
-
-    if let Some((import_path, import_scope, merge_behaviour)) = completion_item.import_data.as_ref() {
-        let import = mod_path_to_ast(&import_path);
-            let mut import_path_without_last_segment = import_path;
-            let _ = import_path_without_last_segment.segments.pop();
-
-            if !import_path_without_last_segment.segments.is_empty() {
-                if lookup.is_none() {
-                    lookup = Some(label.clone());
-                }
-                if insert_text.is_none() {
-                    insert_text = Some(label.clone());
-                }
-                label = format!("{}::{}", import_path_without_last_segment, label);
-            }
-
-            let rewriter = insert_use(&import_scope, import, merge_behaviour);
-            if let Some(old_ast) = rewriter.rewrite_root() {
-                algo::diff(&old_ast, &rewriter.rewrite(&old_ast)).into_text_edit(&mut text_edits);
-            }
-    }
-    */
-
     for mut r in all_results.iter_mut() {
         r.insert_text_format = Some(insert_text_format(completion_item.insert_text_format()));
     }