6 files changed, 152 insertions, 156 deletions

diff --git a/crates/ra_assists/src/assist_ctx.rs b/crates/ra_assists/src/assist_ctx.rs
index 81f999090..5aab5fb8b 100644
--- a/crates/ra_assists/src/assist_ctx.rs
+++ b/crates/ra_assists/src/assist_ctx.rs
@@ -1,5 +1,4 @@
 //! This module defines `AssistCtx` -- the API surface that is exposed to assists.
-use either::Either;
 use hir::{InFile, SourceAnalyzer, SourceBinder};
 use ra_db::{FileRange, SourceDatabase};
 use ra_fmt::{leading_indent, reindent};
@@ -11,12 +10,36 @@ use ra_syntax::{
 };
 use ra_text_edit::TextEditBuilder;
 
-use crate::{AssistAction, AssistId, AssistLabel, ResolvedAssist};
+use crate::{AssistAction, AssistId, AssistLabel, GroupLabel, ResolvedAssist};
 
 #[derive(Clone, Debug)]
-pub(crate) enum Assist {
-    Unresolved { label: AssistLabel },
-    Resolved { assist: ResolvedAssist },
+pub(crate) struct Assist(pub(crate) Vec<AssistInfo>);
+
+#[derive(Clone, Debug)]
+pub(crate) struct AssistInfo {
+    pub(crate) label: AssistLabel,
+    pub(crate) group_label: Option<GroupLabel>,
+    pub(crate) action: Option<AssistAction>,
+}
+
+impl AssistInfo {
+    fn new(label: AssistLabel) -> AssistInfo {
+        AssistInfo { label, group_label: None, action: None }
+    }
+
+    fn resolved(self, action: AssistAction) -> AssistInfo {
+        AssistInfo { action: Some(action), ..self }
+    }
+
+    fn with_group(self, group_label: GroupLabel) -> AssistInfo {
+        AssistInfo { group_label: Some(group_label), ..self }
+    }
+
+    pub(crate) fn into_resolved(self) -> Option<ResolvedAssist> {
+        let label = self.label;
+        let group_label = self.group_label;
+        self.action.map(|action| ResolvedAssist { label, group_label, action })
+    }
 }
 
 pub(crate) type AssistHandler = fn(AssistCtx) -> Option<Assist>;
@@ -84,44 +107,21 @@ impl<'a> AssistCtx<'a> {
     ) -> Option<Assist> {
         let label = AssistLabel::new(label.into(), id);
 
-        let assist = if self.should_compute_edit {
+        let mut info = AssistInfo::new(label);
+        if self.should_compute_edit {
             let action = {
                 let mut edit = ActionBuilder::default();
                 f(&mut edit);
                 edit.build()
             };
-            Assist::Resolved { assist: ResolvedAssist { label, action_data: Either::Left(action) } }
-        } else {
-            Assist::Unresolved { label }
+            info = info.resolved(action)
         };
 
-        Some(assist)
+        Some(Assist(vec![info]))
     }
 
-    pub(crate) fn add_assist_group(
-        self,
-        id: AssistId,
-        label: impl Into<String>,
-        f: impl FnOnce() -> Vec<ActionBuilder>,
-    ) -> Option<Assist> {
-        let label = AssistLabel::new(label.into(), id);
-        let assist = if self.should_compute_edit {
-            let actions = f();
-            assert!(!actions.is_empty(), "Assist cannot have no");
-
-            Assist::Resolved {
-                assist: ResolvedAssist {
-                    label,
-                    action_data: Either::Right(
-                        actions.into_iter().map(ActionBuilder::build).collect(),
-                    ),
-                },
-            }
-        } else {
-            Assist::Unresolved { label }
-        };
-
-        Some(assist)
+    pub(crate) fn add_assist_group(self, group_name: impl Into<String>) -> AssistGroup<'a> {
+        AssistGroup { ctx: self, group_name: group_name.into(), assists: Vec::new() }
     }
 
     pub(crate) fn token_at_offset(&self) -> TokenAtOffset<SyntaxToken> {
@@ -155,20 +155,48 @@ impl<'a> AssistCtx<'a> {
     }
 }
 
+pub(crate) struct AssistGroup<'a> {
+    ctx: AssistCtx<'a>,
+    group_name: String,
+    assists: Vec<AssistInfo>,
+}
+
+impl<'a> AssistGroup<'a> {
+    pub(crate) fn add_assist(
+        &mut self,
+        id: AssistId,
+        label: impl Into<String>,
+        f: impl FnOnce(&mut ActionBuilder),
+    ) {
+        let label = AssistLabel::new(label.into(), id);
+
+        let mut info = AssistInfo::new(label).with_group(GroupLabel(self.group_name.clone()));
+        if self.ctx.should_compute_edit {
+            let action = {
+                let mut edit = ActionBuilder::default();
+                f(&mut edit);
+                edit.build()
+            };
+            info = info.resolved(action)
+        };
+
+        self.assists.push(info)
+    }
+
+    pub(crate) fn finish(self) -> Option<Assist> {
+        assert!(!self.assists.is_empty());
+        Some(Assist(self.assists))
+    }
+}
+
 #[derive(Default)]
 pub(crate) struct ActionBuilder {
     edit: TextEditBuilder,
     cursor_position: Option<TextUnit>,
     target: Option<TextRange>,
-    label: Option<String>,
 }
 
 impl ActionBuilder {
-    /// Adds a custom label to the action, if it needs to be different from the assist label
-    pub(crate) fn label(&mut self, label: impl Into<String>) {
-        self.label = Some(label.into())
-    }
-
     /// Replaces specified `range` of text with a given string.
     pub(crate) fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
         self.edit.replace(range, replace_with.into())
@@ -227,7 +255,6 @@ impl ActionBuilder {
             edit: self.edit.finish(),
             cursor_position: self.cursor_position,
             target: self.target,
-            label: self.label,
         }
     }
 }
