2 files changed, 77 insertions, 48 deletions

diff --git a/crates/ra_assists/src/assist_ctx.rs b/crates/ra_assists/src/assist_ctx.rs
index 81f999090..5924a3fd5 100644
--- a/crates/ra_assists/src/assist_ctx.rs
+++ b/crates/ra_assists/src/assist_ctx.rs
@@ -98,30 +98,8 @@ impl<'a> AssistCtx<'a> {
         Some(assist)
     }
 
-    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,6 +133,67 @@ impl<'a> AssistCtx<'a> {
     }
 }
 
+pub(crate) struct AssistGroup<'a> {
+    ctx: AssistCtx<'a>,
+    group_name: String,
+    assists: Vec<Assist>,
+}
+
+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 assist = if self.ctx.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 }
+        };
+
+        self.assists.push(assist)
+    }
+
+    pub(crate) fn finish(self) -> Option<Assist> {
+        assert!(!self.assists.is_empty());
+        let mut label = match &self.assists[0] {
+            Assist::Unresolved { label } => label.clone(),
+            Assist::Resolved { assist } => assist.label.clone(),
+        };
+        label.label = self.group_name;
+        let assist = if self.ctx.should_compute_edit {
+            Assist::Resolved {
+                assist: ResolvedAssist {
+                    label,
+                    action_data: Either::Right(
+                        self.assists
+                            .into_iter()
+                            .map(|assist| match assist {
+                                Assist::Resolved {
+                                    assist:
+                                        ResolvedAssist { label: _, action_data: Either::Left(it) },
+                                } => it,
+                                _ => unreachable!(),
+                            })
+                            .collect(),
+                    ),
+                },
+            }
+        } else {
+            Assist::Unresolved { label }
+        };
+        Some(assist)
+    }
+}
+
 #[derive(Default)]
 pub(crate) struct ActionBuilder {
     edit: TextEditBuilder,
@@ -164,11 +203,6 @@ pub(crate) struct ActionBuilder {
 }
 
 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())
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)]