1 files changed, 54 insertions, 15 deletions

diff --git a/helix-view/src/document.rs b/helix-view/src/document.rs
index 13ffe794..db12fb92 100644
--- a/helix-view/src/document.rs
+++ b/helix-view/src/document.rs
@@ -9,6 +9,7 @@ use helix_core::text_annotations::TextAnnotations;
 use helix_core::Range;
 use helix_vcs::{DiffHandle, DiffProviderRegistry};
 
+use ::parking_lot::Mutex;
 use serde::de::{self, Deserialize, Deserializer};
 use serde::Serialize;
 use std::borrow::Cow;
@@ -18,7 +19,7 @@ use std::fmt::Display;
 use std::future::Future;
 use std::path::{Path, PathBuf};
 use std::str::FromStr;
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 use std::time::SystemTime;
 
 use helix_core::{
@@ -105,6 +106,13 @@ pub struct DocumentSavedEvent {
 pub type DocumentSavedEventResult = Result<DocumentSavedEvent, anyhow::Error>;
 pub type DocumentSavedEventFuture = BoxFuture<'static, DocumentSavedEventResult>;
 
+#[derive(Debug)]
+pub struct SavePoint {
+    /// The view this savepoint is associated with
+    pub view: ViewId,
+    revert: Mutex<Transaction>,
+}
+
 pub struct Document {
     pub(crate) id: DocumentId,
     text: Rope,
@@ -136,7 +144,7 @@ pub struct Document {
     pub history: Cell<History>,
     pub config: Arc<dyn DynAccess<Config>>,
 
-    pub savepoint: Option<Transaction>,
+    savepoints: Vec<Weak<SavePoint>>,
 
     // Last time we wrote to the file. This will carry the time the file was last opened if there
     // were no saves.
@@ -389,7 +397,7 @@ impl Document {
             diagnostics: Vec::new(),
             version: 0,
             history: Cell::new(History::default()),
-            savepoint: None,
+            savepoints: Vec::new(),
             last_saved_time: SystemTime::now(),
             last_saved_revision: 0,
             modified_since_accessed: false,
@@ -846,11 +854,18 @@ impl Document {
             }
 
             // generate revert to savepoint
-            if self.savepoint.is_some() {
-                take_with(&mut self.savepoint, |prev_revert| {
-                    let revert = transaction.invert(&old_doc);
-                    Some(revert.compose(prev_revert.unwrap()))
-                });
+            if !self.savepoints.is_empty() {
+                let revert = transaction.invert(&old_doc);
+                self.savepoints
+                    .retain_mut(|save_point| match save_point.upgrade() {
+                        Some(savepoint) => {
+                            let mut revert_to_savepoint = savepoint.revert.lock();
+                            *revert_to_savepoint =
+                                revert.clone().compose(mem::take(&mut revert_to_savepoint));
+                            true
+                        }
+                        None => false,
+                    })
             }
 
             // update tree-sitter syntax tree
@@ -940,15 +955,39 @@ impl Document {
         self.undo_redo_impl(view, false)
     }
 
-    pub fn savepoint(&mut self) {
-        self.savepoint =
-            Some(Transaction::new(self.text()).with_selection(self.selection(view.id).clone()));
+    /// Creates a reference counted snapshot (called savpepoint) of the document.
+    ///
+    /// The snapshot will remain valid (and updated) idenfinitly as long as ereferences to it exist.
+    /// Restoring the snapshot will restore the selection and the contents of the document to
+    /// the state it had when this function was called.
+    pub fn savepoint(&mut self, view: &View) -> Arc<SavePoint> {
+        let revert = Transaction::new(self.text()).with_selection(self.selection(view.id).clone());
+        let savepoint = Arc::new(SavePoint {
+            view: view.id,
+            revert: Mutex::new(revert),
+        });
+        self.savepoints.push(Arc::downgrade(&savepoint));
+        savepoint
     }
 
-    pub fn restore(&mut self, view: &mut View) {
-        if let Some(revert) = self.savepoint.take() {
-            self.apply(&revert, view.id);
-        }
+    pub fn restore(&mut self, view: &mut View, savepoint: &SavePoint) {
+        assert_eq!(
+            savepoint.view, view.id,
+            "Savepoint must not be used with a different view!"
+        );
+        // search and remove savepoint using a ptr comparison
+        // this avoids a deadlock as we need to lock the mutex
+        let savepoint_idx = self
+            .savepoints
+            .iter()
+            .position(|savepoint_ref| savepoint_ref.as_ptr() == savepoint as *const _)
+            .expect("Savepoint must belong to this document");
+
+        let savepoint_ref = self.savepoints.remove(savepoint_idx);
+        let mut revert = savepoint.revert.lock();
+        self.apply(&revert, view.id);
+        *revert = Transaction::new(self.text()).with_selection(self.selection(view.id).clone());
+        self.savepoints.push(savepoint_ref)
     }
 
     fn earlier_later_impl(&mut self, view: &mut View, uk: UndoKind, earlier: bool) -> bool {