From aa598ecb75331068f4cacd443512991bc48937dc Mon Sep 17 00:00:00 2001
From: kjeremy <kjeremy@gmail.com>
Date: Mon, 13 Jul 2020 17:41:47 -0400
Subject: Filter assists

---
 crates/ra_assists/src/assist_context.rs | 38 +++++++++++++++++++++++++++++----
 1 file changed, 34 insertions(+), 4 deletions(-)

(limited to 'crates/ra_assists/src/assist_context.rs')

diff --git a/crates/ra_assists/src/assist_context.rs b/crates/ra_assists/src/assist_context.rs
index c33525363..9ca2cfe68 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.
@@ -57,6 +57,7 @@ pub(crate) struct AssistContext<'a> {
     pub(crate) sema: Semantics<'a, RootDatabase>,
     pub(crate) frange: FileRange,
     source_file: SourceFile,
+    allowed: Option<Vec<AssistKind>>,
 }
 
 impl<'a> AssistContext<'a> {
@@ -64,9 +65,10 @@ impl<'a> AssistContext<'a> {
         sema: Semantics<'a, RootDatabase>,
         config: &'a AssistConfig,
         frange: FileRange,
+        allowed: Option<Vec<AssistKind>>,
     ) -> AssistContext<'a> {
         let source_file = sema.parse(frange.file_id);
-        AssistContext { config, sema, frange, source_file }
+        AssistContext { config, sema, frange, source_file, allowed }
     }
 
     pub(crate) fn db(&self) -> &RootDatabase {
@@ -103,14 +105,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.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.allowed.clone(),
+        }
     }
 
     pub(crate) fn finish_unresolved(self) -> Vec<Assist> {
@@ -139,9 +153,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 +168,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,6 +193,13 @@ 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 {
-- 
cgit v1.2.3