path: root/src/rope.rs

use std::fs::File;
use std::io::{ErrorKind, Read, Write};
use std::path::Path;
use std::error::Error;
use std::rc::Rc;

// Set to a small value to help identify bugs in the implementation.
const MAX_NODE_SIZE: usize = 4;

#[derive(Debug, Clone)]
struct Leaf {
    buf: Rc<[u8]>,
    start: u8,
    end: u8,
}

impl Leaf {
    fn bytes(&self) -> &[u8] {
        return &self.buf[usize::from(self.start)..usize::from(self.end)];
    }
}

#[derive(Debug, Clone)]
struct Branch {
    left: Rc<Rope>,
    right: Rc<Rope>,
    len: usize,
    // Number of newline characters under this branch.
    lines: usize,
}

#[derive(Debug, Clone)]
enum Node {
    Leaf(Leaf),
    Branch(Branch),
}

#[derive(Debug, Clone)]
pub struct Rope(Node);

impl Rope {
    fn leaf(buf: Vec<u8>) -> Rope {
        let len = buf.len();
        return Rope(Node::Leaf(Leaf{
            buf: Rc::from(buf),
            start: 0,
            end: u8::try_from(len).expect("buffer too long"),
        }));
    }

    pub fn lines(&self) -> usize {
        match self {
            Rope(Node::Leaf(l)) => {
                let mut count = 0;
                for &c in l.bytes().iter() {
                    if c == b'\n' {
                        count += 1;
                    }
                }
                return count;
            }
            Rope(Node::Branch(b)) => {
                return b.lines;
            }
        }
    }

    pub fn len(&self) -> usize {
        match self {
            Rope(Node::Leaf(l)) => {
                return usize::from(l.end - l.start);
            }
            Rope(Node::Branch(b)) => {
                return b.len;
            }
        }
    }

    pub fn print(&self, out: &mut dyn Write) -> Result<(), std::io::Error> {
        if self.len() == 0 {
            return Ok(());
        }
        // TODO: escape unprintable characters
        match self {
            Rope(Node::Leaf(l)) => {
                out.write_all(l.bytes())?;
                return Ok(());
            }
            Rope(Node::Branch(b)) => {
                b.left.print(out)?;
                b.right.print(out)?;
                return Ok(());
            }
        }
    }

    pub fn line_idx(&self, n: usize) -> usize {
        if n > self.lines() {
            panic!("Index {} out of range 0..{}", n, self.lines()+1);
        }
        if n == 0 {
            return 0;
        }
        match self {
            Rope(Node::Leaf(l)) => {
                let mut nl_count = 0;
                for (i, &c) in l.bytes().iter().enumerate() {
                    if c != b'\n' {
                        continue;
                    }
                    nl_count += 1;
                    if nl_count == n {
                        return i + 1;
                    }
                }
                panic!("unreachable");
            }
            Rope(Node::Branch(b)) => {
                let left_lines = b.left.lines();
                if n <= left_lines {
                    return b.left.line_idx(n);
                }
                return b.left.len() + b.right.line_idx(n - left_lines);
            }
        }
    }

    fn concat(&self, other: &Rope) -> Rope {
        if self.len() == 0 {
            return other.clone();
        }
        if other.len() == 0 {
            return self.clone();
        }
        let len = self.len() + other.len();
        let lines = self.lines() + other.lines();
        return Rope(Node::Branch(Branch {
            left: Rc::new(self.clone()),
            right: Rc::new(other.clone()),
            len: len,
            lines: lines,
        }));
    }

    pub fn insert(&self, pos: usize, c: u8) -> Rope {
        match self {
            Rope(Node::Leaf(l)) => {
                if self.len() < MAX_NODE_SIZE {
                    let mut new_buf = vec![0; self.len() + 1];
                    new_buf[..pos].copy_from_slice(&l.bytes()[..pos]);
                    new_buf[pos] = c;
                    new_buf[pos + 1..].copy_from_slice(&l.bytes()[pos..]);
                    return Rope::leaf(new_buf);
                }
                let mut buf_left = vec![0; pos + 1];
                buf_left[..pos].copy_from_slice(&l.bytes()[..pos]);
                buf_left[pos] = c;
                let mut buf_right = Vec::new();
                buf_right.extend_from_slice(&l.bytes()[pos..]);
                return Rope::leaf(buf_left).concat(&Rope::leaf(buf_right));
            }
            Rope(Node::Branch(b)) => {
                if pos < b.left.len() {
                    return b.left.insert(pos, c).concat(&b.right);
                }
                return b.left.concat(&b.right.insert(pos - b.left.len(), c));
            }
        }
    }

