Simplify import edit calculation

12 files changed, 114 insertions, 111 deletions

diff --git a/crates/completion/src/config.rs b/crates/completion/src/config.rs
index eacdd3449..487c1d0f1 100644
--- a/crates/completion/src/config.rs
+++ b/crates/completion/src/config.rs
@@ -36,12 +36,10 @@ impl CompletionConfig {
         self.snippet_cap = if yes { Some(SnippetCap { _private: () }) } else { None }
     }
 
-    /// Whether the completions' additional edits are calculated later, during a resolve request or not.
-    /// See `CompletionResolveCapability` for the details.
-    pub fn resolve_edits_immediately(&self) -> bool {
-        !self
-            .active_resolve_capabilities
-            .contains(&CompletionResolveCapability::AdditionalTextEdits)
+    /// Whether the completions' additional edits are calculated when sending an initional completions list
+    /// or later, in a separate resolve request.
+    pub fn resolve_additional_edits_lazily(&self) -> bool {
+        self.active_resolve_capabilities.contains(&CompletionResolveCapability::AdditionalTextEdits)
     }
 }
 
diff --git a/crates/completion/src/item.rs b/crates/completion/src/item.rs
index 775245b3b..4e56f28f3 100644
--- a/crates/completion/src/item.rs
+++ b/crates/completion/src/item.rs
@@ -68,7 +68,7 @@ pub struct CompletionItem {
     ref_match: Option<(Mutability, CompletionScore)>,
 
     /// The import data to add to completion's edits.
-    import_to_add: Option<ImportToAdd>,
+    import_to_add: Option<ImportEdit>,
 }
 
 // We use custom debug for CompletionItem to make snapshot tests more readable.
@@ -209,7 +209,7 @@ impl CompletionItem {
             score: None,
             ref_match: None,
             import_to_add: None,
-            resolve_import_immediately: true,
+            resolve_import_lazily: false,
         }
     }
 
@@ -262,27 +262,46 @@ impl CompletionItem {
         self.ref_match
     }
 
-    pub fn import_to_add(&self) -> Option<&ImportToAdd> {
+    pub fn import_to_add(&self) -> Option<&ImportEdit> {
         self.import_to_add.as_ref()
     }
 }
 
 /// An extra import to add after the completion is applied.
 #[derive(Debug, Clone)]
-pub struct ImportToAdd {
+pub struct ImportEdit {
     pub import_path: ModPath,
     pub import_scope: ImportScope,
     pub merge_behaviour: Option<MergeBehaviour>,
 }
 
+impl ImportEdit {
+    /// 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> {
+        let _p = profile::span("ImportEdit::to_edit");
+
+        let rewriter = insert_use::insert_use(
+            &self.import_scope,
+            mod_path_to_ast(&self.import_path),
+            self.merge_behaviour,
+        );
+        let old_ast = rewriter.rewrite_root()?;
+        let mut import_insert = TextEdit::builder();
+        algo::diff(&old_ast, &rewriter.rewrite(&old_ast)).into_text_edit(&mut import_insert);
+
+        Some(import_insert.finish())
+    }
+}
+
 /// A helper to make `CompletionItem`s.
 #[must_use]
 #[derive(Clone)]
 pub(crate) struct Builder {
     source_range: TextRange,
     completion_kind: CompletionKind,
-    import_to_add: Option<ImportToAdd>,
-    resolve_import_immediately: bool,
+    import_to_add: Option<ImportEdit>,
+    resolve_import_lazily: bool,
     label: String,
     insert_text: Option<String>,
     insert_text_format: InsertTextFormat,
@@ -304,7 +323,6 @@ impl Builder {
         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_to_add) = self.import_to_add.as_ref() {
             let mut import_path_without_last_segment = import_to_add.import_path.to_owned();
@@ -319,35 +337,28 @@ impl Builder {
                 }
                 label = format!("{}::{}", import_path_without_last_segment, label);
             }
-
-            if self.resolve_import_immediately {
-                let rewriter = insert_use::insert_use(
-                    &import_to_add.import_scope,
-                    mod_path_to_ast(&import_to_add.import_path),
-                    import_to_add.merge_behaviour,
-                );
-                if let Some(old_ast) = rewriter.rewrite_root() {
-                    algo::diff(&old_ast, &rewriter.rewrite(&old_ast))
-                        .into_text_edit(&mut text_edits);
-                }
-            }
         }
 
