Make undo a bit less weird.

1 files changed, 22 insertions, 17 deletions

diff --git a/src/main.rs b/src/main.rs
index 8999771..5a0d639 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -33,19 +33,20 @@ struct StateHistory {
 
     i: usize,
     history: Vec<State>,
+    curr: State,
 }
 
 impl Deref for StateHistory {
     type Target = State;
 
     fn deref(&self) -> &Self::Target {
-        return &self.history[self.i];
+        return &self.curr;
     }
 }
 
 impl DerefMut for StateHistory {
     fn deref_mut(&mut self) -> &mut Self::Target {
-        return &mut self.history[self.i];
+        return &mut self.curr;
     }
 }
 
@@ -206,10 +207,9 @@ impl StateHistory {
         match (self.mode, mode) {
             (Mode::Normal, Mode::Normal) => (),
             (Mode::Normal, _) => {
-                let curr = self.clone();
                 self.i += 1;
                 self.history.truncate(self.i);
-                self.history.push(curr);
+                self.history.push(self.curr.clone());
             }
             (Mode::Paste, Mode::Paste) => (),
             (Mode::Paste, _) => {
@@ -233,18 +233,20 @@ fn edit(file: &OsStr) -> Result<(), Box<dyn Error>> {
 
     let r = Rope::open(Path::new(file))?;
     let size = term::size()?;
+    let initial_state = State {
+        buf: r,
+        mode: Mode::Normal,
+        row_start: 0,
+        cursor_row: 0,
+        cursor_col: 0,
+        line_offset: 0,
+        line_cols: 0,
+    };
     let mut state = StateHistory {
-        i: 0,
-        history: vec![State {
-            buf: r,
-            mode: Mode::Normal,
-            row_start: 0,
-            cursor_row: 0,
-            cursor_col: 0,
-            line_offset: 0,
-            line_cols: 0,
-        }],
         stdin: BufReader::new(std::io::stdin()),
+        i: 0,
+        history: vec![initial_state.clone()],
+        curr: initial_state,
     };
     let _ = state.repaint_all();
     let _ = write!(stdout(), "\x1b[H");
@@ -268,10 +270,12 @@ fn edit(file: &OsStr) -> Result<(), Box<dyn Error>> {
             }
             Key::CtrlZ => {
                 state.set_mode(Mode::Normal);
-                if state.i == 0 {
-                    continue;
+                if state.i == state.history.len() - 1 {
+                    state.history.push(state.curr.clone());
+                } else if state.i > 0 {
+                    state.i -= 1;
                 }
-                state.i -= 1;
+                state.curr = state.history[state.i].clone();
                 let _ = state.repaint_all();
                 let _ = state.repaint_line_full();
                 let _ = stdout().flush();
@@ -282,6 +286,7 @@ fn edit(file: &OsStr) -> Result<(), Box<dyn Error>> {
                     continue;
                 }
                 state.i += 1;
+                state.curr = state.history[state.i].clone();
                 let _ = state.repaint_all();
                 let _ = state.repaint_line_full();
                 let _ = stdout().flush();