diff --git a/crates/ra_assists/src/doc_tests.rs b/crates/ra_assists/src/doc_tests.rs
index ae0e5605c..c0f9bc1fb 100644
--- a/crates/ra_assists/src/doc_tests.rs
+++ b/crates/ra_assists/src/doc_tests.rs
@@ -30,6 +30,6 @@ fn check(assist_id: &str, before: &str, after: &str) {
             )
         });
 
-    let actual = assist.get_first_action().edit.apply(&before);
+    let actual = assist.action.edit.apply(&before);
     assert_eq_text!(after, &actual);
 }
diff --git a/crates/ra_assists/src/handlers/auto_import.rs b/crates/ra_assists/src/handlers/auto_import.rs
index 4514b8691..d13332f37 100644
--- a/crates/ra_assists/src/handlers/auto_import.rs
+++ b/crates/ra_assists/src/handlers/auto_import.rs
@@ -1,12 +1,8 @@
-use hir::ModPath;
 use ra_ide_db::imports_locator::ImportsLocator;
-use ra_syntax::{
-    ast::{self, AstNode},
-    SyntaxNode,
-};
+use ra_syntax::ast::{self, AstNode};
 
 use crate::{
-    assist_ctx::{ActionBuilder, Assist, AssistCtx},
+    assist_ctx::{Assist, AssistCtx},
     insert_use_statement, AssistId,
 };
 use std::collections::BTreeSet;
@@ -67,19 +63,18 @@ pub(crate) fn auto_import(ctx: AssistCtx) -> Option<Assist> {
         return None;
     }
 
-    ctx.add_assist_group(AssistId("auto_import"), format!("Import {}", name_to_import), || {
-        proposed_imports
-            .into_iter()
-            .map(|import| import_to_action(import, &position, &path_to_import_syntax))
-            .collect()
-    })
-}
-
-fn import_to_action(import: ModPath, position: &SyntaxNode, anchor: &SyntaxNode) -> ActionBuilder {
-    let mut action_builder = ActionBuilder::default();
-    action_builder.label(format!("Import `{}`", &import));
-    insert_use_statement(position, anchor, &import, action_builder.text_edit_builder());
-    action_builder
+    let mut group = ctx.add_assist_group(format!("Import {}", name_to_import));
+    for import in proposed_imports {
+        group.add_assist(AssistId("auto_import"), format!("Import `{}`", &import), |edit| {
+            insert_use_statement(
+                &position,
+                path_to_import_syntax,
+                &import,
+                edit.text_edit_builder(),
+            );
+        });
+    }
+    group.finish()
 }
 
 #[cfg(test)]