-        let original_edit = match self.text_edit {
+        let mut text_edit = match self.text_edit {
             Some(it) => it,
             None => {
                 TextEdit::replace(self.source_range, insert_text.unwrap_or_else(|| label.clone()))
             }
         };
 
-        let mut resulting_edit = text_edits.finish();
-        resulting_edit.union(original_edit).expect("Failed to unite text edits");
+        if !self.resolve_import_lazily {
+            if let Some(import_edit) =
+                self.import_to_add.as_ref().and_then(|import_edit| import_edit.to_text_edit())
+            {
+                text_edit.union(import_edit).expect("Failed to unite import and completion edits");
+            }
+        }
 
         CompletionItem {
             source_range: self.source_range,
             label,
             insert_text_format: self.insert_text_format,
-            text_edit: resulting_edit,
+            text_edit,
             detail: self.detail,
             documentation: self.documentation,
             lookup,
@@ -422,11 +433,11 @@ impl Builder {
     }
     pub(crate) fn add_import(
         mut self,
-        import_to_add: Option<ImportToAdd>,
-        resolve_import_immediately: bool,
+        import_to_add: Option<ImportEdit>,
+        resolve_import_lazily: bool,
     ) -> Builder {
         self.import_to_add = import_to_add;
-        self.resolve_import_immediately = resolve_import_immediately;
+        self.resolve_import_lazily = resolve_import_lazily;
         self
     }
     pub(crate) fn set_ref_match(
diff --git a/crates/completion/src/lib.rs b/crates/completion/src/lib.rs
index c689b0dde..c57203c80 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, CompletionResolveCapability},
-    item::{CompletionItem, CompletionItemKind, CompletionScore, ImportToAdd, InsertTextFormat},
+    item::{CompletionItem, CompletionItemKind, CompletionScore, ImportEdit, InsertTextFormat},
 };
 
 //FIXME: split the following feature into fine-grained features.
diff --git a/crates/completion/src/render.rs b/crates/completion/src/render.rs
index 2b4f1ea14..a6faedb18 100644
--- a/crates/completion/src/render.rs
+++ b/crates/completion/src/render.rs
@@ -16,7 +16,7 @@ use syntax::TextRange;
 use test_utils::mark;
 
 use crate::{
-    config::SnippetCap, item::ImportToAdd, CompletionContext, CompletionItem, CompletionItemKind,
+    config::SnippetCap, item::ImportEdit, CompletionContext, CompletionItem, CompletionItemKind,
     CompletionKind, CompletionScore,
 };
 
@@ -56,7 +56,7 @@ pub(crate) fn render_resolution_with_import<'a>(
     let local_name = import_path.segments.last()?.to_string();
     Render::new(ctx).render_resolution(
         local_name,
-        Some(ImportToAdd { import_path, import_scope, merge_behaviour }),
+        Some(ImportEdit { import_path, import_scope, merge_behaviour }),
         resolution,
     )
 }
