Keep track of cursor position.

2 files changed, 172 insertions, 31 deletions

diff --git a/src/main.rs b/src/main.rs
index 592f89c..de25e9c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,10 +3,15 @@ mod term;
 
 use rope::Rope;
 use std::ffi::{OsStr, OsString};
-use std::io::Write;
+use std::io::{Stdout, Write};
 use std::path::Path;
 use term::Key;
 
+fn move_cursor(stdout: &mut Stdout, row: usize, col: usize) {
+    let _ = write!(stdout, "\x1b[{};{}H", row + 1, col + 1);
+    let _ = stdout.flush();
+}
+
 fn edit(file: &OsStr) -> Result<(), String> {
     let _raw_handle = match term::raw() {
         Some(h) => h,
@@ -28,14 +33,25 @@ fn edit(file: &OsStr) -> Result<(), String> {
         }
     };
     let mut stdout = std::io::stdout();
+    let _ = stdout.write(b"\x1b[J");
     for i in 0..usize::from(size.ws_row) {
         if i > 0 {
             let _ = stdout.write(b"\r\n");
         }
-        let _ = r.print_line(&mut stdout, i);
+        let line = match r.line(i) {
+            Some(line) => line,
+            None => {
+                break;
+            }
+        };
+        let _ = line.print(&mut stdout);
     }
+    let _ = stdout.write(b"\x1b[H");
     let _ = stdout.flush();
 
+    let mut cursor_row = 0;
+    let mut cursor_col = 0;
+
     let mut stdin = std::io::stdin();
     loop {
         let c = match term::read_key(&mut stdin) {
@@ -49,20 +65,59 @@ fn edit(file: &OsStr) -> Result<(), String> {
                 break;
             }
             Key::Up => {
-                let _ = stdout.write(b"\x1b[A");
-                let _ = stdout.flush();
+                if cursor_row == 0 {
+                    continue;
+                }
+                cursor_row -= 1;
+
+                let len = r
+                    .line(cursor_row)
+                    .expect("cursor_row should always be valid")
+                    .len();
+                if cursor_col > len {
+                    move_cursor(&mut stdout, cursor_row, len);
+                } else {
+                    move_cursor(&mut stdout, cursor_row, cursor_col);
+                }
             }
             Key::Down => {
-                let _ = stdout.write(b"\x1b[B");
-                let _ = stdout.flush();
+                if cursor_row == usize::from(size.ws_row - 1) || cursor_row == r.lines() {
+                    continue;
+                }
+                cursor_row += 1;
+
+                let len = r
+                    .line(cursor_row)
+                    .expect("cursor_row should always be valid")
+                    .len();
+                if cursor_col > len {
+                    move_cursor(&mut stdout, cursor_row, len);
+                } else {
+                    move_cursor(&mut stdout, cursor_row, cursor_col);
+                }
             }
             Key::Left => {
-                let _ = stdout.write(b"\x1b[D");
-                let _ = stdout.flush();
+                let len = r
+                    .line(cursor_row)
+                    .expect("cursor_row should always be valid")
+                    .len();
+                cursor_col = std::cmp::min(cursor_col, len);
+                if cursor_col == 0 {
+                    continue;
+                }
+                cursor_col -= 1;
+                move_cursor(&mut stdout, cursor_row, cursor_col);
             }
             Key::Right => {
-                let _ = stdout.write(b"\x1b[C");
-                let _ = stdout.flush();
+                let len = r
+                    .line(cursor_row)
+                    .expect("cursor_row should always be valid")
+                    .len();
+                if cursor_col == usize::from(size.ws_col - 1) || cursor_col >= len {
+                    continue;
+                }
+                cursor_col += 1;
+                move_cursor(&mut stdout, cursor_row, cursor_col);
             }
             Key::Char(c) => {
                 print!("{} ", c);
diff --git a/src/rope.rs b/src/rope.rs
index be56568..3162e44 100644
--- a/src/rope.rs
+++ b/src/rope.rs
@@ -10,6 +10,7 @@ const MAX_NODE_SIZE: usize = 4;
 struct Branch {
     left: Rc<Node>,
     right: Rc<Node>,
+    len: usize,
     // Number of newline characters under this branch.
     lines: usize,
 }
@@ -56,44 +57,78 @@ impl Node {
         }
     }
 
-    fn print_line(&self, out: &mut dyn Write, n: usize) -> Result<(), std::io::Error> {
+    fn len(&self) -> usize {
+        match self {
+            Node::Leaf(v) => {
+                return v.len();
+            }
+            Node::Branch(b) => {
+                return b.len;
+            }
+        }
+    }
+
+    fn print(&self, out: &mut dyn Write, start: usize, end: usize) -> Result<(), std::io::Error> {
         // TODO: escape unprintable characters
         match self {
             Node::Leaf(v) => {
+                if let Err(err) = out.write(&v[start..end]) {
+                    return Err(err);
+                }
+                return Ok(());
+            }
+            Node::Branch(b) => {
+                let left_len = b.left.len();
+                if start < left_len {
+                    if let Err(err) = b.left.print(out, start, std::cmp::min(left_len, end)) {
+                        return Err(err);
+                    }
+                }
+                if end > left_len {
+                    let mut right_start = 0;
+                    if start > left_len {
+                        right_start = start - left_len;
+                    }
+                    if let Err(err) = b.right.print(out, right_start, end - left_len) {
+                        return Err(err);
+                    }
+                }
+                return Ok(());
+            }
+        }
+    }
+
+    fn line_idx(&self, n: usize) -> Option<usize> {
+        if n == 0 {
+            return Some(0);
+        }
+        match self {
+            Node::Leaf(v) => {
                 let mut nl_count = 0;
-                let mut start_pos = 0;
-                let mut end_pos = v.len();
                 for (i, &c) in v.iter().enumerate() {
                     if c != b'\n' {
                         continue;
                     }
                     nl_count += 1;
                     if nl_count == n {
-                        start_pos = i + 1;
+                        return Some(i + 1);
                     }
-                    if nl_count == n + 1 {
-                        end_pos = i;
-                        break;
-                    }
-                }
-                if let Err(err) = out.write(&v[start_pos..end_pos]) {
-                    return Err(err);
                 }
-                return Ok(());
+                return None;
             }
             Node::Branch(b) => {
                 let left_lines = b.left.lines();
                 if n <= left_lines {
-                    if let Err(err) = b.left.print_line(out, n) {
-                        return Err(err);
-                    }
+                    return b.left.line_idx(n);
                 }
-                if n >= left_lines {
-                    if let Err(err) = b.right.print_line(out, n - left_lines) {
-                        return Err(err);
+                match b.right.line_idx(n - left_lines) {
+                    Some(i) => {
+                        return Some(b.left.len() + i);
+                    }
+                    None => {
+                        return None;
                     }
                 }
-                return Ok(());
             }
         }
     }
@@ -107,6 +142,7 @@ impl Display for Node {
     }
 }
 
+#[derive(Clone)]
 pub struct Rope(Rc<Node>);
 
 impl Rope {
@@ -118,10 +154,12 @@ impl Rope {
         if other.empty() {
             return Rope(me);
         }
+        let len = me.len() + other.len();
         let lines = me.lines() + other.lines();
         return Rope(Rc::new(Node::Branch(Branch {
             left: me,
             right: other,
+            len: len,
             lines: lines,
         })));
     }
@@ -154,9 +192,40 @@ impl Rope {
         return Ok(rope);
     }
 
-    pub fn print_line(&self, out: &mut dyn Write, n: usize) -> Result<(), std::io::Error> {
+    pub fn len(&self) -> usize {
         let Rope(me) = self;
-        return me.print_line(out, n);
+        return me.len();
+    }
+
+    pub fn lines(&self) -> usize {
+        let Rope(me) = self;
+        return me.lines();
+    }
+
+    pub fn line(&self, n: usize) -> Option<Slice> {
+        let Rope(me) = self;
+        let start = match me.line_idx(n) {
+            Some(start) => start,
+            None => {
+                return None;
+            }
+        };
+        match me.line_idx(n + 1) {
+            Some(end) => {
+                return Some(Slice {
+                    start: start,
+                    end: end - 1,
+                    buf: self.clone(),
+                });
+            }
+            None => {
+                return Some(Slice {
+                    start: start,
+                    end: self.len(),
+                    buf: self.clone(),
+                });
+            }
+        }
     }
 }
 
@@ -166,3 +235,20 @@ impl Display for Rope {
         return write!(f, "{}", n);
     }
 }
+
+pub struct Slice {
+    start: usize,
+    end: usize,
+    buf: Rope,
+}
+
+impl Slice {
+    pub fn print(&self, out: &mut dyn Write) -> Result<(), std::io::Error> {
+        let Rope(buf) = &self.buf;
+        return buf.print(out, self.start, self.end);
+    }
+
+    pub fn len(&self) -> usize {
+        return self.end - self.start;
+    }
+}