Group modal assistant edits into the same transaction

Co-Authored-By: Kyle Caverly <kyle@zed.dev>

crates/ai/src/refactor.rs

@@ -85,35 +85,51 @@ impl RefactoringAssistant {
                         anyhow::Ok(())
                     });
 
+                    let mut last_transaction = None;
                     while let Some(hunks) = hunks_rx.next().await {
                         editor.update(&mut cx, |editor, cx| {
                             let mut highlights = Vec::new();
 
                             editor.buffer().update(cx, |buffer, cx| {
+                                buffer.finalize_last_transaction(cx);
+
                                 buffer.start_transaction(cx);
-                                for hunk in hunks {
-                                    match hunk {
+                                buffer.edit(
+                                    hunks.into_iter().filter_map(|hunk| match hunk {
                                         crate::diff::Hunk::Insert { text } => {
                                             let edit_start = snapshot.anchor_after(edit_start);
-                                            buffer.edit([(edit_start..edit_start, text)], None, cx);
+                                            Some((edit_start..edit_start, text))
                                         }
                                         crate::diff::Hunk::Remove { len } => {
                                             let edit_end = edit_start + len;
                                             let edit_range = snapshot.anchor_after(edit_start)
                                                 ..snapshot.anchor_before(edit_end);
-                                            buffer.edit([(edit_range, "")], None, cx);
                                             edit_start = edit_end;
+                                            Some((edit_range, String::new()))
                                         }
                                         crate::diff::Hunk::Keep { len } => {
                                             let edit_end = edit_start + len;
                                             let edit_range = snapshot.anchor_after(edit_start)
                                                 ..snapshot.anchor_before(edit_end);
-                                            highlights.push(edit_range);
                                             edit_start += len;
+                                            highlights.push(edit_range);
+                                            None
                                         }
+                                    }),
+                                    None,
+                                    cx,
+                                );
+                                if let Some(transaction) = buffer.end_transaction(cx) {
+                                    if let Some(last_transaction) = last_transaction {
+                                        buffer.merge_transaction_into(
+                                            last_transaction,
+                                            transaction,
+                                            cx,
+                                        );
                                     }
+                                    last_transaction = Some(transaction);
+                                    buffer.finalize_last_transaction(cx);
                                 }
-                                buffer.end_transaction(cx);
                             });
 
                             editor.highlight_text::<Self>(
@@ -199,7 +215,7 @@ impl RefactoringModal {
     fn deploy(workspace: &mut Workspace, _: &Refactor, cx: &mut ViewContext<Workspace>) {
         if let Some(active_editor) = workspace
             .active_item(cx)
-            .and_then(|item| Some(item.downcast::<Editor>()?.downgrade()))
+            .and_then(|item| Some(item.act_as::<Editor>(cx)?.downgrade()))
         {
             workspace.toggle_modal(cx, |_, cx| {
                 let prompt_editor = cx.add_view(|cx| {

crates/editor/src/multi_buffer.rs

@@ -615,6 +615,42 @@ impl MultiBuffer {
         }
     }
 
+    pub fn merge_transaction_into(
+        &mut self,
+        transaction: TransactionId,
+        destination: TransactionId,
+        cx: &mut ModelContext<Self>,
+    ) {
+        if let Some(buffer) = self.as_singleton() {
+            buffer.update(cx, |buffer, _| {
+                buffer.merge_transaction_into(transaction, destination)
+            });
+        } else {
+            if let Some(transaction) = self.history.remove_transaction(transaction) {
+                if let Some(destination) = self.history.transaction_mut(destination) {
+                    for (buffer_id, buffer_transaction_id) in transaction.buffer_transactions {
+                        if let Some(destination_buffer_transaction_id) =
+                            destination.buffer_transactions.get(&buffer_id)
+                        {
+                            if let Some(state) = self.buffers.borrow().get(&buffer_id) {
+                                state.buffer.update(cx, |buffer, _| {
+                                    buffer.merge_transaction_into(
+                                        buffer_transaction_id,
+                                        *destination_buffer_transaction_id,
+                                    )
+                                });
+                            }
+                        } else {
+                            destination
+                                .buffer_transactions
+                                .insert(buffer_id, buffer_transaction_id);
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     pub fn finalize_last_transaction(&mut self, cx: &mut ModelContext<Self>) {
         self.history.finalize_last_transaction();
         for BufferState { buffer, .. } in self.buffers.borrow().values() {
@@ -3333,6 +3369,35 @@ impl History {
         }
     }
 
+    fn remove_transaction(&mut self, transaction_id: TransactionId) -> Option<Transaction> {
+        if let Some(ix) = self
+            .undo_stack
+            .iter()
+            .rposition(|transaction| transaction.id == transaction_id)
+        {
+            Some(self.undo_stack.remove(ix))
+        } else if let Some(ix) = self
+            .redo_stack
+            .iter()
+            .rposition(|transaction| transaction.id == transaction_id)
+        {
+            Some(self.redo_stack.remove(ix))
+        } else {
+            None
+        }
+    }
+
+    fn transaction_mut(&mut self, transaction_id: TransactionId) -> Option<&mut Transaction> {
+        self.undo_stack
+            .iter_mut()
+            .find(|transaction| transaction.id == transaction_id)
+            .or_else(|| {
+                self.redo_stack
+                    .iter_mut()
+                    .find(|transaction| transaction.id == transaction_id)
+            })
+    }
+
     fn pop_undo(&mut self) -> Option<&mut Transaction> {
         assert_eq!(self.transaction_depth, 0);
         if let Some(transaction) = self.undo_stack.pop() {

crates/language/src/buffer.rs

@@ -1294,6 +1294,14 @@ impl Buffer {
         self.text.forget_transaction(transaction_id);
     }
 
+    pub fn merge_transaction_into(
+        &mut self,
+        transaction: TransactionId,
+        destination: TransactionId,
+    ) {
+        self.text.merge_transaction_into(transaction, destination);
+    }
+
     pub fn wait_for_edits(
         &mut self,
         edit_ids: impl IntoIterator<Item = clock::Local>,

crates/text/src/text.rs

@@ -278,20 +278,43 @@ impl History {
         &self.redo_stack[redo_stack_start_len..]
     }
 
-    fn forget(&mut self, transaction_id: TransactionId) {
+    fn forget(&mut self, transaction_id: TransactionId) -> Option<Transaction> {
         assert_eq!(self.transaction_depth, 0);
         if let Some(entry_ix) = self
             .undo_stack
             .iter()
             .rposition(|entry| entry.transaction.id == transaction_id)
         {
-            self.undo_stack.remove(entry_ix);
+            Some(self.undo_stack.remove(entry_ix).transaction)
         } else if let Some(entry_ix) = self
             .redo_stack
             .iter()
             .rposition(|entry| entry.transaction.id == transaction_id)
         {
-            self.undo_stack.remove(entry_ix);
+            Some(self.redo_stack.remove(entry_ix).transaction)
+        } else {
+            None
+        }
+    }
+
+    fn transaction_mut(&mut self, transaction_id: TransactionId) -> Option<&mut Transaction> {
+        let entry = self
+            .undo_stack
+            .iter_mut()
+            .rfind(|entry| entry.transaction.id == transaction_id)
+            .or_else(|| {
+                self.redo_stack
+                    .iter_mut()
+                    .rfind(|entry| entry.transaction.id == transaction_id)
+            })?;
+        Some(&mut entry.transaction)
+    }
+
+    fn merge_transaction_into(&mut self, transaction: TransactionId, destination: TransactionId) {
+        if let Some(transaction) = self.forget(transaction) {
+            if let Some(destination) = self.transaction_mut(destination) {
+                destination.edit_ids.extend(transaction.edit_ids);
+            }
         }
     }
 
@@ -1202,6 +1225,15 @@ impl Buffer {
         self.history.forget(transaction_id);
     }
 
+    pub fn merge_transaction_into(
+        &mut self,
+        transaction: TransactionId,
+        destination: TransactionId,
+    ) {
+        self.history
+            .merge_transaction_into(transaction, destination);
+    }
+
     pub fn redo(&mut self) -> Option<(TransactionId, Operation)> {
         if let Some(entry) = self.history.pop_redo() {
             let transaction = entry.transaction.clone();