@@ -147,7 +147,7 @@ impl<'a> Render<'a> {
     fn render_resolution(
         self,
         local_name: String,
-        import_to_add: Option<ImportToAdd>,
+        import_to_add: Option<ImportEdit>,
         resolution: &ScopeDef,
     ) -> Option<CompletionItem> {
         let _p = profile::span("render_resolution");
@@ -194,7 +194,10 @@ impl<'a> Render<'a> {
                     local_name,
                 )
                 .kind(CompletionItemKind::UnresolvedReference)
-                .add_import(import_to_add, self.ctx.completion.config.resolve_edits_immediately())
+                .add_import(
+                    import_to_add,
+                    self.ctx.completion.config.resolve_additional_edits_lazily(),
+                )
                 .build();
                 return Some(item);
             }
@@ -249,7 +252,7 @@ impl<'a> Render<'a> {
 
         let item = item
             .kind(kind)
-            .add_import(import_to_add, self.ctx.completion.config.resolve_edits_immediately())
+            .add_import(import_to_add, self.ctx.completion.config.resolve_additional_edits_lazily())
             .set_documentation(docs)
             .set_ref_match(ref_match)
             .build();
diff --git a/crates/completion/src/render/enum_variant.rs b/crates/completion/src/render/enum_variant.rs
index 4a91fe3c7..e979a090b 100644
--- a/crates/completion/src/render/enum_variant.rs
+++ b/crates/completion/src/render/enum_variant.rs
@@ -5,13 +5,13 @@ use itertools::Itertools;
 use test_utils::mark;
 
 use crate::{
-    item::{CompletionItem, CompletionItemKind, CompletionKind, ImportToAdd},
+    item::{CompletionItem, CompletionItemKind, CompletionKind, ImportEdit},
     render::{builder_ext::Params, RenderContext},
 };
 
 pub(crate) fn render_enum_variant<'a>(
     ctx: RenderContext<'a>,
-    import_to_add: Option<ImportToAdd>,
+    import_to_add: Option<ImportEdit>,
     local_name: Option<String>,
     variant: hir::EnumVariant,
     path: Option<ModPath>,
@@ -62,7 +62,7 @@ impl<'a> EnumVariantRender<'a> {
         }
     }
 
-    fn render(self, import_to_add: Option<ImportToAdd>) -> CompletionItem {
+    fn render(self, import_to_add: Option<ImportEdit>) -> CompletionItem {
         let mut builder = CompletionItem::new(
             CompletionKind::Reference,
             self.ctx.source_range(),
@@ -71,7 +71,7 @@ impl<'a> EnumVariantRender<'a> {
         .kind(CompletionItemKind::EnumVariant)
         .set_documentation(self.variant.docs(self.ctx.db()))
         .set_deprecated(self.ctx.is_deprecated(self.variant))
-        .add_import(import_to_add, self.ctx.completion.config.resolve_edits_immediately())
+        .add_import(import_to_add, self.ctx.completion.config.resolve_additional_edits_lazily())
         .detail(self.detail());
 
         if self.variant_kind == StructKind::Tuple {
diff --git a/crates/completion/src/render/function.rs b/crates/completion/src/render/function.rs
index 20f2b9b7e..dd2c999ef 100644
--- a/crates/completion/src/render/function.rs
+++ b/crates/completion/src/render/function.rs
@@ -5,13 +5,13 @@ use syntax::{ast::Fn, display::function_declaration};
 use test_utils::mark;
 
 use crate::{
-    item::{CompletionItem, CompletionItemKind, CompletionKind, ImportToAdd},
+    item::{CompletionItem, CompletionItemKind, CompletionKind, ImportEdit},
     render::{builder_ext::Params, RenderContext},
 };
 
 pub(crate) fn render_fn<'a>(
     ctx: RenderContext<'a>,
-    import_to_add: Option<ImportToAdd>,
+    import_to_add: Option<ImportEdit>,
     local_name: Option<String>,
     fn_: hir::Function,
 ) -> CompletionItem {
@@ -39,7 +39,7 @@ impl<'a> FunctionRender<'a> {
         FunctionRender { ctx, name, func: fn_, ast_node }
     }
 
-    fn render(self, import_to_add: Option<ImportToAdd>) -> CompletionItem {
+    fn render(self, import_to_add: Option<ImportEdit>) -> CompletionItem {
         let params = self.params();
         CompletionItem::new(CompletionKind::Reference, self.ctx.source_range(), self.name.clone())
             .kind(self.kind())
@@ -47,7 +47,7 @@ impl<'a> FunctionRender<'a> {
             .set_deprecated(self.ctx.is_deprecated(self.func))
             .detail(self.detail())
             .add_call_parens(self.ctx.completion, self.name, params)
-            .add_import(import_to_add, self.ctx.completion.config.resolve_edits_immediately())
+            .add_import(import_to_add, self.ctx.completion.config.resolve_additional_edits_lazily())
             .build()
     }
 
diff --git a/crates/completion/src/render/macro_.rs b/crates/completion/src/render/macro_.rs
index be7c53659..bdbc642ca 100644
--- a/crates/completion/src/render/macro_.rs
+++ b/crates/completion/src/render/macro_.rs
@@ -5,13 +5,13 @@ use syntax::display::macro_label;
 use test_utils::mark;
 
 use crate::{
-    item::{CompletionItem, CompletionItemKind, CompletionKind, ImportToAdd},
+    item::{CompletionItem, CompletionItemKind, CompletionKind, ImportEdit},
     render::RenderContext,
 };
 
 pub(crate) fn render_macro<'a>(
     ctx: RenderContext<'a>,
-    import_to_add: Option<ImportToAdd>,
+    import_to_add: Option<ImportEdit>,
     name: String,
     macro_: hir::MacroDef,
 ) -> Option<CompletionItem> {
@@ -38,7 +38,7 @@ impl<'a> MacroRender<'a> {
         MacroRender { ctx, name, macro_, docs, bra, ket }
     }
 
-    fn render(&self, import_to_add: Option<ImportToAdd>) -> Option<CompletionItem> {
+    fn render(&self, import_to_add: Option<ImportEdit>) -> Option<CompletionItem> {
         // FIXME: Currently proc-macro do not have ast-node,
         // such that it does not have source
         if self.macro_.is_proc_macro() {
@@ -50,7 +50,10 @@ impl<'a> MacroRender<'a> {
                 .kind(CompletionItemKind::Macro)
                 .set_documentation(self.docs.clone())
                 .set_deprecated(self.ctx.is_deprecated(self.macro_))
-                .add_import(import_to_add, self.ctx.completion.config.resolve_edits_immediately())
+                .add_import(
+                    import_to_add,
+                    self.ctx.completion.config.resolve_additional_edits_lazily(),
+                )
                 .detail(self.detail());
 
         let needs_bang = self.needs_bang();
diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs
index d1a27f3a5..9e38d6506 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, ImportToAdd, InsertTextFormat,
+    CompletionScore, ImportEdit, InsertTextFormat,
 };
 pub use ide_db::{
     call_info::CallInfo,
diff --git a/crates/rust-analyzer/src/caps.rs b/crates/rust-analyzer/src/caps.rs
index 0b6ca76e2..1e0ee10e4 100644
--- a/crates/rust-analyzer/src/caps.rs
+++ b/crates/rust-analyzer/src/caps.rs
@@ -104,7 +104,7 @@ fn completions_resolve_provider(client_caps: &ClientCapabilities) -> Option<bool
 }
 
 /// Parses client capabilities and returns all completion resolve capabilities rust-analyzer supports.
-pub fn enabled_completions_resolve_capabilities(
+pub(crate) fn enabled_completions_resolve_capabilities(
     caps: &ClientCapabilities,
 ) -> Option<FxHashSet<CompletionResolveCapability>> {
     Some(
@@ -118,13 +118,11 @@ pub fn enabled_completions_resolve_capabilities(
             .as_ref()?
             .properties
             .iter()
-            .filter_map(|cap_string| {
-                Some(match cap_string.as_str() {
-                    "additionalTextEdits" => CompletionResolveCapability::AdditionalTextEdits,
-                    "detail" => CompletionResolveCapability::Detail,
-                    "documentation" => CompletionResolveCapability::Documentation,
-                    _unsupported => return None,
-                })
+            .filter_map(|cap_string| match cap_string.as_str() {
+                "additionalTextEdits" => Some(CompletionResolveCapability::AdditionalTextEdits),
+                "detail" => Some(CompletionResolveCapability::Detail),
+                "documentation" => Some(CompletionResolveCapability::Documentation),
+                _unsupported => None,
             })
             .collect(),
     )
diff --git a/crates/rust-analyzer/src/global_state.rs b/crates/rust-analyzer/src/global_state.rs
index e12651937..a35abe86d 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, CompletionItem, FileId};
+use ide::{Analysis, AnalysisHost, Change, FileId, ImportEdit};
 use ide_db::base_db::{CrateId, VfsPath};
 use lsp_types::{SemanticTokens, Url};
 use parking_lot::{Mutex, RwLock};
@@ -53,7 +53,7 @@ pub(crate) type ReqQueue = lsp_server::ReqQueue<(String, Instant), ReqHandler>;
 
 pub(crate) struct CompletionResolveData {
     pub(crate) file_id: FileId,
-    pub(crate) item: CompletionItem,
+    pub(crate) import_edit: ImportEdit,
 }
 
 /// `GlobalState` is the primary mutable state of the language server
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index 55c7b0c66..3eb5f26bc 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -8,10 +8,9 @@ use std::{
 };
 
 use ide::{
-    FileId, FilePosition, FileRange, HoverAction, HoverGotoTypeData, ImportToAdd, LineIndex,
+    FileId, FilePosition, FileRange, HoverAction, HoverGotoTypeData, ImportEdit, LineIndex,
     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::{
@@ -581,13 +580,21 @@ pub(crate) fn handle_completion(
             let mut new_completion_items =
                 to_proto::completion_item(&line_index, line_endings, item.clone());
 
-            if !snap.config.completion.active_resolve_capabilities.is_empty() {
-                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, CompletionResolveData { file_id: position.file_id, item });
-                for new_item in &mut new_completion_items {
-                    new_item.data = Some(item_id.clone());
+            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));
+                    for new_item in &mut new_completion_items {
+                        new_item.data = Some(item_id.clone());
+                    }
                 }
             }
 
@@ -601,12 +608,17 @@ pub(crate) fn handle_completion(
     Ok(Some(completion_list.into()))
 }
 
-pub(crate) fn handle_resolve_completion(
+pub(crate) fn handle_completion_resolve(
     global_state: &mut GlobalState,
     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() {
+        return Ok(original_completion);
+    }
+
     let server_completion_data = match original_completion
         .data
         .as_ref()
@@ -620,18 +632,16 @@ pub(crate) fn handle_resolve_completion(
     };
 
     let snap = &global_state.snapshot();
-    for supported_completion_resolve_cap in &snap.config.completion.active_resolve_capabilities {
+    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 => {
-                // FIXME actually add all additional edits here?
-                if let Some(import_to_add) = server_completion_data.item.import_to_add() {
-                    append_import_edits(
-                        &mut original_completion,
-                        import_to_add,
-                        snap.analysis.file_line_index(server_completion_data.file_id)?.as_ref(),
-                        snap.file_line_endings(server_completion_data.file_id),
-                    );
-                }
+                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?
             _ => {}
@@ -1601,41 +1611,21 @@ fn should_skip_target(runnable: &Runnable, cargo_spec: Option<&CargoTargetSpec>)
 
 fn append_import_edits(
     completion: &mut lsp_types::CompletionItem,
-    import_to_add: &ImportToAdd,
+    import_to_add: &ImportEdit,
     line_index: &LineIndex,
     line_endings: LineEndings,
 ) {
-    let new_edits = import_into_edits(import_to_add, line_index, line_endings);
+    let import_edits = import_to_add.to_text_edit().map(|import_edit| {
+        import_edit
+            .into_iter()
+            .map(|indel| to_proto::text_edit(line_index, line_endings, indel))
+            .collect_vec()
+    });
     if let Some(original_additional_edits) = completion.additional_text_edits.as_mut() {
-        if let Some(mut new_edits) = new_edits {
+        if let Some(mut new_edits) = import_edits {
             original_additional_edits.extend(new_edits.drain(..))
         }
     } else {
-        completion.additional_text_edits = new_edits;
+        completion.additional_text_edits = import_edits;
     }
 }
-
-fn import_into_edits(
-    import_to_add: &ImportToAdd,
-    line_index: &LineIndex,
-    line_endings: LineEndings,
-) -> Option<Vec<lsp_types::TextEdit>> {
-    let _p = profile::span("add_import_edits");
-
-    let rewriter = insert_use::insert_use(
-        &import_to_add.import_scope,
-        mod_path_to_ast(&import_to_add.import_path),
-        import_to_add.merge_behaviour,
-    );
-    let old_ast = rewriter.rewrite_root()?;
-    let mut import_insert = TextEdit::builder();
-    algo::diff(&old_ast, &rewriter.rewrite(&old_ast)).into_text_edit(&mut import_insert);
-    let import_edit = import_insert.finish();
-
-    Some(
-        import_edit
-            .into_iter()
-            .map(|indel| to_proto::text_edit(line_index, line_endings, indel))
-            .collect_vec(),
-    )
-}
diff --git a/crates/rust-analyzer/src/main_loop.rs b/crates/rust-analyzer/src/main_loop.rs
index 21c58d959..db30b3dce 100644
--- a/crates/rust-analyzer/src/main_loop.rs
+++ b/crates/rust-analyzer/src/main_loop.rs
@@ -438,7 +438,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_resolve_completion,
+                handlers::handle_completion_resolve,
             )?
             .on::<lsp_ext::AnalyzerStatus>(handlers::handle_analyzer_status)
             .on::<lsp_ext::SyntaxTree>(handlers::handle_syntax_tree)