    pub fn open(path: &Path) -> Result<Rope, Box<dyn Error>> {
        let mut f = File::open(path).map_err(|err| format!("open file {}: {}", path.display(), err))?;
        let mut rope = Rope::leaf(Vec::new());
        loop {
            let mut buf = vec![0; MAX_NODE_SIZE];
            let n = match f.read(&mut buf) {
                Ok(n) => n,
                Err(err) => {
                    if err.kind() == ErrorKind::Interrupted {
                        continue;
                    }
                    return Err(Box::from(format!("read file {}: {}", path.display(), err)));
                }
            };
            if n == 0 {
                break;
            }
            buf.truncate(n);
            rope = rope.concat(&Rope::leaf(buf));
        }
        // TODO: rebalance
        return Ok(rope);
    }

    pub fn save(&self, path: &Path) -> Result<(), Box<dyn Error>> {
        let mut f = File::create(path).map_err(|err| format!("save: create file {}: {}", path.display(), err))?;
        self.print(&mut f).map_err(|err| format!("save: write file {}: {}", path.display(), err))?;
        f.sync_all().map_err(|err| format!("save: write file {}: {}", path.display(), err))?;
        return Ok(());
    }

    pub fn slice(&self, start: usize, end: usize) -> Rope {
        if start > self.len() {
            panic!("Index {} out of range 0..{}", start, self.len()+1);
        }
        if end > self.len() {
            panic!("Index {} out of range 0..{}", end, self.len());
        }
        if start > end {
            panic!("Slice start index {} is greater than end index {}", start, end);
        }
        match self {
            Rope(Node::Leaf(l)) => {
                return Rope(Node::Leaf(Leaf{
                    buf: l.buf.clone(),
                    start: l.start + u8::try_from(start).expect("buffer too long"),
                    end: l.start + u8::try_from(end).expect("buffer too long"),
                }));
            }
            Rope(Node::Branch(b)) => {
                let mut left = Rope::leaf(Vec::new());
                if start < b.left.len() {
                    left = b.left.slice(start, std::cmp::min(end, b.left.len()));
                }
                let mut right = Rope::leaf(Vec::new());
                if end > b.left.len() {
                    let mut right_start = 0;
                    if start > b.left.len() {
                        right_start = start - b.left.len();
                    }
                    right = b.right.slice(right_start, end - b.left.len());
                }
                return left.concat(&right);
            }
        }
    }

    pub fn line(&self, n: usize) -> Rope {
        if n > self.lines() {
            panic!("Index {} out of range 0..{}", n, self.lines()+1);
        }
        let start = self.line_idx(n);
        let mut end = self.len();
        if n < self.lines() {
            end = self.line_idx(n+1)-1;
        }
        return self.slice(start, end);
    }

    fn is_char_boundary(&self, index: usize) -> bool {
        match self {
            Rope(Node::Leaf(l)) => {
                return l.bytes()[index] & 0xc0 != 0x80;
            }
            Rope(Node::Branch(b)) => {
                if index < b.left.len() {
                    return b.left.is_char_boundary(index);
                }
                return b.right.is_char_boundary(index - b.left.len());
            }
        }
    }

    pub fn floor_char_boundary(&self, index: usize) -> usize {
        for i in (0..index+1).rev() {
            if self.is_char_boundary(i) {
                return i;
            }
        }
        panic!("I'm not valid UTF-8: {:?}", self);
    }

    pub fn ceil_char_boundary(&self, index: usize) -> usize {
        for i in index..self.len() {
            if self.is_char_boundary(i) {
                return i;
            }
        }
        return self.len();
    }

    pub fn char(&self, index: usize) -> usize {
        if index == 0 {
            return 0;
        }
        let mut count = 1;
        for i in 1..self.len() {
            if !self.is_char_boundary(i) {
                continue;
            }
            if count == index {
                return i;
            }
            count += 1;
        }
        return self.len();
    }
}