Add a way to resolve certain assists

3 files changed, 39 insertions, 16 deletions

diff --git a/crates/ide_assists/src/assist_context.rs b/crates/ide_assists/src/assist_context.rs
index 8714e4978..19e9f179e 100644
--- a/crates/ide_assists/src/assist_context.rs
+++ b/crates/ide_assists/src/assist_context.rs
@@ -19,7 +19,9 @@ use syntax::{
 };
 use text_edit::{TextEdit, TextEditBuilder};
 
-use crate::{assist_config::AssistConfig, Assist, AssistId, AssistKind, GroupLabel};
+use crate::{
+    assist_config::AssistConfig, Assist, AssistId, AssistKind, AssistResolveStrategy, GroupLabel,
+};
 
 /// `AssistContext` allows to apply an assist or check if it could be applied.
 ///
@@ -105,14 +107,14 @@ impl<'a> AssistContext<'a> {
 }
 
 pub(crate) struct Assists {
-    resolve: bool,
     file: FileId,
+    resolve: AssistResolveStrategy,
     buf: Vec<Assist>,
     allowed: Option<Vec<AssistKind>>,
 }
 
 impl Assists {
-    pub(crate) fn new(ctx: &AssistContext, resolve: bool) -> Assists {
+    pub(crate) fn new(ctx: &AssistContext, resolve: AssistResolveStrategy) -> Assists {
         Assists {
             resolve,
             file: ctx.frange.file_id,
@@ -158,7 +160,7 @@ impl Assists {
     }
 
     fn add_impl(&mut self, mut assist: Assist, f: impl FnOnce(&mut AssistBuilder)) -> Option<()> {
-        let source_change = if self.resolve {
+        let source_change = if self.resolve.should_resolve(&assist.id) {
             let mut builder = AssistBuilder::new(self.file);
             f(&mut builder);
             Some(builder.finish())
diff --git a/crates/ide_assists/src/lib.rs b/crates/ide_assists/src/lib.rs
index 88ae5c9a9..397d2a3d0 100644
--- a/crates/ide_assists/src/lib.rs
+++ b/crates/ide_assists/src/lib.rs
@@ -26,7 +26,7 @@ pub(crate) use crate::assist_context::{AssistContext, Assists};
 
 pub use assist_config::AssistConfig;
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum AssistKind {
     // FIXME: does the None variant make sense? Probably not.
     None,
@@ -60,9 +60,27 @@ impl AssistKind {
 
 /// Unique identifier of the assist, should not be shown to the user
 /// directly.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct AssistId(pub &'static str, pub AssistKind);
 
+// TODO kb docs
+#[derive(Debug, Clone, Copy)]
+pub enum AssistResolveStrategy {
+    None,
+    All,
+    Single(AssistId),
+}
+
+impl AssistResolveStrategy {
+    pub fn should_resolve(&self, id: &AssistId) -> bool {
+        match self {
+            AssistResolveStrategy::None => false,
+            AssistResolveStrategy::All => true,
+            AssistResolveStrategy::Single(id_to_resolve) => id_to_resolve == id,
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct GroupLabel(pub String);
 
@@ -91,7 +109,7 @@ impl Assist {
     pub fn get(
         db: &RootDatabase,
         config: &AssistConfig,
-        resolve: bool,
+        resolve: AssistResolveStrategy,
         range: FileRange,
     ) -> Vec<Assist> {
         let sema = Semantics::new(db);
diff --git a/crates/ide_assists/src/tests.rs b/crates/ide_assists/src/tests.rs
index 6f4f97361..227687766 100644
--- a/crates/ide_assists/src/tests.rs
+++ b/crates/ide_assists/src/tests.rs
@@ -12,7 +12,10 @@ use stdx::{format_to, trim_indent};
 use syntax::TextRange;
 use test_utils::{assert_eq_text, extract_offset};
 
-use crate::{handlers::Handler, Assist, AssistConfig, AssistContext, AssistKind, Assists};
+use crate::{
+    handlers::Handler, Assist, AssistConfig, AssistContext, AssistKind, AssistResolveStrategy,
+    Assists,
+};
 
 pub(crate) const TEST_CONFIG: AssistConfig = AssistConfig {
     snippet_cap: SnippetCap::new(true),
@@ -65,14 +68,14 @@ fn check_doc_test(assist_id: &str, before: &str, after: &str) {
     let before = db.file_text(file_id).to_string();
     let frange = FileRange { file_id, range: selection.into() };
 
-    let assist = Assist::get(&db, &TEST_CONFIG, true, frange)
+    let assist = Assist::get(&db, &TEST_CONFIG, AssistResolveStrategy::All, frange)
         .into_iter()
         .find(|assist| assist.id.0 == assist_id)
         .unwrap_or_else(|| {
             panic!(
                 "\n\nAssist is not applicable: {}\nAvailable assists: {}",
                 assist_id,
-                Assist::get(&db, &TEST_CONFIG, false, frange)
+                Assist::get(&db, &TEST_CONFIG, AssistResolveStrategy::None, frange)
                     .into_iter()
                     .map(|assist| assist.id.0)
                     .collect::<Vec<_>>()
@@ -108,7 +111,7 @@ fn check(handler: Handler, before: &str, expected: ExpectedResult, assist_label:
     let sema = Semantics::new(&db);
     let config = TEST_CONFIG;
     let ctx = AssistContext::new(sema, &config, frange);
-    let mut acc = Assists::new(&ctx, true);
+    let mut acc = Assists::new(&ctx, AssistResolveStrategy::All);
     handler(&mut acc, &ctx);
     let mut res = acc.finish();
 
@@ -186,7 +189,7 @@ fn assist_order_field_struct() {
     let (before_cursor_pos, before) = extract_offset(before);
     let (db, file_id) = with_single_file(&before);
     let frange = FileRange { file_id, range: TextRange::empty(before_cursor_pos) };
-    let assists = Assist::get(&db, &TEST_CONFIG, false, frange);
+    let assists = Assist::get(&db, &TEST_CONFIG, AssistResolveStrategy::None, frange);
     let mut assists = assists.iter();
 
     assert_eq!(assists.next().expect("expected assist").label, "Change visibility to pub(crate)");
@@ -211,7 +214,7 @@ pub fn test_some_range(a: int) -> bool {
 "#,
     );
 
-    let assists = Assist::get(&db, &TEST_CONFIG, false, frange);
+    let assists = Assist::get(&db, &TEST_CONFIG, AssistResolveStrategy::None, frange);
     let expected = labels(&assists);
 
     expect![[r#"
@@ -240,7 +243,7 @@ pub fn test_some_range(a: int) -> bool {
         let mut cfg = TEST_CONFIG;
         cfg.allowed = Some(vec![AssistKind::Refactor]);
 
-        let assists = Assist::get(&db, &cfg, false, frange);
+        let assists = Assist::get(&db, &cfg, AssistResolveStrategy::None, frange);
         let expected = labels(&assists);
 
         expect![[r#"
@@ -255,7 +258,7 @@ pub fn test_some_range(a: int) -> bool {
     {
         let mut cfg = TEST_CONFIG;
         cfg.allowed = Some(vec![AssistKind::RefactorExtract]);
-        let assists = Assist::get(&db, &cfg, false, frange);
+        let assists = Assist::get(&db, &cfg, AssistResolveStrategy::None, frange);
         let expected = labels(&assists);
 
         expect![[r#"
@@ -268,7 +271,7 @@ pub fn test_some_range(a: int) -> bool {
     {
         let mut cfg = TEST_CONFIG;
         cfg.allowed = Some(vec![AssistKind::QuickFix]);
-        let assists = Assist::get(&db, &cfg, false, frange);
+        let assists = Assist::get(&db, &cfg, AssistResolveStrategy::None, frange);
         let expected = labels(&assists);
 
         expect![[r#""#]].assert_eq(&expected);