2 files changed, 46 insertions, 12 deletions

diff --git a/src/main.rs b/src/main.rs
index a4fe9ae..c1131f1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -15,6 +15,7 @@ enum Mode {
     Normal,
 }
 
+#[derive(Debug, Clone)]
 struct State {
     buf: Rope,
     mode: Mode,
@@ -23,6 +24,7 @@ struct State {
     cursor_col: usize,
     line_offset: usize,
     line_cols: usize,
+    prev: Option<Box<State>>,
 }
 
 fn parse_status_report(buf: &[u8]) -> Result<(usize, usize), Box<dyn Error>> {
@@ -71,7 +73,7 @@ impl State {
 
         write!(stdout, "\x1b[H\x1b[J")?;
         for i in self.row_start..self.row_start + usize::from(size.ws_row) {
-            if i > 0 {
+            if i > self.row_start {
                 write!(stdout, "\r\n")?;
             }
             if i > self.buf.lines() {
@@ -175,7 +177,7 @@ impl State {
         if self.cursor_col > self.line_cols {
             self.cursor_col = self.line_cols;
         }
-        if self.cursor_col == 0 && self.cursor_row == 0 {
+        if self.cursor_col == 0 && self.cursor_row == 0 && self.row_start == 0 {
             return Ok(());
         }
         if self.cursor_col == 0 {
@@ -189,6 +191,10 @@ impl State {
     }
 
     fn right(&mut self) -> Result<(), Box<dyn Error>> {
+        if self.cursor_col >= self.line_cols && self.row_start + self.cursor_row == self.buf.lines()
+        {
+            return Ok(());
+        }
         if self.cursor_col >= self.line_cols {
             self.down()?;
             self.cursor_col = 0;
@@ -208,6 +214,13 @@ impl State {
                 }
                 let _ = self.repaint_line_full();
             }
+            (Mode::Normal, Mode::Normal) => (),
+            (Mode::Normal, _) => {
+                let prev = self.prev.take();
+                let mut curr = Box::new(self.clone());
+                curr.prev = prev;
+                self.prev = Some(curr);
+            }
             _ => (),
         }
         self.mode = mode;
@@ -227,6 +240,7 @@ fn edit(file: &OsStr) -> Result<(), Box<dyn Error>> {
         cursor_col: 0,
         line_offset: 0,
         line_cols: 0,
+        prev: None,
     };
     let _ = state.repaint_all();
     let _ = write!(stdout(), "\x1b[H");
@@ -242,6 +256,16 @@ fn edit(file: &OsStr) -> Result<(), Box<dyn Error>> {
                 state.set_mode(Mode::Normal);
                 let _ = state.buf.save(Path::new(file));
             }
+            Key::CtrlZ => {
+                state.set_mode(Mode::Normal);
+                let Some(prev) = state.prev else {
+                    continue;
+                };
+                state = *prev;
+                let _ = state.repaint_all();
+                let _ = state.repaint_line_full();
+                let _ = stdout().flush();
+            }
             Key::Up => {
                 state.set_mode(Mode::Normal);
                 let _ = state.up();
@@ -264,20 +288,25 @@ fn edit(file: &OsStr) -> Result<(), Box<dyn Error>> {
             }
             Key::Backspace => {
                 state.set_mode(Mode::Backspace);
-                if state.row_start == 0 && state.cursor_row == 0 && state.cursor_col == 0 {
+                if state.row_start == 0 && state.cursor_row == 0 && state.line_offset == 0 {
                     continue;
                 }
-                let offset =
-                    state.buf.line_idx(state.row_start + state.cursor_row) + state.line_offset;
-                let need_refresh_all = state.cursor_col == 0;
-                let _ = state.left();
+                let line_start = state.buf.line_idx(state.row_start + state.cursor_row);
+                let gap_end = line_start + state.line_offset;
+                let gap_start = state.buf.floor_char_boundary(gap_end - 1);
+                if state.line_offset == 0 {
+                    let _ = state.left();
+                } else {
+                    let _ = write!(stdout(), "\x1b[G");
+                    let _ = state.buf.slice(line_start, gap_start).print(&mut stdout());
+                    let (_, col) = state.cursor_pos()?;
+                    state.cursor_col = col;
+                }
                 state.buf = state
                     .buf
-                    .slice(0, state.buf.floor_char_boundary(offset - 1))
-                    .concat(&state.buf.slice(offset, state.buf.len()));
-                if need_refresh_all {
-                    let _ = state.repaint_all();
-                }
+                    .slice(0, gap_start)
+                    .concat(&state.buf.slice(gap_end, state.buf.len()));
+                let _ = state.repaint_all();
                 let _ = state.repaint_line_full();
                 let _ = stdout().flush();
             }
diff --git a/src/term.rs b/src/term.rs
index 9c7e2cc..3a56091 100644
--- a/src/term.rs
+++ b/src/term.rs
@@ -68,6 +68,7 @@ pub fn raw() -> Result<RawHandle, Box<dyn Error>> {
 pub enum Key {
     CtrlQ,
     CtrlS,
+    CtrlZ,
     Up,
     Down,
     Left,
@@ -79,6 +80,7 @@ pub enum Key {
 pub fn read_key() -> Result<Key, Box<dyn Error>> {
     const CTRL_Q: u8 = 17;
     const CTRL_S: u8 = 19;
+    const CTRL_Z: u8 = 26;
     const BACKSPACE: u8 = 127;
     const ESC: u8 = 27;
 
@@ -91,6 +93,9 @@ pub fn read_key() -> Result<Key, Box<dyn Error>> {
         if buf[0] == CTRL_S {
             return Ok(Key::CtrlS);
         }
+        if buf[0] == CTRL_Z {
+            return Ok(Key::CtrlZ);
+        }
         if buf[0] == BACKSPACE {
             return Ok(Key::Backspace);
         }