Make multiple import edits work for completions

Cargo.lock

@@ -633,6 +633,7 @@ dependencies = [
  "once_cell",
  "profile",
  "rustc-hash",
+ "smallvec",
  "sourcegen",
  "stdx",
  "syntax",

crates/ide/src/lib.rs

@@ -533,19 +533,10 @@ impl Analysis {
         &self,
         config: &CompletionConfig,
         position: FilePosition,
-        full_import_path: &str,
+        imports: impl IntoIterator<Item = (String, String)> + std::panic::UnwindSafe,
-        imported_name: String,
     ) -> Cancellable<Vec<TextEdit>> {
         Ok(self
-            .with_db(|db| {
+            .with_db(|db| ide_completion::resolve_completion_edits(db, config, position, imports))?
-                ide_completion::resolve_completion_edits(
-                    db,
-                    config,
-                    position,
-                    full_import_path,
-                    imported_name,
-                )
-            })?
             .unwrap_or_default())
     }
 

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" }

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 imports = match snippet.imports(ctx, &import_scope) {
-        let import = match snippet.imports(ctx, &import_scope) {
+            Ok(imports) => imports,
-            Ok(mut imports) => imports.pop(),
             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

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 imports = match snip.imports(ctx, &import_scope) {
-        let import = match snip.imports(ctx, &import_scope) {
+            Ok(imports) => imports,
-            Ok(mut imports) => imports.pop(),
             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

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> {
+    pub fn imports_to_add(&self) -> &[ImportEdit] {
-        self.import_to_add.as_ref()
+        &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
+        if let [import_edit] = &*self.imports_to_add {
-            .import_to_add
+            // snippets can have multiple imports, but normal completions only have up to one
-            .as_ref()
+            if let Some(original_path) = import_edit.import.original_path.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()));
-            lookup = lookup.or_else(|| Some(label.clone()));
+                format_to!(label, " (use {})", original_path)
-            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 {
+    pub(crate) fn add_import(&mut self, import_to_add: ImportEdit) -> &mut Builder {
-        self.import_to_add = import_to_add;
+        self.imports_to_add.push(import_to_add);
         self
     }
     pub(crate) fn ref_match(&mut self, mutability: Mutability) -> &mut Builder {

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,
+    imports: impl IntoIterator<Item = (String, String)>,
-    imported_name: 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(
+    Some(
-        &ctx.sema,
+        imports
-        current_crate,
+            .into_iter()
-        NameToImport::Exact(imported_name),
+            .filter_map(|(full_import_path, imported_name)| {
-        items_locator::AssocItemSearch::Include,
+                let (import_path, item_to_import) = items_locator::items_with_name(
-        Some(items_locator::DEFAULT_QUERY_SEARCH_LIMIT.inner()),
+                    &ctx.sema,
-    )
+                    current_crate,
-    .filter_map(|candidate| {
+                    NameToImport::Exact(imported_name),
-        current_module
+                    items_locator::AssocItemSearch::Include,
-            .find_use_path_prefixed(db, candidate, config.insert_use.prefix_kind)
+                    Some(items_locator::DEFAULT_QUERY_SEARCH_LIMIT.inner()),
-            .zip(Some(candidate))
+                )
-    })
+                .filter_map(|candidate| {
-    .find(|(mod_path, _)| mod_path.to_string() == full_import_path)?;
+                    current_module
-    let import =
+                        .find_use_path_prefixed(db, candidate, config.insert_use.prefix_kind)
-        LocatedImport::new(import_path.clone(), item_to_import, item_to_import, Some(import_path));
+                        .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(),
+    )
 }

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())
 }
 

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());

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 {

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();

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,

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
-        completion.import_to_add().and_then(|edit| edit.to_text_edit(config.insert_use))
+        .imports_to_add()
-    {
+        .iter()
-        combined_edit.union(import_text_edit).expect(
+        .filter_map(|edit| edit.to_text_edit(config.insert_use))
-            "Failed to apply completion resolve changes: change ranges overlap, but should not",
+        .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)

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>

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
-            resolve_data.imported_name,
+                .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)))

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,
 }

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() {
+        if let imports @ [_, ..] = item.imports_to_add() {
-            let import_path = &import_edit.import.import_path;
+            let imports: Vec<_> = imports
-            if let Some(import_name) = import_path.segments().last() {
+                .iter()
-                let data = lsp_ext::CompletionResolveData {
+                .filter_map(|import_edit| {
-                    position: tdpp.clone(),
+                    let import_path = &import_edit.import.import_path;
-                    full_import_path: import_path.to_string(),
+                    let import_name = import_path.segments().last()?;
-                    imported_name: import_name.to_string(),
+                    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());
             }
         }