4 files changed, 56 insertions, 107 deletions

diff --git a/crates/ra_assists/src/assist_context.rs b/crates/ra_assists/src/assist_context.rs
index 1925db8b2..edd8255f4 100644
--- a/crates/ra_assists/src/assist_context.rs
+++ b/crates/ra_assists/src/assist_context.rs
@@ -1,5 +1,7 @@
 //! See `AssistContext`
 
+use std::mem;
+
 use algo::find_covering_element;
 use hir::Semantics;
 use ra_db::{FileId, FileRange};
@@ -19,7 +21,6 @@ use crate::{
     assist_config::{AssistConfig, SnippetCap},
     Assist, AssistId, GroupLabel, ResolvedAssist,
 };
-use rustc_hash::FxHashMap;
 
 /// `AssistContext` allows to apply an assist or check if it could be applied.
 ///
@@ -139,16 +140,6 @@ impl Assists {
         let label = Assist::new(id, label.into(), None, target);
         self.add_impl(label, f)
     }
-    pub(crate) fn add_in_multiple_files(
-        &mut self,
-        id: AssistId,
-        label: impl Into<String>,
-        target: TextRange,
-        f: impl FnOnce(&mut AssistDirector),
-    ) -> Option<()> {
-        let label = Assist::new(id, label.into(), None, target);
-        self.add_impl_multiple_files(label, f)
-    }
     pub(crate) fn add_group(
         &mut self,
         group: &GroupLabel,
@@ -173,31 +164,6 @@ impl Assists {
         Some(())
     }
 
-    fn add_impl_multiple_files(
-        &mut self,
-        label: Assist,
-        f: impl FnOnce(&mut AssistDirector),
-    ) -> Option<()> {
-        if !self.resolve {
-            self.buf.push((label, None));
-            return None;
-        }
-        let mut director = AssistDirector::default();
-        f(&mut director);
-        let changes = director.finish();
-        let file_edits: Vec<SourceFileEdit> =
-            changes.into_iter().map(|mut change| change.source_file_edits.pop().unwrap()).collect();
-
-        let source_change = SourceChange {
-            source_file_edits: file_edits,
-            file_system_edits: vec![],
-            is_snippet: false,
-        };
-
-        self.buf.push((label, Some(source_change)));
-        Some(())
-    }
-
     fn finish(mut self) -> Vec<(Assist, Option<SourceChange>)> {
         self.buf.sort_by_key(|(label, _edit)| label.target.len());
         self.buf
@@ -206,13 +172,32 @@ impl Assists {
 
 pub(crate) struct AssistBuilder {
     edit: TextEditBuilder,
-    file: FileId,
+    file_id: FileId,
     is_snippet: bool,
+    edits: Vec<SourceFileEdit>,
 }
 
 impl AssistBuilder {
-    pub(crate) fn new(file: FileId) -> AssistBuilder {
-        AssistBuilder { edit: TextEditBuilder::default(), file, is_snippet: false }
+    pub(crate) fn new(file_id: FileId) -> AssistBuilder {
+        AssistBuilder {
+            edit: TextEditBuilder::default(),
+            file_id,
+            is_snippet: false,
+            edits: Vec::new(),
+        }
+    }
+
+    pub(crate) fn edit_file(&mut self, file_id: FileId) {
+        self.file_id = file_id;
+    }
+
+    fn commit(&mut self) {
+        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);
+        }
     }
 
     /// Remove specified `range` of text.
@@ -270,48 +255,18 @@ impl AssistBuilder {
         algo::diff(&node, &new).into_text_edit(&mut self.edit)
     }
 
-    // FIXME: better API
-    pub(crate) fn set_file(&mut self, assist_file: FileId) {
-        self.file = assist_file;
-    }
-
     // FIXME: kill this API
     /// Get access to the raw `TextEditBuilder`.
     pub(crate) fn text_edit_builder(&mut self) -> &mut TextEditBuilder {
         &mut self.edit
     }
 
-    fn finish(self) -> SourceChange {
-        let edit = self.edit.finish();
-        let source_file_edit = SourceFileEdit { file_id: self.file, edit };
-        let mut res: SourceChange = source_file_edit.into();
+    fn finish(mut self) -> SourceChange {
+        self.commit();
+        let mut res: SourceChange = mem::take(&mut self.edits).into();
         if self.is_snippet {
             res.is_snippet = true;
         }
         res
     }
 }
-
-pub(crate) struct AssistDirector {
-    builders: FxHashMap<FileId, AssistBuilder>,
-}
-
-impl AssistDirector {
-    pub(crate) fn perform(&mut self, file_id: FileId, f: impl FnOnce(&mut AssistBuilder)) {
-        let mut builder = self.builders.entry(file_id).or_insert(AssistBuilder::new(file_id));
-        f(&mut builder);
-    }
-
-    fn finish(self) -> Vec<SourceChange> {
-        self.builders
-            .into_iter()
-            .map(|(_, builder)| builder.finish())
-            .collect::<Vec<SourceChange>>()
-    }
-}
-
-impl Default for AssistDirector {
-    fn default() -> Self {
-        AssistDirector { builders: FxHashMap::default() }
-    }
-}
diff --git a/crates/ra_assists/src/handlers/add_function.rs b/crates/ra_assists/src/handlers/add_function.rs
index 24f931a85..1cfbd75aa 100644
--- a/crates/ra_assists/src/handlers/add_function.rs
+++ b/crates/ra_assists/src/handlers/add_function.rs
@@ -64,7 +64,7 @@ pub(crate) fn add_function(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
     let target = call.syntax().text_range();
     acc.add(AssistId("add_function"), "Add function", target, |builder| {
         let function_template = function_builder.render();
-        builder.set_file(function_template.file);
+        builder.edit_file(function_template.file);
         let new_fn = function_template.to_string(ctx.config.snippet_cap);
         match ctx.config.snippet_cap {
             Some(cap) => builder.insert_snippet(cap, function_template.insert_offset, new_fn),
diff --git a/crates/ra_assists/src/handlers/extract_struct_from_enum_variant.rs b/crates/ra_assists/src/handlers/extract_struct_from_enum_variant.rs
index 2c455a1fd..44db7917a 100644
--- a/crates/ra_assists/src/handlers/extract_struct_from_enum_variant.rs
+++ b/crates/ra_assists/src/handlers/extract_struct_from_enum_variant.rs
@@ -1,20 +1,17 @@
+use hir::{EnumVariant, Module, ModuleDef, Name};
+use ra_db::FileId;
+use ra_fmt::leading_indent;
 use ra_ide_db::{defs::Definition, search::Reference, RootDatabase};
 use ra_syntax::{
     algo::find_node_at_offset,
-    ast::{self, AstNode, NameOwner},
+    ast::{self, ArgListOwner, AstNode, NameOwner, VisibilityOwner},
     SourceFile, SyntaxNode, TextRange, TextSize,
 };
+use rustc_hash::FxHashSet;
 
 use crate::{
-    assist_context::{AssistBuilder, AssistDirector},
-    utils::insert_use_statement,
-    AssistContext, AssistId, Assists,
+    assist_context::AssistBuilder, utils::insert_use_statement, AssistContext, AssistId, Assists,
 };
-use ast::{ArgListOwner, VisibilityOwner};
-use hir::{EnumVariant, Module, ModuleDef, Name};
-use ra_db::FileId;
-use ra_fmt::leading_indent;
-use rustc_hash::FxHashSet;
 
 // Assist: extract_struct_from_enum_variant
 //
@@ -50,11 +47,11 @@ pub(crate) fn extract_struct_from_enum_variant(
     let enum_module_def = ModuleDef::from(enum_hir);
     let current_module = enum_hir.module(ctx.db);
     let target = variant.syntax().text_range();
-    acc.add_in_multiple_files(
+    acc.add(
         AssistId("extract_struct_from_enum_variant"),
         "Extract struct from enum variant",
         target,
-        |edit| {
+        |builder| {
             let definition = Definition::ModuleDef(ModuleDef::EnumVariant(variant_hir));
             let res = definition.find_usages(&ctx.db, None);
             let start_offset = variant.parent_enum().syntax().text_range().start();
@@ -64,7 +61,7 @@ pub(crate) fn extract_struct_from_enum_variant(
                 let source_file = ctx.sema.parse(reference.file_range.file_id);
                 update_reference(
                     ctx,
-                    edit,
+                    builder,
                     reference,
                     &source_file,
                     &enum_module_def,
@@ -73,7 +70,7 @@ pub(crate) fn extract_struct_from_enum_variant(
                 );
             }
             extract_struct_def(
-                edit,
+                builder,
                 enum_ast.syntax(),
                 &variant_name,
                 &field_list.to_string(),
@@ -82,7 +79,7 @@ pub(crate) fn extract_struct_from_enum_variant(
                 &visibility,
             );
             let list_range = field_list.syntax().text_range();
-            update_variant(edit, &variant_name, ctx.frange.file_id, list_range);
+            update_variant(builder, &variant_name, ctx.frange.file_id, list_range);
         },
     )
 }
@@ -115,7 +112,7 @@ fn insert_import(
 }
 
 fn extract_struct_def(
-    edit: &mut AssistDirector,
+    builder: &mut AssistBuilder,
     enum_ast: &SyntaxNode,
     variant_name: &str,
     variant_list: &str,
@@ -142,14 +139,13 @@ fn extract_struct_def(
         list_with_visibility(variant_list),
         indent
     );
-    edit.perform(file_id, |builder| {
-        builder.insert(start_offset, struct_def);
-    });
+    builder.edit_file(file_id);
+    builder.insert(start_offset, struct_def);
     Some(())
 }
 
 fn update_variant(
-    edit: &mut AssistDirector,
+    builder: &mut AssistBuilder,
     variant_name: &str,
     file_id: FileId,
     list_range: TextRange,
@@ -158,15 +154,14 @@ fn update_variant(
         list_range.start().checked_add(TextSize::from(1))?,
         list_range.end().checked_sub(TextSize::from(1))?,
     );
-    edit.perform(file_id, |builder| {
-        builder.replace(inside_variant_range, variant_name);
-    });
+    builder.edit_file(file_id);
+    builder.replace(inside_variant_range, variant_name);
     Some(())
 }
 
 fn update_reference(
     ctx: &AssistContext,
-    edit: &mut AssistDirector,
+    builder: &mut AssistBuilder,
     reference: Reference,
     source_file: &SourceFile,
     enum_module_def: &ModuleDef,
@@ -186,16 +181,15 @@ fn update_reference(
         list_range.start().checked_add(TextSize::from(1))?,
         list_range.end().checked_sub(TextSize::from(1))?,
     );
-    edit.perform(reference.file_range.file_id, |builder| {
-        if !visited_modules_set.contains(&module) {
-            if insert_import(ctx, builder, &path_expr, &module, enum_module_def, variant_hir_name)
-                .is_some()
-            {
-                visited_modules_set.insert(module);
-            }
+    builder.edit_file(reference.file_range.file_id);
+    if !visited_modules_set.contains(&module) {
+        if insert_import(ctx, builder, &path_expr, &module, enum_module_def, variant_hir_name)
+            .is_some()
+        {
+            visited_modules_set.insert(module);
         }
-        builder.replace(inside_list_range, format!("{}{}", segment, list));
-    });
+    }
+    builder.replace(inside_list_range, format!("{}{}", segment, list));
     Some(())
 }
 
diff --git a/crates/ra_assists/src/handlers/fix_visibility.rs b/crates/ra_assists/src/handlers/fix_visibility.rs
index 9ec42f568..531b3560f 100644
--- a/crates/ra_assists/src/handlers/fix_visibility.rs
+++ b/crates/ra_assists/src/handlers/fix_visibility.rs
@@ -63,7 +63,7 @@ fn add_vis_to_referenced_module_def(acc: &mut Assists, ctx: &AssistContext) -> O
     };
 
     acc.add(AssistId("fix_visibility"), assist_label, target, |builder| {
-        builder.set_file(target_file);
+        builder.edit_file(target_file);
         match ctx.config.snippet_cap {
             Some(cap) => builder.insert_snippet(cap, offset, format!("$0{} ", missing_visibility)),
             None => builder.insert(offset, format!("{} ", missing_visibility)),
@@ -106,7 +106,7 @@ fn add_vis_to_referenced_record_field(acc: &mut Assists, ctx: &AssistContext) ->
         format!("Change visibility of {}.{} to {}", parent_name, target_name, missing_visibility);
 
     acc.add(AssistId("fix_visibility"), assist_label, target, |builder| {
-        builder.set_file(target_file);
+        builder.edit_file(target_file);
         match ctx.config.snippet_cap {
             Some(cap) => builder.insert_snippet(cap, offset, format!("$0{} ", missing_visibility)),
             None => builder.insert(offset, format!("{} ", missing_visibility)),