1 files changed, 20 insertions, 18 deletions

diff --git a/crates/assists/src/assist_context.rs b/crates/assists/src/assist_context.rs
index fcfe2d6ee..69499ea32 100644
--- a/crates/assists/src/assist_context.rs
+++ b/crates/assists/src/assist_context.rs
@@ -208,7 +208,7 @@ pub(crate) struct AssistBuilder {
     edit: TextEditBuilder,
     file_id: FileId,
     is_snippet: bool,
-    change: SourceChange,
+    source_file_edits: Vec<SourceFileEdit>,
 }
 
 impl AssistBuilder {
@@ -217,20 +217,27 @@ impl AssistBuilder {
             edit: TextEdit::builder(),
             file_id,
             is_snippet: false,
-            change: SourceChange::default(),
+            source_file_edits: Vec::default(),
         }
     }
 
     pub(crate) fn edit_file(&mut self, file_id: FileId) {
+        self.commit();
         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.change.source_file_edits.iter().any(|it| it.file_id == new_edit.file_id));
-            self.change.source_file_edits.push(new_edit);
+            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 }),
+            }
         }
     }
 
@@ -270,23 +277,18 @@ impl AssistBuilder {
         algo::diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
     }
     pub(crate) fn rewrite(&mut self, rewriter: SyntaxRewriter) {
-        let node = rewriter.rewrite_root().unwrap();
-        let new = rewriter.rewrite(&node);
-        algo::diff(&node, &new).into_text_edit(&mut self.edit);
-    }
-
-    // FIXME: kill this API
-    /// Get access to the raw `TextEditBuilder`.
-    pub(crate) fn text_edit_builder(&mut self) -> &mut TextEditBuilder {
-        &mut self.edit
+        if let Some(node) = rewriter.rewrite_root() {
+            let new = rewriter.rewrite(&node);
+            algo::diff(&node, &new).into_text_edit(&mut self.edit);
+        }
     }
 
     fn finish(mut self) -> SourceChange {
         self.commit();
-        let mut change = mem::take(&mut self.change);
-        if self.is_snippet {
-            change.is_snippet = true;
+        SourceChange {
+            source_file_edits: mem::take(&mut self.source_file_edits),
+            file_system_edits: Default::default(),
+            is_snippet: self.is_snippet,
         }
-        change
     }
 }