Keep track of cursor position.

1 files changed, 107 insertions, 21 deletions

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;
+    }
+}