1 files changed, 32 insertions, 17 deletions

diff --git a/crates/ide_assists/src/assist_context.rs b/crates/ide_assists/src/assist_context.rs
index 8714e4978..112939948 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())
@@ -185,7 +187,29 @@ pub(crate) struct AssistBuilder {
     source_change: SourceChange,
 
     /// Maps the original, immutable `SyntaxNode` to a `clone_for_update` twin.
-    mutated_tree: Option<(SyntaxNode, SyntaxNode)>,
+    mutated_tree: Option<TreeMutator>,
+}
+
+pub(crate) struct TreeMutator {
+    immutable: SyntaxNode,
+    mutable_clone: SyntaxNode,
+}
+
+impl TreeMutator {
+    pub(crate) fn new(immutable: &SyntaxNode) -> TreeMutator {
+        let immutable = immutable.ancestors().last().unwrap();
+        let mutable_clone = immutable.clone_for_update();
+        TreeMutator { immutable, mutable_clone }
+    }
+
+    pub(crate) fn make_mut<N: AstNode>(&self, node: &N) -> N {
+        N::cast(self.make_syntax_mut(node.syntax())).unwrap()
+    }
+
+    pub(crate) fn make_syntax_mut(&self, node: &SyntaxNode) -> SyntaxNode {
+        let ptr = SyntaxNodePtr::new(node);
+        ptr.to_node(&self.mutable_clone)
+    }
 }
 
 impl AssistBuilder {
@@ -204,8 +228,8 @@ impl AssistBuilder {
     }
 
     fn commit(&mut self) {
-        if let Some((old, new)) = self.mutated_tree.take() {
-            algo::diff(&old, &new).into_text_edit(&mut self.edit)
+        if let Some(tm) = self.mutated_tree.take() {
+            algo::diff(&tm.immutable, &tm.mutable_clone).into_text_edit(&mut self.edit)
         }
 
         let edit = mem::take(&mut self.edit).finish();
@@ -228,16 +252,7 @@ impl AssistBuilder {
     /// phase, and then get their mutable couterparts using `make_mut` in the
     /// mutable state.
     pub(crate) fn make_mut(&mut self, node: SyntaxNode) -> SyntaxNode {
-        let root = &self
-            .mutated_tree
-            .get_or_insert_with(|| {
-                let immutable = node.ancestors().last().unwrap();
-                let mutable = immutable.clone_for_update();
-                (immutable, mutable)
-            })
-            .1;
-        let ptr = SyntaxNodePtr::new(&&node);
-        ptr.to_node(root)
+        self.mutated_tree.get_or_insert_with(|| TreeMutator::new(&node)).make_syntax_mut(&node)
     }
 
     /// Remove specified `range` of text.