bitbottle 0.10.0

use std::cell::RefCell;
use std::io::{Read, Write};
use std::rc::Rc;


/// Wrap a `Write` and count the bytes written.
pub struct CountingWriter<W: Write> {
    writer: W,
    pub count: Rc<RefCell<u64>>,
}

impl<W: Write> CountingWriter<W> {
    /// Return a new `CountingWriter` and a `RefCell` to its byte count.
    pub fn new(writer: W) -> (CountingWriter<W>, Rc<RefCell<u64>>) {
        let count = Rc::new(RefCell::new(0));
        (CountingWriter { writer, count: count.clone() }, count)
    }

    /// Dissolve back into the original `Write` and return the final byte
    /// count.
    pub fn into_inner(self) -> (W, u64) {
        (self.writer, *self.count.borrow())
    }
}

impl<W: Write> Write for CountingWriter<W> {
    fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
        *self.count.borrow_mut() += data.len() as u64;
        self.writer.write(data)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.writer.flush()
    }
}


/// Wrap a `Read` and count the bytes read.
pub struct CountingReader<R: Read> {
    reader: R,
    count: Rc<RefCell<u64>>,
}

impl<R: Read> CountingReader<R> {
    /// Return a new `CountingReader` and a `RefCell` to its byte count.
    pub fn new(reader: R) -> (CountingReader<R>, Rc<RefCell<u64>>) {
        let count = Rc::new(RefCell::new(0));
        (CountingReader { reader, count: count.clone() }, count)
    }

    /// Dissolve back into the original `Read` and return the final byte
    /// count.
    pub fn into_inner(self) -> (R, u64) {
        (self.reader, *self.count.borrow())
    }
}

impl<R: Read> Read for CountingReader<R> {
    fn read(&mut self, data: &mut [u8]) -> std::io::Result<usize> {
        let n = self.reader.read(data)?;
        *self.count.borrow_mut() += n as u64;
        Ok(n)
    }
}