1 files changed, 51 insertions, 9 deletions

diff --git a/crates/ra_assists/src/assist_ctx.rs b/crates/ra_assists/src/assist_ctx.rs
index 28152f724..43f0d664b 100644
--- a/crates/ra_assists/src/assist_ctx.rs
+++ b/crates/ra_assists/src/assist_ctx.rs
@@ -1,5 +1,6 @@
 //! This module defines `AssistCtx` -- the API surface that is exposed to assists.
-use hir::{db::HirDatabase, InFile, SourceAnalyzer};
+use either::Either;
+use hir::{db::HirDatabase, InFile, SourceAnalyzer, SourceBinder};
 use ra_db::FileRange;
 use ra_fmt::{leading_indent, reindent};
 use ra_syntax::{
@@ -9,12 +10,12 @@ use ra_syntax::{
 };
 use ra_text_edit::TextEditBuilder;
 
-use crate::{AssistAction, AssistId, AssistLabel};
+use crate::{AssistAction, AssistId, AssistLabel, ResolvedAssist};
 
 #[derive(Clone, Debug)]
 pub(crate) enum Assist {
     Unresolved { label: AssistLabel },
-    Resolved { label: AssistLabel, action: AssistAction },
+    Resolved { assist: ResolvedAssist },
 }
 
 /// `AssistCtx` allows to apply an assist or check if it could be applied.
@@ -81,16 +82,45 @@ impl<'a, DB: HirDatabase> AssistCtx<'a, DB> {
         self,
         id: AssistId,
         label: impl Into<String>,
-        f: impl FnOnce(&mut AssistBuilder),
+        f: impl FnOnce(&mut ActionBuilder),
     ) -> Option<Assist> {
         let label = AssistLabel { label: label.into(), id };
+        assert!(label.label.chars().nth(0).unwrap().is_uppercase());
+
         let assist = if self.should_compute_edit {
             let action = {
-                let mut edit = AssistBuilder::default();
+                let mut edit = ActionBuilder::default();
                 f(&mut edit);
                 edit.build()
             };
-            Assist::Resolved { label, action }
+            Assist::Resolved { assist: ResolvedAssist { label, action_data: Either::Left(action) } }
+        } else {
+            Assist::Unresolved { label }
+        };
+
+        Some(assist)
+    }
+
+    #[allow(dead_code)] // will be used for auto import assist with multiple actions
+    pub(crate) fn add_assist_group(
+        self,
+        id: AssistId,
+        label: impl Into<String>,
+        f: impl FnOnce() -> Vec<ActionBuilder>,
+    ) -> Option<Assist> {
+        let label = AssistLabel { label: 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 }
         };
@@ -112,12 +142,16 @@ impl<'a, DB: HirDatabase> AssistCtx<'a, DB> {
     pub(crate) fn covering_element(&self) -> SyntaxElement {
         find_covering_element(self.source_file.syntax(), self.frange.range)
     }
+    pub(crate) fn source_binder(&self) -> SourceBinder<'a, DB> {
+        SourceBinder::new(self.db)
+    }
     pub(crate) fn source_analyzer(
         &self,
         node: &SyntaxNode,
         offset: Option<TextUnit>,
     ) -> SourceAnalyzer {
-        SourceAnalyzer::new(self.db, InFile::new(self.frange.file_id.into(), node), offset)
+        let src = InFile::new(self.frange.file_id.into(), node);
+        self.source_binder().analyze(src, offset)
     }
 
     pub(crate) fn covering_node_for_range(&self, range: TextRange) -> SyntaxElement {
@@ -126,13 +160,20 @@ impl<'a, DB: HirDatabase> AssistCtx<'a, DB> {
 }
 
 #[derive(Default)]
-pub(crate) struct AssistBuilder {
+pub(crate) struct ActionBuilder {
     edit: TextEditBuilder,
     cursor_position: Option<TextUnit>,
     target: Option<TextRange>,
+    label: Option<String>,
 }
 
-impl AssistBuilder {
+impl ActionBuilder {
+    #[allow(dead_code)]
+    /// 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())
@@ -191,6 +232,7 @@ impl AssistBuilder {
             edit: self.edit.finish(),
             cursor_position: self.cursor_position,
             target: self.target,
+            label: self.label,
         }
     }
 }