diff --git a/crates/ra_assists/src/lib.rs b/crates/ra_assists/src/lib.rs
index f79189ae8..828a8e9e8 100644
--- a/crates/ra_assists/src/lib.rs
+++ b/crates/ra_assists/src/lib.rs
@@ -12,9 +12,6 @@ mod doc_tests;
 mod utils;
 pub mod ast_transform;
 
-use std::cmp::Ordering;
-
-use either::Either;
 use ra_db::FileRange;
 use ra_ide_db::RootDatabase;
 use ra_syntax::{TextRange, TextUnit};
@@ -35,6 +32,9 @@ pub struct AssistLabel {
     pub id: AssistId,
 }
 
+#[derive(Clone, Debug)]
+pub struct GroupLabel(pub String);
+
 impl AssistLabel {
     pub(crate) fn new(label: String, id: AssistId) -> AssistLabel {
         // FIXME: make fields private, so that this invariant can't be broken
@@ -45,7 +45,6 @@ impl AssistLabel {
 
 #[derive(Debug, Clone)]
 pub struct AssistAction {
-    pub label: Option<String>,
     pub edit: TextEdit,
     pub cursor_position: Option<TextUnit>,
     // FIXME: This belongs to `AssistLabel`
@@ -55,16 +54,8 @@ pub struct AssistAction {
 #[derive(Debug, Clone)]
 pub struct ResolvedAssist {
     pub label: AssistLabel,
-    pub action_data: Either<AssistAction, Vec<AssistAction>>,
-}
-
-impl ResolvedAssist {
-    pub fn get_first_action(&self) -> AssistAction {
-        match &self.action_data {
-            Either::Left(action) => action.clone(),
-            Either::Right(actions) => actions[0].clone(),
-        }
-    }
+    pub group_label: Option<GroupLabel>,
+    pub action: AssistAction,
 }
 
 /// Return all the assists applicable at the given position.
@@ -76,10 +67,8 @@ pub fn unresolved_assists(db: &RootDatabase, range: FileRange) -> Vec<AssistLabe
     handlers::all()
         .iter()
         .filter_map(|f| f(ctx.clone()))
-        .map(|a| match a {
-            Assist::Unresolved { label } => label,
-            Assist::Resolved { .. } => unreachable!(),
-        })
+        .flat_map(|it| it.0)
+        .map(|a| a.label)
         .collect()
 }
 
@@ -92,24 +81,13 @@ pub fn resolved_assists(db: &RootDatabase, range: FileRange) -> Vec<ResolvedAssi
     let mut a = handlers::all()
         .iter()
         .filter_map(|f| f(ctx.clone()))
-        .map(|a| match a {
-            Assist::Resolved { assist } => assist,
-            Assist::Unresolved { .. } => unreachable!(),
-        })
+        .flat_map(|it| it.0)
+        .map(|it| it.into_resolved().unwrap())
         .collect::<Vec<_>>();
-    sort_assists(&mut a);
+    a.sort_by_key(|it| it.action.target.map_or(TextUnit::from(!0u32), |it| it.len()));
     a
 }
 
-fn sort_assists(assists: &mut [ResolvedAssist]) {
-    assists.sort_by(|a, b| match (a.get_first_action().target, b.get_first_action().target) {
-        (Some(a), Some(b)) => a.len().cmp(&b.len()),
-        (Some(_), None) => Ordering::Less,
-        (None, Some(_)) => Ordering::Greater,
-        (None, None) => Ordering::Equal,
-    });
-}
-
 mod handlers {
     use crate::AssistHandler;
 
@@ -184,7 +162,7 @@ mod helpers {
     use ra_syntax::TextRange;
     use test_utils::{add_cursor, assert_eq_text, extract_offset, extract_range};
 
-    use crate::{Assist, AssistCtx, AssistHandler};
+    use crate::{AssistCtx, AssistHandler};
 
     pub(crate) fn with_single_file(text: &str) -> (RootDatabase, FileId) {
         let (mut db, file_id) = RootDatabase::with_single_file(text);
@@ -202,10 +180,7 @@ mod helpers {
             FileRange { file_id, range: TextRange::offset_len(before_cursor_pos, 0.into()) };
         let assist =
             assist(AssistCtx::new(&db, frange, true)).expect("code action is not applicable");
-        let action = match assist {
-            Assist::Unresolved { .. } => unreachable!(),
-            Assist::Resolved { assist } => assist.get_first_action(),
-        };
+        let action = assist.0[0].action.clone().unwrap();
 
         let actual = action.edit.apply(&before);
         let actual_cursor_pos = match action.cursor_position {
@@ -225,10 +200,7 @@ mod helpers {
         let frange = FileRange { file_id, range };
         let assist =
             assist(AssistCtx::new(&db, frange, true)).expect("code action is not applicable");
-        let action = match assist {
-            Assist::Unresolved { .. } => unreachable!(),
-            Assist::Resolved { assist } => assist.get_first_action(),
-        };
+        let action = assist.0[0].action.clone().unwrap();
 
         let mut actual = action.edit.apply(&before);
         if let Some(pos) = action.cursor_position {
@@ -244,10 +216,7 @@ mod helpers {
             FileRange { file_id, range: TextRange::offset_len(before_cursor_pos, 0.into()) };
         let assist =
             assist(AssistCtx::new(&db, frange, true)).expect("code action is not applicable");
-        let action = match assist {
-            Assist::Unresolved { .. } => unreachable!(),
-            Assist::Resolved { assist } => assist.get_first_action(),
-        };
+        let action = assist.0[0].action.clone().unwrap();
 
         let range = action.target.expect("expected target on action");
         assert_eq_text!(&before[range.start().to_usize()..range.end().to_usize()], target);
@@ -259,10 +228,7 @@ mod helpers {
         let frange = FileRange { file_id, range };
         let assist =
             assist(AssistCtx::new(&db, frange, true)).expect("code action is not applicable");
-        let action = match assist {
-            Assist::Unresolved { .. } => unreachable!(),
-            Assist::Resolved { assist } => assist.get_first_action(),
-        };
+        let action = assist.0[0].action.clone().unwrap();
 
         let range = action.target.expect("expected target on action");
         assert_eq_text!(&before[range.start().to_usize()..range.end().to_usize()], target);
diff --git a/crates/ra_ide/src/assists.rs b/crates/ra_ide/src/assists.rs
index b60b1a60d..40d56a4f7 100644
--- a/crates/ra_ide/src/assists.rs
+++ b/crates/ra_ide/src/assists.rs
@@ -1,6 +1,5 @@
 //! FIXME: write short doc here
 
-use either::Either;
 use ra_assists::{resolved_assists, AssistAction, AssistLabel};
 use ra_db::{FilePosition, FileRange};
 use ra_ide_db::RootDatabase;
@@ -13,7 +12,8 @@ pub use ra_assists::AssistId;
 pub struct Assist {
     pub id: AssistId,
     pub label: String,
-    pub change_data: Either<SourceChange, Vec<SourceChange>>,
+    pub group_label: Option<String>,
+    pub source_change: SourceChange,
 }
 
 pub(crate) fn assists(db: &RootDatabase, frange: FileRange) -> Vec<Assist> {
@@ -25,17 +25,8 @@ pub(crate) fn assists(db: &RootDatabase, frange: FileRange) -> Vec<Assist> {
             Assist {
                 id: assist_label.id,
                 label: assist_label.label.clone(),
-                change_data: match assist.action_data {
-                    Either::Left(action) => {
-                        Either::Left(action_to_edit(action, file_id, assist_label))
-                    }
-                    Either::Right(actions) => Either::Right(
-                        actions
-                            .into_iter()
-                            .map(|action| action_to_edit(action, file_id, assist_label))
-                            .collect(),
-                    ),
-                },
+                group_label: assist.group_label.map(|it| it.0),
+                source_change: action_to_edit(assist.action, file_id, assist_label),
             }
         })
         .collect()
@@ -47,9 +38,6 @@ fn action_to_edit(
     assist_label: &AssistLabel,
 ) -> SourceChange {
     let file_edit = SourceFileEdit { file_id, edit: action.edit };
-    SourceChange::source_file_edit(
-        action.label.unwrap_or_else(|| assist_label.label.clone()),
-        file_edit,
-    )
-    .with_cursor_opt(action.cursor_position.map(|offset| FilePosition { offset, file_id }))
+    SourceChange::source_file_edit(assist_label.label.clone(), file_edit)
+        .with_cursor_opt(action.cursor_position.map(|offset| FilePosition { offset, file_id }))
 }
diff --git a/crates/ra_lsp_server/src/main_loop/handlers.rs b/crates/ra_lsp_server/src/main_loop/handlers.rs
index 282f6e8fc..65e8bc856 100644
--- a/crates/ra_lsp_server/src/main_loop/handlers.rs
+++ b/crates/ra_lsp_server/src/main_loop/handlers.rs
@@ -2,20 +2,21 @@
 //! The majority of requests are fulfilled by calling into the `ra_ide` crate.
 
 use std::{
+    collections::hash_map::Entry,
     fmt::Write as _,
     io::Write as _,
     process::{self, Stdio},
 };
 
-use either::Either;
 use lsp_server::ErrorCode;
 use lsp_types::{
     CallHierarchyIncomingCall, CallHierarchyIncomingCallsParams, CallHierarchyItem,
     CallHierarchyOutgoingCall, CallHierarchyOutgoingCallsParams, CallHierarchyPrepareParams,
-    CodeAction, CodeActionResponse, CodeLens, Command, CompletionItem, Diagnostic,
-    DocumentFormattingParams, DocumentHighlight, DocumentSymbol, FoldingRange, FoldingRangeParams,
-    Hover, HoverContents, Location, MarkupContent, MarkupKind, Position, PrepareRenameResponse,
-    Range, RenameParams, SymbolInformation, TextDocumentIdentifier, TextEdit, WorkspaceEdit,
+    CodeAction, CodeActionOrCommand, CodeActionResponse, CodeLens, Command, CompletionItem,
+    Diagnostic, DocumentFormattingParams, DocumentHighlight, DocumentSymbol, FoldingRange,
+    FoldingRangeParams, Hover, HoverContents, Location, MarkupContent, MarkupKind, Position,
+    PrepareRenameResponse, Range, RenameParams, SymbolInformation, TextDocumentIdentifier,
+    TextEdit, WorkspaceEdit,
 };
 use ra_ide::{
     AssistId, FileId, FilePosition, FileRange, Query, RangeInfo, Runnable, RunnableKind,
@@ -685,34 +686,53 @@ pub fn handle_code_action(
         res.push(fix.action.clone());
     }
 
+    let mut groups = FxHashMap::default();
     for assist in world.analysis().assists(FileRange { file_id, range })?.into_iter() {
-        let title = assist.label.clone();
+        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 = match assist.change_data {
-            Either::Left(change) => Command {
-                title,
-                command: "rust-analyzer.applySourceChange".to_string(),
-                arguments: Some(vec![to_value(change.try_conv_with(&world)?)?]),
-            },
-            Either::Right(changes) => Command {
-                title,
-                command: "rust-analyzer.selectAndApplySourceChange".to_string(),
-                arguments: Some(vec![to_value(
-                    changes
-                        .into_iter()
-                        .map(|change| change.try_conv_with(&world))
-                        .collect::<Result<Vec<_>>>()?,
-                )?]),
-            },
+        let command = Command {
+            title: assist.label.clone(),
+            command: command.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,
         };
 
         let action = CodeAction {
-            title: command.title.clone(),
-            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,
-            },
+            title,
+            kind,
             diagnostics: None,
             edit: None,
             command: Some(command),