2 files changed, 19 insertions, 34 deletions

diff --git a/crates/assists/src/assist_context.rs b/crates/assists/src/assist_context.rs
index de4b32573..321fe77f3 100644
--- a/crates/assists/src/assist_context.rs
+++ b/crates/assists/src/assist_context.rs
@@ -10,7 +10,7 @@ use ide_db::{
 };
 use ide_db::{
     label::Label,
-    source_change::{FileSystemEdit, SourceChange, SourceFileEdits},
+    source_change::{FileSystemEdit, SourceChange},
     RootDatabase,
 };
 use syntax::{
@@ -180,20 +180,12 @@ impl Assists {
 pub(crate) struct AssistBuilder {
     edit: TextEditBuilder,
     file_id: FileId,
-    is_snippet: bool,
-    source_file_edits: SourceFileEdits,
-    file_system_edits: Vec<FileSystemEdit>,
+    source_change: SourceChange,
 }
 
 impl AssistBuilder {
     pub(crate) fn new(file_id: FileId) -> AssistBuilder {
-        AssistBuilder {
-            edit: TextEdit::builder(),
-            file_id,
-            is_snippet: false,
-            source_file_edits: SourceFileEdits::default(),
-            file_system_edits: Vec::default(),
-        }
+        AssistBuilder { edit: TextEdit::builder(), file_id, source_change: SourceChange::default() }
     }
 
     pub(crate) fn edit_file(&mut self, file_id: FileId) {
@@ -204,7 +196,7 @@ impl AssistBuilder {
     fn commit(&mut self) {
         let edit = mem::take(&mut self.edit).finish();
         if !edit.is_empty() {
-            self.source_file_edits.insert(self.file_id, edit);
+            self.source_change.insert_source_edit(self.file_id, edit);
         }
     }
 
@@ -223,7 +215,7 @@ impl AssistBuilder {
         offset: TextSize,
         snippet: impl Into<String>,
     ) {
-        self.is_snippet = true;
+        self.source_change.is_snippet = true;
         self.insert(offset, snippet);
     }
     /// Replaces specified `range` of text with a given string.
@@ -237,7 +229,7 @@ impl AssistBuilder {
         range: TextRange,
         snippet: impl Into<String>,
     ) {
-        self.is_snippet = true;
+        self.source_change.is_snippet = true;
         self.replace(range, snippet);
     }
     pub(crate) fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
@@ -252,15 +244,11 @@ impl AssistBuilder {
     pub(crate) fn create_file(&mut self, dst: AnchoredPathBuf, content: impl Into<String>) {
         let file_system_edit =
             FileSystemEdit::CreateFile { dst: dst.clone(), initial_contents: content.into() };
-        self.file_system_edits.push(file_system_edit);
+        self.source_change.push_file_system_edit(file_system_edit);
     }
 
     fn finish(mut self) -> SourceChange {
         self.commit();
-        SourceChange {
-            source_file_edits: mem::take(&mut self.source_file_edits),
-            file_system_edits: mem::take(&mut self.file_system_edits),
-            is_snippet: self.is_snippet,
-        }
+        mem::take(&mut self.source_change)
     }
 }
diff --git a/crates/assists/src/tests.rs b/crates/assists/src/tests.rs
index d13d6ad31..71431b406 100644
--- a/crates/assists/src/tests.rs
+++ b/crates/assists/src/tests.rs
@@ -80,7 +80,7 @@ fn check_doc_test(assist_id: &str, before: &str, after: &str) {
     let actual = {
         let source_change = assist.source_change.unwrap();
         let mut actual = before;
-        if let Some(source_file_edit) = source_change.source_file_edits.edits.get(&file_id) {
+        if let Some(source_file_edit) = source_change.get_source_edit(file_id) {
             source_file_edit.apply(&mut actual);
         }
         actual
@@ -120,7 +120,7 @@ fn check(handler: Handler, before: &str, expected: ExpectedResult, assist_label:
                 && source_change.file_system_edits.len() == 0;
 
             let mut buf = String::new();
-            for (file_id, edit) in source_change.source_file_edits.edits {
+            for (file_id, edit) in source_change.source_file_edits {
                 let mut text = db.file_text(file_id).as_ref().to_owned();
                 edit.apply(&mut text);
                 if !skip_header {
@@ -132,18 +132,15 @@ fn check(handler: Handler, before: &str, expected: ExpectedResult, assist_label:
                 buf.push_str(&text);
             }
 
-            for file_system_edit in source_change.file_system_edits.clone() {
-                match file_system_edit {
-                    FileSystemEdit::CreateFile { dst, initial_contents } => {
-                        let sr = db.file_source_root(dst.anchor);
-                        let sr = db.source_root(sr);
-                        let mut base = sr.path_for_file(&dst.anchor).unwrap().clone();
-                        base.pop();
-                        let created_file_path = format!("{}{}", base.to_string(), &dst.path[1..]);
-                        format_to!(buf, "//- {}\n", created_file_path);
-                        buf.push_str(&initial_contents);
-                    }
-                    _ => (),
+            for file_system_edit in source_change.file_system_edits {
+                if let FileSystemEdit::CreateFile { dst, initial_contents } = file_system_edit {
+                    let sr = db.file_source_root(dst.anchor);
+                    let sr = db.source_root(sr);
+                    let mut base = sr.path_for_file(&dst.anchor).unwrap().clone();
+                    base.pop();
+                    let created_file_path = format!("{}{}", base.to_string(), &dst.path[1..]);
+                    format_to!(buf, "//- {}\n", created_file_path);
+                    buf.push_str(&initial_contents);
                 }
             }