2 files changed, 12 insertions, 23 deletions

diff --git a/crates/assists/src/assist_context.rs b/crates/assists/src/assist_context.rs
index 91cc63427..de4b32573 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, SourceFileEdit},
+    source_change::{FileSystemEdit, SourceChange, SourceFileEdits},
     RootDatabase,
 };
 use syntax::{
@@ -181,7 +181,7 @@ pub(crate) struct AssistBuilder {
     edit: TextEditBuilder,
     file_id: FileId,
     is_snippet: bool,
-    source_file_edits: Vec<SourceFileEdit>,
+    source_file_edits: SourceFileEdits,
     file_system_edits: Vec<FileSystemEdit>,
 }
 
@@ -191,7 +191,7 @@ impl AssistBuilder {
             edit: TextEdit::builder(),
             file_id,
             is_snippet: false,
-            source_file_edits: Vec::default(),
+            source_file_edits: SourceFileEdits::default(),
             file_system_edits: Vec::default(),
         }
     }
@@ -204,15 +204,7 @@ impl AssistBuilder {
     fn commit(&mut self) {
         let edit = mem::take(&mut self.edit).finish();
         if !edit.is_empty() {
-            match self.source_file_edits.binary_search_by_key(&self.file_id, |edit| edit.file_id) {
-                Ok(idx) => self.source_file_edits[idx]
-                    .edit
-                    .union(edit)
-                    .expect("overlapping edits for same file"),
-                Err(idx) => self
-                    .source_file_edits
-                    .insert(idx, SourceFileEdit { file_id: self.file_id, edit }),
-            }
+            self.source_file_edits.insert(self.file_id, edit);
         }
     }
 
diff --git a/crates/assists/src/tests.rs b/crates/assists/src/tests.rs
index fef29a0b8..d13d6ad31 100644
--- a/crates/assists/src/tests.rs
+++ b/crates/assists/src/tests.rs
@@ -80,10 +80,8 @@ fn check_doc_test(assist_id: &str, before: &str, after: &str) {
     let actual = {
         let source_change = assist.source_change.unwrap();
         let mut actual = before;
-        for source_file_edit in source_change.source_file_edits {
-            if source_file_edit.file_id == file_id {
-                source_file_edit.edit.apply(&mut actual)
-            }
+        if let Some(source_file_edit) = source_change.source_file_edits.edits.get(&file_id) {
+            source_file_edit.apply(&mut actual);
         }
         actual
     };
@@ -116,20 +114,19 @@ fn check(handler: Handler, before: &str, expected: ExpectedResult, assist_label:
 
     match (assist, expected) {
         (Some(assist), ExpectedResult::After(after)) => {
-            let mut source_change = assist.source_change.unwrap();
+            let source_change = assist.source_change.unwrap();
             assert!(!source_change.source_file_edits.is_empty());
             let skip_header = source_change.source_file_edits.len() == 1
                 && source_change.file_system_edits.len() == 0;
-            source_change.source_file_edits.sort_by_key(|it| it.file_id);
 
             let mut buf = String::new();
-            for source_file_edit in source_change.source_file_edits {
-                let mut text = db.file_text(source_file_edit.file_id).as_ref().to_owned();
-                source_file_edit.edit.apply(&mut text);
+            for (file_id, edit) in source_change.source_file_edits.edits {
+                let mut text = db.file_text(file_id).as_ref().to_owned();
+                edit.apply(&mut text);
                 if !skip_header {
-                    let sr = db.file_source_root(source_file_edit.file_id);
+                    let sr = db.file_source_root(file_id);
                     let sr = db.source_root(sr);
-                    let path = sr.path_for_file(&source_file_edit.file_id).unwrap();
+                    let path = sr.path_for_file(&file_id).unwrap();
                     format_to!(buf, "//- {}\n", path)
                 }
                 buf.push_str(&text);