3 files changed, 73 insertions, 66 deletions

diff --git a/crates/ra_assists/src/assist_ctx.rs b/crates/ra_assists/src/assist_ctx.rs
index c25d2e323..62182cf03 100644
--- a/crates/ra_assists/src/assist_ctx.rs
+++ b/crates/ra_assists/src/assist_ctx.rs
@@ -166,8 +166,11 @@ impl<'a> AssistGroup<'a> {
     }
 
     pub(crate) fn finish(self) -> Option<Assist> {
-        assert!(!self.assists.is_empty());
-        Some(Assist(self.assists))
+        if self.assists.is_empty() {
+            None
+        } else {
+            Some(Assist(self.assists))
+        }
     }
 }
 
diff --git a/crates/ra_assists/src/handlers/auto_import.rs b/crates/ra_assists/src/handlers/auto_import.rs
index c8bf181f9..4b62aac48 100644
--- a/crates/ra_assists/src/handlers/auto_import.rs
+++ b/crates/ra_assists/src/handlers/auto_import.rs
@@ -42,12 +42,7 @@ pub(crate) fn auto_import(ctx: AssistCtx) -> Option<Assist> {
         return None;
     }
 
-    let assist_group_name = if proposed_imports.len() == 1 {
-        format!("Import `{}`", proposed_imports.iter().next().unwrap())
-    } else {
-        auto_import_assets.get_import_group_message()
-    };
-    let mut group = ctx.add_assist_group(assist_group_name);
+    let mut group = ctx.add_assist_group(auto_import_assets.get_import_group_message());
     for import in proposed_imports {
         group.add_assist(AssistId("auto_import"), format!("Import `{}`", &import), |edit| {
             edit.target(auto_import_assets.syntax_under_caret.text_range());
diff --git a/crates/rust-analyzer/src/main_loop/handlers.rs b/crates/rust-analyzer/src/main_loop/handlers.rs
index f51263f22..b5db1fd38 100644
--- a/crates/rust-analyzer/src/main_loop/handlers.rs
+++ b/crates/rust-analyzer/src/main_loop/handlers.rs
@@ -3,7 +3,6 @@
 //! `ra_ide` crate.
 
 use std::{
-    collections::hash_map::Entry,
     fmt::Write as _,
     io::Write as _,
     process::{self, Stdio},
@@ -13,15 +12,15 @@ use lsp_server::ErrorCode;
 use lsp_types::{
     CallHierarchyIncomingCall, CallHierarchyIncomingCallsParams, CallHierarchyItem,
     CallHierarchyOutgoingCall, CallHierarchyOutgoingCallsParams, CallHierarchyPrepareParams,
-    CodeAction, CodeActionOrCommand, CodeActionResponse, CodeLens, Command, CompletionItem,
-    Diagnostic, DocumentFormattingParams, DocumentHighlight, DocumentSymbol, FoldingRange,
-    FoldingRangeParams, Hover, HoverContents, Location, MarkupContent, MarkupKind, Position,
-    PrepareRenameResponse, Range, RenameParams, SemanticTokens, SemanticTokensParams,
-    SemanticTokensRangeParams, SemanticTokensRangeResult, SemanticTokensResult, SymbolInformation,
-    TextDocumentIdentifier, TextEdit, WorkspaceEdit,
+    CodeAction, CodeActionResponse, CodeLens, Command, CompletionItem, Diagnostic,
+    DocumentFormattingParams, DocumentHighlight, DocumentSymbol, FoldingRange, FoldingRangeParams,
+    Hover, HoverContents, Location, MarkupContent, MarkupKind, Position, PrepareRenameResponse,
+    Range, RenameParams, SemanticTokens, SemanticTokensParams, SemanticTokensRangeParams,
+    SemanticTokensRangeResult, SemanticTokensResult, SymbolInformation, TextDocumentIdentifier,
+    TextEdit, WorkspaceEdit,
 };
 use ra_ide::{
-    AssistId, FileId, FilePosition, FileRange, Query, RangeInfo, Runnable, RunnableKind,
+    Assist, AssistId, FileId, FilePosition, FileRange, Query, RangeInfo, Runnable, RunnableKind,
     SearchScope,
 };
 use ra_prof::profile;
@@ -649,6 +648,31 @@ pub fn handle_formatting(
     }]))
 }
 
