Make multiple import edits work for completions

17 files changed, 141 insertions, 87 deletions

diff --git a/Cargo.lock b/Cargo.lock
index d1058c02c72..9ad34f71605 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -633,6 +633,7 @@ dependencies = [
  "once_cell",
  "profile",
  "rustc-hash",
+ "smallvec",
  "sourcegen",
  "stdx",
  "syntax",
diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs
index 2c7e1837c3f..c825e4e9cb1 100644
--- a/crates/ide/src/lib.rs
+++ b/crates/ide/src/lib.rs
@@ -533,19 +533,10 @@ impl Analysis {
         &self,
         config: &CompletionConfig,
         position: FilePosition,
-        full_import_path: &str,
-        imported_name: String,
+        imports: impl IntoIterator<Item = (String, String)> + std::panic::UnwindSafe,
     ) -> Cancellable<Vec<TextEdit>> {
         Ok(self
-            .with_db(|db| {
-                ide_completion::resolve_completion_edits(
-                    db,
-                    config,
-                    position,
-                    full_import_path,
-                    imported_name,
-                )
-            })?
+            .with_db(|db| ide_completion::resolve_completion_edits(db, config, position, imports))?
             .unwrap_or_default())
     }
 
diff --git a/crates/ide_completion/Cargo.toml b/crates/ide_completion/Cargo.toml
index 8d910015602..0d4413978d4 100644
--- a/crates/ide_completion/Cargo.toml
+++ b/crates/ide_completion/Cargo.toml
@@ -14,6 +14,7 @@ itertools = "0.10.0"
 rustc-hash = "1.1.0"
 either = "1.6.1"
 once_cell = "1.7"
+smallvec = "1.4"
 
 stdx = { path = "../stdx", version = "0.0.0" }
 syntax = { path = "../syntax", version = "0.0.0" }
diff --git a/crates/ide_completion/src/completions/postfix.rs b/crates/ide_completion/src/completions/postfix.rs
index f83001c22d2..44f2aec51b6 100644
--- a/crates/ide_completion/src/completions/postfix.rs
+++ b/crates/ide_completion/src/completions/postfix.rs
@@ -231,9 +231,8 @@ fn add_custom_postfix_completions(
     let import_scope =
         ImportScope::find_insert_use_container_with_macros(&ctx.token.parent()?, &ctx.sema)?;
     ctx.config.postfix_snippets.iter().for_each(|snippet| {
-        // FIXME: Support multiple imports
-        let import = match snippet.imports(ctx, &import_scope) {
-            Ok(mut imports) => imports.pop(),
+        let imports = match snippet.imports(ctx, &import_scope) {
+            Ok(imports) => imports,
             Err(_) => return,
         };
         let mut builder = postfix_snippet(
@@ -241,7 +240,9 @@ fn add_custom_postfix_completions(
             snippet.description.as_deref().unwrap_or_default(),
             &format!("{}", snippet.snippet(&receiver_text)),
         );
-        builder.add_import(import);
+        for import in imports.into_iter() {
+            builder.add_import(import);
+        }
         builder.add_to(acc);
     });
     None
@@ -480,7 +481,7 @@ fn main() {
                     &["ControlFlow::Break($receiver)".into()],
                     &[],
                     &["core::ops::ControlFlow".into()],
-                    None,
+                    crate::PostfixSnippetScope::Expr,
                 )
                 .unwrap()],
                 ..TEST_CONFIG
diff --git a/crates/ide_completion/src/completions/snippet.rs b/crates/ide_completion/src/completions/snippet.rs
index 9812f25b402..1dace20102c 100644
--- a/crates/ide_completion/src/completions/snippet.rs
+++ b/crates/ide_completion/src/completions/snippet.rs
@@ -104,13 +104,15 @@ fn add_custom_completions(
     let import_scope =
         ImportScope::find_insert_use_container_with_macros(&ctx.token.parent()?, &ctx.sema)?;
     ctx.config.snippets.iter().filter(|snip| snip.scope == scope).for_each(|snip| {
-        // FIXME: Support multiple imports
-        let import = match snip.imports(ctx, &import_scope) {
-            Ok(mut imports) => imports.pop(),
+        let imports = match snip.imports(ctx, &import_scope) {
+            Ok(imports) => imports,
             Err(_) => return,
         };
         let mut builder = snippet(ctx, cap, &snip.label, &snip.snippet);
-        builder.add_import(import).detail(snip.description.as_deref().unwrap_or_default());
+        for import in imports.into_iter() {
+            builder.add_import(import);
+        }
+        builder.detail(snip.description.as_deref().unwrap_or_default());
         builder.add_to(acc);
     });
     None
@@ -132,7 +134,7 @@ mod tests {
                     &["ControlFlow::Break(())".into()],
                     &[],
                     &["core::ops::ControlFlow".into()],
-                    None,
+                    crate::SnippetScope::Expr,
                 )
                 .unwrap()],
                 ..TEST_CONFIG
diff --git a/crates/ide_completion/src/item.rs b/crates/ide_completion/src/item.rs
index 2bc69d5657e..4c75bd69000 100644
--- a/crates/ide_completion/src/item.rs
+++ b/crates/ide_completion/src/item.rs
@@ -11,6 +11,7 @@ use ide_db::{
     },
     SymbolKind,
 };
+use smallvec::SmallVec;
 use stdx::{format_to, impl_from, never};
 use syntax::{algo, TextRange};
 use text_edit::TextEdit;
@@ -76,7 +77,7 @@ pub struct CompletionItem {
     ref_match: Option<Mutability>,
 
     /// The import data to add to completion's edits.
-    import_to_add: Option<ImportEdit>,
+    import_to_add: SmallVec<[ImportEdit; 1]>,
 }
 
 // We use custom debug for CompletionItem to make snapshot tests more readable.
@@ -305,7 +306,7 @@ impl CompletionItem {
             trigger_call_info: None,
             relevance: CompletionRelevance::default(),
             ref_match: None,
-            import_to_add: None,
+            imports_to_add: Default::default(),
         }
     }
 
@@ -364,8 +365,8 @@ impl CompletionItem {
         self.ref_match.map(|mutability| (mutability, relevance))
     }
 
-    pub fn import_to_add(&self) -> Option<&ImportEdit> {
-        self.import_to_add.as_ref()
+    pub fn imports_to_add(&self) -> &[ImportEdit] {
+        &self.import_to_add
     }
 }
 
@@ -398,7 +399,7 @@ impl ImportEdit {
 pub(crate) struct Builder {
     source_range: TextRange,
     completion_kind: CompletionKind,
-    import_to_add: Option<ImportEdit>,
+    imports_to_add: SmallVec<[ImportEdit; 1]>,
     trait_name: Option<String>,
     label: String,
     insert_text: Option<String>,
@@ -422,14 +423,13 @@ impl Builder {
         let mut lookup = self.lookup;
         let mut insert_text = self.insert_text;
 
-        if let Some(original_path) = self
-            .import_to_add
-            .as_ref()
-            .and_then(|import_edit| import_edit.import.original_path.as_ref())
-        {
-            lookup = lookup.or_else(|| Some(label.clone()));
-            insert_text = insert_text.or_else(|| Some(label.clone()));
-            format_to!(label, " (use {})", original_path)
+        if let [import_edit] = &*self.imports_to_add {
+            // snippets can have multiple imports, but normal completions only have up to one
+            if let Some(original_path) = import_edit.import.original_path.as_ref() {
+                lookup = lookup.or_else(|| Some(label.clone()));
+                insert_text = insert_text.or_else(|| Some(label.clone()));
+                format_to!(label, " (use {})", original_path)
+            }
         } else if let Some(trait_name) = self.trait_name {
             insert_text = insert_text.or_else(|| Some(label.clone()));
             format_to!(label, " (as {})", trait_name)
@@ -456,7 +456,7 @@ impl Builder {
             trigger_call_info: self.trigger_call_info.unwrap_or(false),
             relevance: self.relevance,
             ref_match: self.ref_match,
-            import_to_add: self.import_to_add,
+            import_to_add: self.imports_to_add,
         }
     }
     pub(crate) fn lookup_by(&mut self, lookup: impl Into<String>) -> &mut Builder {
@@ -527,8 +527,8 @@ impl Builder {
         self.trigger_call_info = Some(true);
         self
     }
-    pub(crate) fn add_import(&mut self, import_to_add: Option<ImportEdit>) -> &mut Builder {
-        self.import_to_add = import_to_add;
+    pub(crate) fn add_import(&mut self, import_to_add: ImportEdit) -> &mut Builder {
+        self.imports_to_add.push(import_to_add);
         self
     }
     pub(crate) fn ref_match(&mut self, mutability: Mutability) -> &mut Builder {
diff --git a/crates/ide_completion/src/lib.rs b/crates/ide_completion/src/lib.rs
index 56464c07ce2..bbfdadafc72 100644
--- a/crates/ide_completion/src/lib.rs
+++ b/crates/ide_completion/src/lib.rs
@@ -175,8 +175,7 @@ pub fn resolve_completion_edits(
     db: &RootDatabase,
     config: &CompletionConfig,
     position: FilePosition,
-    full_import_path: &str,
-    imported_name: String,
+    imports: impl IntoIterator<Item = (String, String)>,
 ) -> Option<Vec<TextEdit>> {
     let ctx = CompletionContext::new(db, position, config)?;
     let position_for_import = position_for_import(&ctx, None)?;
@@ -185,21 +184,34 @@ pub fn resolve_completion_edits(
     let current_module = ctx.sema.scope(position_for_import).module()?;
     let current_crate = current_module.krate();
 
-    let (import_path, item_to_import) = items_locator::items_with_name(
-        &ctx.sema,
-        current_crate,
-        NameToImport::Exact(imported_name),
-        items_locator::AssocItemSearch::Include,
-        Some(items_locator::DEFAULT_QUERY_SEARCH_LIMIT.inner()),
-    )
-    .filter_map(|candidate| {
-        current_module
-            .find_use_path_prefixed(db, candidate, config.insert_use.prefix_kind)
-            .zip(Some(candidate))
-    })
-    .find(|(mod_path, _)| mod_path.to_string() == full_import_path)?;
-    let import =
-        LocatedImport::new(import_path.clone(), item_to_import, item_to_import, Some(import_path));
+    Some(
+        imports
+            .into_iter()
+            .filter_map(|(full_import_path, imported_name)| {
+                let (import_path, item_to_import) = items_locator::items_with_name(
+                    &ctx.sema,
+                    current_crate,
+                    NameToImport::Exact(imported_name),
+                    items_locator::AssocItemSearch::Include,
+                    Some(items_locator::DEFAULT_QUERY_SEARCH_LIMIT.inner()),
+                )
+                .filter_map(|candidate| {
+                    current_module
+                        .find_use_path_prefixed(db, candidate, config.insert_use.prefix_kind)
+                        .zip(Some(candidate))
+                })
+                .find(|(mod_path, _)| mod_path.to_string() == full_import_path)?;
+                let import = LocatedImport::new(
+                    import_path.clone(),
+                    item_to_import,
+                    item_to_import,
+                    Some(import_path),
+                );
 
-    ImportEdit { import, scope }.to_text_edit(config.insert_use).map(|edit| vec![edit])
+                ImportEdit { import, scope: scope.clone() }
+                    .to_text_edit(config.insert_use)
+                    .map(|edit| edit)
+            })
+            .collect(),
+    )
 }
diff --git a/crates/ide_completion/src/render.rs b/crates/ide_completion/src/render.rs
index 62a5fac5349..58443f566ef 100644
--- a/crates/ide_completion/src/render.rs
+++ b/crates/ide_completion/src/render.rs
@@ -212,7 +212,10 @@ fn render_resolution_(
                 ctx.source_range(),
                 local_name.to_string(),
             );
-            item.kind(CompletionItemKind::UnresolvedReference).add_import(import_to_add);
+            item.kind(CompletionItemKind::UnresolvedReference);
+            if let Some(import_to_add) = import_to_add {
+                item.add_import(import_to_add);
+            }
             return Some(item.build());
         }
     };
@@ -258,9 +261,12 @@ fn render_resolution_(
         }
     }
     item.kind(kind)
-        .add_import(import_to_add)
         .set_documentation(scope_def_docs(ctx.db(), resolution))
         .set_deprecated(scope_def_is_deprecated(&ctx, resolution));
+
+    if let Some(import_to_add) = import_to_add {
+        item.add_import(import_to_add);
+    }
     Some(item.build())
 }
 
diff --git a/crates/ide_completion/src/render/enum_variant.rs b/crates/ide_completion/src/render/enum_variant.rs
index d5cfd8bba46..2ba86eaa0af 100644
--- a/crates/ide_completion/src/render/enum_variant.rs
+++ b/crates/ide_completion/src/render/enum_variant.rs
@@ -68,9 +68,12 @@ impl<'a> EnumRender<'a> {
         item.kind(SymbolKind::Variant)
             .set_documentation(self.variant.docs(self.ctx.db()))
             .set_deprecated(self.ctx.is_deprecated(self.variant))
-            .add_import(import_to_add)
             .detail(self.detail());
 
+        if let Some(import_to_add) = import_to_add {
+            item.add_import(import_to_add);
+        }
+
         if self.variant_kind == hir::StructKind::Tuple {
             cov_mark::hit!(inserts_parens_for_tuple_enums);
             let params = Params::Anonymous(self.variant.fields(self.ctx.db()).len());
diff --git a/crates/ide_completion/src/render/function.rs b/crates/ide_completion/src/render/function.rs
index 95244a758de..4e4663b857a 100644
--- a/crates/ide_completion/src/render/function.rs
+++ b/crates/ide_completion/src/render/function.rs
@@ -98,7 +98,10 @@ impl<'a> FunctionRender<'a> {
             }
         }
 
-        item.add_import(import_to_add).lookup_by(self.name);
+        if let Some(import_to_add) = import_to_add {
+            item.add_import(import_to_add);
+        }
+        item.lookup_by(self.name);
 
         let ret_type = self.func.ret_type(self.ctx.db());
         item.set_relevance(CompletionRelevance {
diff --git a/crates/ide_completion/src/render/macro_.rs b/crates/ide_completion/src/render/macro_.rs
index d1b549df1bb..196b667baac 100644
--- a/crates/ide_completion/src/render/macro_.rs
+++ b/crates/ide_completion/src/render/macro_.rs
@@ -51,9 +51,12 @@ impl<'a> MacroRender<'a> {
         item.kind(SymbolKind::Macro)
             .set_documentation(self.docs.clone())
             .set_deprecated(self.ctx.is_deprecated(self.macro_))
-            .add_import(import_to_add)
             .set_detail(self.detail());
 
+        if let Some(import_to_add) = import_to_add {
+            item.add_import(import_to_add);
+        }
+
         let needs_bang = !(self.ctx.completion.in_use_tree()
             || matches!(self.ctx.completion.path_call_kind(), Some(CallKind::Mac)));
         let has_parens = self.ctx.completion.path_call_kind().is_some();
diff --git a/crates/ide_completion/src/snippet.rs b/crates/ide_completion/src/snippet.rs
index 1bcb128fa9c..8100487a72b 100644
--- a/crates/ide_completion/src/snippet.rs
+++ b/crates/ide_completion/src/snippet.rs
@@ -41,7 +41,7 @@ impl Snippet {
         snippet: &[String],
         description: &[String],
         requires: &[String],
-        scope: Option<SnippetScope>,
+        scope: SnippetScope,
     ) -> Option<Self> {
         // validate that these are indeed simple paths
         if requires.iter().any(|path| match ast::Path::parse(path) {
@@ -57,7 +57,7 @@ impl Snippet {
         let description = description.iter().join("\n");
         let description = if description.is_empty() { None } else { Some(description) };
         Some(Snippet {
-            scope: scope.unwrap_or(SnippetScope::Expr),
+            scope,
             label,
             snippet,
             description,
@@ -89,7 +89,7 @@ impl PostfixSnippet {
         snippet: &[String],
         description: &[String],
         requires: &[String],
-        scope: Option<PostfixSnippetScope>,
+        scope: PostfixSnippetScope,
     ) -> Option<Self> {
         // validate that these are indeed simple paths
         if requires.iter().any(|path| match ast::Path::parse(path) {
@@ -105,7 +105,7 @@ impl PostfixSnippet {
         let description = description.iter().join("\n");
         let description = if description.is_empty() { None } else { Some(description) };
         Some(PostfixSnippet {
-            scope: scope.unwrap_or(PostfixSnippetScope::Expr),
+            scope,
             label,
             snippet,
             description,
diff --git a/crates/ide_completion/src/tests.rs b/crates/ide_completion/src/tests.rs
index 9f5ef2a9aa7..6f4121dd80a 100644
--- a/crates/ide_completion/src/tests.rs
+++ b/crates/ide_completion/src/tests.rs
@@ -183,13 +183,15 @@ pub(crate) fn check_edit_with_config(
     let mut actual = db.file_text(position.file_id).to_string();
 
     let mut combined_edit = completion.text_edit().to_owned();
-    if let Some(import_text_edit) =
-        completion.import_to_add().and_then(|edit| edit.to_text_edit(config.insert_use))
-    {
-        combined_edit.union(import_text_edit).expect(
-            "Failed to apply completion resolve changes: change ranges overlap, but should not",
-        )
-    }
+    completion
+        .imports_to_add()
+        .iter()
+        .filter_map(|edit| edit.to_text_edit(config.insert_use))
+        .for_each(|text_edit| {
+            combined_edit.union(text_edit).expect(
+                "Failed to apply completion resolve changes: change ranges overlap, but should not",
+            )
+        });
 
     combined_edit.apply(&mut actual);
     assert_eq_text!(&ra_fixture_after, &actual)
diff --git a/crates/rust-analyzer/src/config.rs b/crates/rust-analyzer/src/config.rs
index 58539543ee8..150fc8e7026 100644
--- a/crates/rust-analyzer/src/config.rs
+++ b/crates/rust-analyzer/src/config.rs
@@ -462,10 +462,10 @@ impl Config {
                     &desc.snippet,
                     &desc.description,
                     &desc.requires,
-                    desc.scope.map(|scope| match scope {
+                    match desc.scope {
                         PostfixSnippetScopeDef::Expr => PostfixSnippetScope::Expr,
                         PostfixSnippetScopeDef::Type => PostfixSnippetScope::Type,
-                    }),
+                    },
                 )
             })
             .collect();
@@ -479,10 +479,10 @@ impl Config {
                     &desc.snippet,
                     &desc.description,
                     &desc.requires,
-                    desc.scope.map(|scope| match scope {
+                    match desc.scope {
                         SnippetScopeDef::Expr => SnippetScope::Expr,
                         SnippetScopeDef::Item => SnippetScope::Item,
-                    }),
+                    },
                 )
             })
             .collect();
@@ -954,17 +954,31 @@ impl Config {
 }
 
 #[derive(Deserialize, Debug, Clone, Copy)]
+#[serde(rename_all = "snake_case")]
 enum PostfixSnippetScopeDef {
     Expr,
     Type,
 }
 
+impl Default for PostfixSnippetScopeDef {
+    fn default() -> Self {
+        PostfixSnippetScopeDef::Expr
+    }
+}
+
 #[derive(Deserialize, Debug, Clone, Copy)]
+#[serde(rename_all = "snake_case")]
 enum SnippetScopeDef {
     Expr,
     Item,
 }
 
+impl Default for SnippetScopeDef {
+    fn default() -> Self {
+        SnippetScopeDef::Expr
+    }
+}
+
 #[derive(Deserialize, Debug, Clone)]
 struct PostfixSnippetDef {
     #[serde(deserialize_with = "single_or_array")]
@@ -973,7 +987,8 @@ struct PostfixSnippetDef {
     snippet: Vec<String>,
     #[serde(deserialize_with = "single_or_array")]
     requires: Vec<String>,
-    scope: Option<PostfixSnippetScopeDef>,
+    #[serde(default)]
+    scope: PostfixSnippetScopeDef,
 }
 
 #[derive(Deserialize, Debug, Clone)]
@@ -984,7 +999,8 @@ struct SnippetDef {
     snippet: Vec<String>,
     #[serde(deserialize_with = "single_or_array")]
     requires: Vec<String>,
-    scope: Option<SnippetScopeDef>,
+    #[serde(default)]
+    scope: SnippetScopeDef,
 }
 
 fn single_or_array<'de, D>(deserializer: D) -> Result<Vec<String>, D::Error>
diff --git a/crates/rust-analyzer/src/handlers.rs b/crates/rust-analyzer/src/handlers.rs
index e62bb9499fa..6ad5cb53342 100644
--- a/crates/rust-analyzer/src/handlers.rs
+++ b/crates/rust-analyzer/src/handlers.rs
@@ -787,8 +787,10 @@ pub(crate) fn handle_completion_resolve(
         .resolve_completion_edits(
             &snap.config.completion(),
             FilePosition { file_id, offset },
-            &resolve_data.full_import_path,
-            resolve_data.imported_name,
+            resolve_data
+                .imports
+                .into_iter()
+                .map(|import| (import.full_import_path, import.imported_name)),
         )?
         .into_iter()
         .flat_map(|edit| edit.into_iter().map(|indel| to_proto::text_edit(&line_index, indel)))
diff --git a/crates/rust-analyzer/src/lsp_ext.rs b/crates/rust-analyzer/src/lsp_ext.rs
index 521691d5ec9..19137b942eb 100644
--- a/crates/rust-analyzer/src/lsp_ext.rs
+++ b/crates/rust-analyzer/src/lsp_ext.rs
@@ -520,6 +520,11 @@ pub enum WorkspaceSymbolSearchKind {
 #[derive(Debug, Serialize, Deserialize)]
 pub struct CompletionResolveData {
     pub position: lsp_types::TextDocumentPositionParams,
+    pub imports: Vec<CompletionImport>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct CompletionImport {
     pub full_import_path: String,
     pub imported_name: String,
 }
diff --git a/crates/rust-analyzer/src/to_proto.rs b/crates/rust-analyzer/src/to_proto.rs
index 59a768397fe..fc3e25064f3 100644
--- a/crates/rust-analyzer/src/to_proto.rs
+++ b/crates/rust-analyzer/src/to_proto.rs
@@ -270,14 +270,20 @@ fn completion_item(
         lsp_item.insert_text_format = Some(lsp_types::InsertTextFormat::Snippet);
     }
     if config.completion().enable_imports_on_the_fly {
-        if let Some(import_edit) = item.import_to_add() {
-            let import_path = &import_edit.import.import_path;
-            if let Some(import_name) = import_path.segments().last() {
-                let data = lsp_ext::CompletionResolveData {
-                    position: tdpp.clone(),
-                    full_import_path: import_path.to_string(),
-                    imported_name: import_name.to_string(),
-                };
+        if let imports @ [_, ..] = item.imports_to_add() {
+            let imports: Vec<_> = imports
+                .iter()
+                .filter_map(|import_edit| {
+                    let import_path = &import_edit.import.import_path;
+                    let import_name = import_path.segments().last()?;
+                    Some(lsp_ext::CompletionImport {
+                        full_import_path: import_path.to_string(),
+                        imported_name: import_name.to_string(),
+                    })
+                })
+                .collect();
+            if !imports.is_empty() {
+                let data = lsp_ext::CompletionResolveData { position: tdpp.clone(), imports };
                 lsp_item.data = Some(to_value(data).unwrap());
             }
         }