1 files changed, 38 insertions, 10 deletions

diff --git a/crates/ra_assists/src/assist_context.rs b/crates/ra_assists/src/assist_context.rs
index 3640bb4d2..3407df856 100644
--- a/crates/ra_assists/src/assist_context.rs
+++ b/crates/ra_assists/src/assist_context.rs
@@ -19,7 +19,7 @@ use ra_text_edit::TextEditBuilder;
 
 use crate::{
     assist_config::{AssistConfig, SnippetCap},
-    Assist, AssistId, GroupLabel, ResolvedAssist,
+    Assist, AssistId, AssistKind, GroupLabel, ResolvedAssist,
 };
 
 /// `AssistContext` allows to apply an assist or check if it could be applied.
@@ -103,14 +103,26 @@ pub(crate) struct Assists {
     resolve: bool,
     file: FileId,
     buf: Vec<(Assist, Option<SourceChange>)>,
+    allowed: Option<Vec<AssistKind>>,
 }
 
 impl Assists {
     pub(crate) fn new_resolved(ctx: &AssistContext) -> Assists {
-        Assists { resolve: true, file: ctx.frange.file_id, buf: Vec::new() }
+        Assists {
+            resolve: true,
+            file: ctx.frange.file_id,
+            buf: Vec::new(),
+            allowed: ctx.config.allowed.clone(),
+        }
     }
+
     pub(crate) fn new_unresolved(ctx: &AssistContext) -> Assists {
-        Assists { resolve: false, file: ctx.frange.file_id, buf: Vec::new() }
+        Assists {
+            resolve: false,
+            file: ctx.frange.file_id,
+            buf: Vec::new(),
+            allowed: ctx.config.allowed.clone(),
+        }
     }
 
     pub(crate) fn finish_unresolved(self) -> Vec<Assist> {
@@ -139,9 +151,13 @@ impl Assists {
         target: TextRange,
         f: impl FnOnce(&mut AssistBuilder),
     ) -> Option<()> {
+        if !self.is_allowed(&id) {
+            return None;
+        }
         let label = Assist::new(id, label.into(), None, target);
         self.add_impl(label, f)
     }
+
     pub(crate) fn add_group(
         &mut self,
         group: &GroupLabel,
@@ -150,9 +166,14 @@ impl Assists {
         target: TextRange,
         f: impl FnOnce(&mut AssistBuilder),
     ) -> Option<()> {
+        if !self.is_allowed(&id) {
+            return None;
+        }
+
         let label = Assist::new(id, label.into(), Some(group.clone()), target);
         self.add_impl(label, f)
     }
+
     fn add_impl(&mut self, label: Assist, f: impl FnOnce(&mut AssistBuilder)) -> Option<()> {
         let source_change = if self.resolve {
             let mut builder = AssistBuilder::new(self.file);
@@ -170,13 +191,20 @@ impl Assists {
         self.buf.sort_by_key(|(label, _edit)| label.target.len());
         self.buf
     }
+
+    fn is_allowed(&self, id: &AssistId) -> bool {
+        match &self.allowed {
+            Some(allowed) => allowed.iter().any(|kind| kind.contains(id.1)),
+            None => true,
+        }
+    }
 }
 
 pub(crate) struct AssistBuilder {
     edit: TextEditBuilder,
     file_id: FileId,
     is_snippet: bool,
-    edits: Vec<SourceFileEdit>,
+    change: SourceChange,
 }
 
 impl AssistBuilder {
@@ -185,7 +213,7 @@ impl AssistBuilder {
             edit: TextEditBuilder::default(),
             file_id,
             is_snippet: false,
-            edits: Vec::new(),
+            change: SourceChange::default(),
         }
     }
 
@@ -197,8 +225,8 @@ impl AssistBuilder {
         let edit = mem::take(&mut self.edit).finish();
         if !edit.is_empty() {
             let new_edit = SourceFileEdit { file_id: self.file_id, edit };
-            assert!(!self.edits.iter().any(|it| it.file_id == new_edit.file_id));
-            self.edits.push(new_edit);
+            assert!(!self.change.source_file_edits.iter().any(|it| it.file_id == new_edit.file_id));
+            self.change.source_file_edits.push(new_edit);
         }
     }
 
@@ -265,10 +293,10 @@ impl AssistBuilder {
 
     fn finish(mut self) -> SourceChange {
         self.commit();
-        let mut res: SourceChange = mem::take(&mut self.edits).into();
+        let mut change = mem::take(&mut self.change);
         if self.is_snippet {
-            res.is_snippet = true;
+            change.is_snippet = true;
         }
-        res
+        change
     }
 }