+fn create_single_code_action(assist: Assist, world: &WorldSnapshot) -> Result<CodeAction> {
+    let arg = to_value(assist.source_change.try_conv_with(world)?)?;
+    let title = assist.label;
+    let command = Command {
+        title: title.clone(),
+        command: "rust-analyzer.applySourceChange".to_string(),
+        arguments: Some(vec![arg]),
+    };
+
+    let kind = match assist.id {
+        AssistId("introduce_variable") => Some("refactor.extract.variable".to_string()),
+        AssistId("add_custom_impl") => Some("refactor.rewrite.add_custom_impl".to_string()),
+        _ => None,
+    };
+
+    Ok(CodeAction {
+        title,
+        kind,
+        diagnostics: None,
+        edit: None,
+        command: Some(command),
+        is_preferred: None,
+    })
+}
+
 pub fn handle_code_action(
     world: WorldSnapshot,
     params: req::CodeActionParams,
@@ -695,59 +719,44 @@ pub fn handle_code_action(
         res.push(fix.action.clone());
     }
 
-    let mut groups = FxHashMap::default();
+    let mut grouped_assists: FxHashMap<String, Vec<Assist>> = FxHashMap::default();
     for assist in world.analysis().assists(FileRange { file_id, range })?.into_iter() {
-        let arg = to_value(assist.source_change.try_conv_with(&world)?)?;
-
-        let (command, title, arg) = match assist.group_label {
-            None => ("rust-analyzer.applySourceChange", assist.label.clone(), arg),
-
-            // Group all assists with the same `group_label` into a single CodeAction.
-            Some(group_label) => {
-                match groups.entry(group_label.clone()) {
-                    Entry::Occupied(entry) => {
-                        let idx: usize = *entry.get();
-                        match &mut res[idx] {
-                            CodeActionOrCommand::CodeAction(CodeAction {
-                                command: Some(Command { arguments: Some(arguments), .. }),
-                                ..
-                            }) => match arguments.as_mut_slice() {
-                                [serde_json::Value::Array(arguments)] => arguments.push(arg),
-                                _ => panic!("invalid group"),
-                            },
-                            _ => panic!("invalid group"),
-                        }
-                        continue;
-                    }
-                    Entry::Vacant(entry) => {
-                        entry.insert(res.len());
-                    }
-                }
-                ("rust-analyzer.selectAndApplySourceChange", group_label, to_value(vec![arg])?)
-            }
-        };
-
-        let command = Command {
-            title: assist.label.clone(),
-            command: command.to_string(),
-            arguments: Some(vec![arg]),
-        };
+        match &assist.group_label {
+            Some(label) => grouped_assists.entry(label.to_owned()).or_default().push(assist),
+            None => res.push(create_single_code_action(assist, &world)?.into()),
+        }
+    }
 
-        let kind = match assist.id {
-            AssistId("introduce_variable") => Some("refactor.extract.variable".to_string()),
-            AssistId("add_custom_impl") => Some("refactor.rewrite.add_custom_impl".to_string()),
-            _ => None,
-        };
+    for (group_label, assists) in grouped_assists {
+        if assists.len() == 1 {
+            res.push(
+                create_single_code_action(assists.into_iter().next().unwrap(), &world)?.into(),
+            );
+        } else {
+            let title = group_label;
+
+            let mut arguments = Vec::with_capacity(assists.len());
+            for assist in assists {
+                arguments.push(to_value(assist.source_change.try_conv_with(&world)?)?);
+            }
 
-        let action = CodeAction {
-            title,
-            kind,
-            diagnostics: None,
-            edit: None,
-            command: Some(command),
-            is_preferred: None,
-        };
-        res.push(action.into());
+            let command = Some(Command {
+                title: title.clone(),
+                command: "rust-analyzer.selectAndApplySourceChange".to_string(),
+                arguments: Some(vec![serde_json::Value::Array(arguments)]),
+            });
+            res.push(
+                CodeAction {
+                    title,
+                    kind: None,
+                    diagnostics: None,
+                    edit: None,
+                    command,
+                    is_preferred: None,
+                }
+                .into(),
+            );
+        }
     }
 
     Ok(Some(res))