From ea502c8665332932b2311df3852b5ac8df6509af Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bla=C5=BE=20Hrastnik?= <blaz@mxxn.io>
Date: Mon, 21 Dec 2020 13:58:54 +0900
Subject: [PATCH] fix change -> change -> undo -> change -> undo -> undo.

---
 helix-term/src/commands.rs | 16 +++++---------
 helix-view/src/document.rs | 44 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index a8aea9141..7722eca57 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -660,21 +660,15 @@ pub fn insert_char_prompt(prompt: &mut Prompt, c: char) {
 
 // Undo / Redo
 
-pub fn undo(cx: &mut Context) {
-    if let Some(revert) = cx.view.doc.history.undo() {
-        cx.view.doc.version += 1;
-        cx.view.doc.apply(&revert);
-        // TODO: undo/redo needs to avoid storing in self.changes/self.old_state
-    }
+// TODO: each command could simply return a Option<transaction>, then the higher level handles
+// storing it?
 
-    // TODO: each command could simply return a Option<transaction>, then the higher level handles storing it?
+pub fn undo(cx: &mut Context) {
+    cx.view.doc.undo();
 }
 
 pub fn redo(cx: &mut Context) {
-    if let Some(transaction) = cx.view.doc.history.redo() {
-        cx.view.doc.version += 1;
-        cx.view.doc.apply(&transaction);
-    }
+    cx.view.doc.redo();
 }
 
 // Yank / Paste
diff --git a/helix-view/src/document.rs b/helix-view/src/document.rs
index 0dc40c5ab..d6246c341 100644
--- a/helix-view/src/document.rs
+++ b/helix-view/src/document.rs
@@ -185,6 +185,50 @@ impl Document {
         success
     }
 
+    pub fn undo(&mut self) -> bool {
+        if let Some(transaction) = self.history.undo() {
+            let old_doc = self.text().clone();
+            self.version += 1;
+            let success = transaction.apply(&mut self.state);
+
+            // update tree-sitter syntax tree
+            if let Some(syntax) = &mut self.syntax {
+                // TODO: no unwrap
+                syntax
+                    .update(&old_doc, &self.state.doc, transaction.changes())
+                    .unwrap();
+            }
+
+            // reset changeset to fix len
+            self.changes = ChangeSet::new(self.text());
+
+            return success;
+        }
+        false
+    }
+
+    pub fn redo(&mut self) -> bool {
+        if let Some(transaction) = self.history.redo() {
+            let old_doc = self.text().clone();
+            self.version += 1;
+            let success = transaction.apply(&mut self.state);
+
+            // update tree-sitter syntax tree
+            if let Some(syntax) = &mut self.syntax {
+                // TODO: no unwrap
+                syntax
+                    .update(&old_doc, &self.state.doc, transaction.changes())
+                    .unwrap();
+            }
+
+            // reset changeset to fix len
+            self.changes = ChangeSet::new(self.text());
+
+            return success;
+        }
+        false
+    }
+
     #[inline]
     pub fn mode(&self) -> Mode {
         self.mode