From 52e1e014ad08ecba01843c36f4ff6dac905933c5 Mon Sep 17 00:00:00 2001
From: Antonio Scandurra <me@as-cii.com>
Date: Mon, 28 Aug 2023 14:42:41 +0200
Subject: [PATCH] Allow redoing edits performed by inline assistant after
 cancelling it

---
 crates/ai/src/assistant.rs        |  2 +-
 crates/editor/src/multi_buffer.rs | 24 ++++++++++++++++++------
 crates/language/src/buffer.rs     |  4 ++--
 crates/text/src/text.rs           | 29 +++++++++++++++++++++--------
 4 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/crates/ai/src/assistant.rs b/crates/ai/src/assistant.rs
index 4dad12ad08..952c924292 100644
--- a/crates/ai/src/assistant.rs
+++ b/crates/ai/src/assistant.rs
@@ -442,7 +442,7 @@ impl AssistantPanel {
                     if let Some(transaction_id) = pending_assist.transaction_id {
                         editor.update(cx, |editor, cx| {
                             editor.buffer().update(cx, |buffer, cx| {
-                                buffer.undo_and_forget(transaction_id, cx)
+                                buffer.undo_transaction(transaction_id, cx)
                             });
                         });
                     }
diff --git a/crates/editor/src/multi_buffer.rs b/crates/editor/src/multi_buffer.rs
index c7c4028995..0c499c16c4 100644
--- a/crates/editor/src/multi_buffer.rs
+++ b/crates/editor/src/multi_buffer.rs
@@ -824,13 +824,15 @@ impl MultiBuffer {
         None
     }
 
-    pub fn undo_and_forget(&mut self, transaction_id: TransactionId, cx: &mut ModelContext<Self>) {
+    pub fn undo_transaction(&mut self, transaction_id: TransactionId, cx: &mut ModelContext<Self>) {
         if let Some(buffer) = self.as_singleton() {
-            buffer.update(cx, |buffer, cx| buffer.undo_and_forget(transaction_id, cx));
-        } else if let Some(transaction) = self.history.forget(transaction_id) {
-            for (buffer_id, transaction_id) in transaction.buffer_transactions {
-                if let Some(BufferState { buffer, .. }) = self.buffers.borrow().get(&buffer_id) {
-                    buffer.update(cx, |buffer, cx| buffer.undo_and_forget(transaction_id, cx));
+            buffer.update(cx, |buffer, cx| buffer.undo_transaction(transaction_id, cx));
+        } else if let Some(transaction) = self.history.remove_from_undo(transaction_id) {
+            for (buffer_id, transaction_id) in &transaction.buffer_transactions {
+                if let Some(BufferState { buffer, .. }) = self.buffers.borrow().get(buffer_id) {
+                    buffer.update(cx, |buffer, cx| {
+                        buffer.undo_transaction(*transaction_id, cx)
+                    });
                 }
             }
         }
@@ -3454,6 +3456,16 @@ impl History {
         }
     }
 
+    fn remove_from_undo(&mut self, transaction_id: TransactionId) -> Option<&Transaction> {
+        let ix = self
+            .undo_stack
+            .iter()
+            .rposition(|transaction| transaction.id == transaction_id)?;
+        let transaction = self.undo_stack.remove(ix);
+        self.redo_stack.push(transaction);
+        self.redo_stack.last()
+    }
+
     fn group(&mut self) -> Option<TransactionId> {
         let mut count = 0;
         let mut transactions = self.undo_stack.iter();
diff --git a/crates/language/src/buffer.rs b/crates/language/src/buffer.rs
index 2ed99d8526..4310f84830 100644
--- a/crates/language/src/buffer.rs
+++ b/crates/language/src/buffer.rs
@@ -1668,14 +1668,14 @@ impl Buffer {
         }
     }
 
-    pub fn undo_and_forget(
+    pub fn undo_transaction(
         &mut self,
         transaction_id: TransactionId,
         cx: &mut ModelContext<Self>,
     ) -> bool {
         let was_dirty = self.is_dirty();
         let old_version = self.version.clone();
-        if let Some(operation) = self.text.undo_and_forget(transaction_id) {
+        if let Some(operation) = self.text.undo_transaction(transaction_id) {
             self.send_operation(Operation::Buffer(operation), cx);
             self.did_edit(&old_version, was_dirty, cx);
             true
diff --git a/crates/text/src/text.rs b/crates/text/src/text.rs
index 43bcef0825..6a00ea12db 100644
--- a/crates/text/src/text.rs
+++ b/crates/text/src/text.rs
@@ -264,7 +264,19 @@ impl History {
         }
     }
 
-    fn remove_from_undo(&mut self, transaction_id: TransactionId) -> &[HistoryEntry] {
+    fn remove_from_undo(&mut self, transaction_id: TransactionId) -> Option<&HistoryEntry> {
+        assert_eq!(self.transaction_depth, 0);
+
+        let entry_ix = self
+            .undo_stack
+            .iter()
+            .rposition(|entry| entry.transaction.id == transaction_id)?;
+        let entry = self.undo_stack.remove(entry_ix);
+        self.redo_stack.push(entry);
+        self.redo_stack.last()
+    }
+
+    fn remove_from_undo_until(&mut self, transaction_id: TransactionId) -> &[HistoryEntry] {
         assert_eq!(self.transaction_depth, 0);
 
         let redo_stack_start_len = self.redo_stack.len();
@@ -1207,19 +1219,20 @@ impl Buffer {
         }
     }
 
-    pub fn undo_and_forget(&mut self, transaction_id: TransactionId) -> Option<Operation> {
-        if let Some(transaction) = self.history.forget(transaction_id) {
-            self.undo_or_redo(transaction).log_err()
-        } else {
-            None
-        }
+    pub fn undo_transaction(&mut self, transaction_id: TransactionId) -> Option<Operation> {
+        let transaction = self
+            .history
+            .remove_from_undo(transaction_id)?
+            .transaction
+            .clone();
+        self.undo_or_redo(transaction).log_err()
     }
 
     #[allow(clippy::needless_collect)]
     pub fn undo_to_transaction(&mut self, transaction_id: TransactionId) -> Vec<Operation> {
         let transactions = self
             .history
-            .remove_from_undo(transaction_id)
+            .remove_from_undo_until(transaction_id)
             .iter()
             .map(|entry| entry.transaction.clone())
             .collect::<Vec<_>>();