2017-10-09 20:59:52 +08:00
|
|
|
use std::io::prelude::*;
|
|
|
|
use std::io;
|
|
|
|
use std::cmp;
|
|
|
|
|
|
|
|
pub trait ReadSeek: Read + Seek {}
|
|
|
|
impl<T> ReadSeek for T where T: Read + Seek {}
|
|
|
|
|
|
|
|
pub trait ReadWriteSeek: Read + Write + Seek {}
|
|
|
|
impl<T> ReadWriteSeek for T where T: Read + Write + Seek {}
|
|
|
|
|
|
|
|
const BUF_SIZE: usize = 512;
|
|
|
|
|
|
|
|
pub struct BufStream<T: Read+Write+Seek> {
|
|
|
|
inner: T,
|
|
|
|
buf: [u8; BUF_SIZE],
|
2017-10-10 20:48:57 +08:00
|
|
|
len: usize,
|
|
|
|
pos: usize,
|
|
|
|
write: bool,
|
2017-10-09 20:59:52 +08:00
|
|
|
}
|
|
|
|
|
2017-10-21 22:25:04 +08:00
|
|
|
/// The BufStream struct adds buffering to underlying file or device.
|
|
|
|
///
|
|
|
|
/// It's basically composition of BufReader and BufWritter.
|
2017-10-09 20:59:52 +08:00
|
|
|
impl<T: Read+Write+Seek> BufStream<T> {
|
2017-10-21 22:25:04 +08:00
|
|
|
/// Creates new BufStream object for given stream.
|
2017-10-09 20:59:52 +08:00
|
|
|
pub fn new(inner: T) -> Self {
|
|
|
|
BufStream::<T> {
|
|
|
|
inner,
|
|
|
|
buf: [0; BUF_SIZE],
|
2017-10-10 20:48:57 +08:00
|
|
|
pos: 0,
|
|
|
|
len: 0,
|
|
|
|
write: false,
|
2017-10-09 20:59:52 +08:00
|
|
|
}
|
|
|
|
}
|
2017-10-25 23:20:27 +08:00
|
|
|
|
2017-10-10 20:48:57 +08:00
|
|
|
fn flush_buf(&mut self) -> io::Result<()> {
|
|
|
|
if self.write {
|
|
|
|
self.inner.write_all(&self.buf[..self.pos])?;
|
|
|
|
self.pos = 0;
|
|
|
|
}
|
|
|
|
Ok(())
|
|
|
|
}
|
2017-10-25 23:20:27 +08:00
|
|
|
|
2017-10-10 20:48:57 +08:00
|
|
|
fn make_reader(&mut self) -> io::Result<()> {
|
|
|
|
if self.write {
|
|
|
|
self.flush_buf()?;
|
|
|
|
self.write = false;
|
|
|
|
self.len = 0;
|
|
|
|
self.pos = 0;
|
|
|
|
}
|
|
|
|
Ok(())
|
|
|
|
}
|
2017-10-25 23:20:27 +08:00
|
|
|
|
2017-10-10 20:48:57 +08:00
|
|
|
fn make_writter(&mut self) -> io::Result<()> {
|
|
|
|
if !self.write {
|
|
|
|
self.inner.seek(io::SeekFrom::Current(-(self.len as i64 - self.pos as i64)))?;
|
|
|
|
self.write = true;
|
|
|
|
self.len = 0;
|
|
|
|
self.pos = 0;
|
|
|
|
}
|
|
|
|
Ok(())
|
|
|
|
}
|
2017-10-09 20:59:52 +08:00
|
|
|
}
|
|
|
|
|
2017-10-10 20:48:57 +08:00
|
|
|
impl<T: Read+Write+Seek> BufRead for BufStream<T> {
|
|
|
|
fn fill_buf(&mut self) -> io::Result<&[u8]> {
|
|
|
|
self.make_reader()?;
|
|
|
|
if self.pos >= self.len {
|
|
|
|
self.len = self.inner.read(&mut self.buf)?;
|
|
|
|
self.pos = 0;
|
2017-10-09 20:59:52 +08:00
|
|
|
}
|
2017-10-10 20:48:57 +08:00
|
|
|
Ok(&self.buf[self.pos..self.len])
|
|
|
|
}
|
2017-10-25 23:20:27 +08:00
|
|
|
|
2017-10-10 20:48:57 +08:00
|
|
|
fn consume(&mut self, amt: usize) {
|
|
|
|
self.pos = cmp::min(self.pos + amt, self.len);
|
2017-10-09 20:59:52 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2017-10-10 20:48:57 +08:00
|
|
|
impl<T: Read+Write+Seek> Read for BufStream<T> {
|
|
|
|
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
|
|
|
|
// Make sure we are in read mode
|
|
|
|
self.make_reader()?;
|
|
|
|
// Check if this read is bigger than buffer size
|
|
|
|
if self.pos == self.len && buf.len() >= BUF_SIZE {
|
|
|
|
return self.inner.read(buf);
|
2017-10-09 20:59:52 +08:00
|
|
|
}
|
2017-10-10 20:48:57 +08:00
|
|
|
let nread = {
|
|
|
|
let mut rem = self.fill_buf()?;
|
|
|
|
rem.read(buf)?
|
|
|
|
};
|
|
|
|
self.consume(nread);
|
|
|
|
Ok(nread)
|
2017-10-09 20:59:52 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<T: Read+Write+Seek> Write for BufStream<T> {
|
|
|
|
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
2017-10-10 20:48:57 +08:00
|
|
|
// Make sure we are in write mode
|
|
|
|
self.make_writter()?;
|
|
|
|
if self.pos + buf.len() > BUF_SIZE {
|
|
|
|
self.flush_buf()?;
|
|
|
|
if buf.len() >= BUF_SIZE {
|
|
|
|
return self.inner.write(buf);
|
2017-10-09 20:59:52 +08:00
|
|
|
}
|
|
|
|
}
|
2017-10-10 20:48:57 +08:00
|
|
|
let written = (&mut self.buf[self.pos..]).write(buf)?;
|
|
|
|
self.pos += written;
|
|
|
|
Ok(written)
|
2017-10-09 20:59:52 +08:00
|
|
|
}
|
2017-10-25 23:20:27 +08:00
|
|
|
|
2017-10-09 20:59:52 +08:00
|
|
|
fn flush(&mut self) -> io::Result<()> {
|
2017-10-10 20:48:57 +08:00
|
|
|
self.flush_buf()?;
|
2017-10-09 20:59:52 +08:00
|
|
|
self.inner.flush()
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<T: Read+Write+Seek> Seek for BufStream<T> {
|
|
|
|
fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
|
2017-10-10 20:48:57 +08:00
|
|
|
self.flush_buf()?;
|
2017-10-09 20:59:52 +08:00
|
|
|
let new_pos = match pos {
|
2017-10-10 20:48:57 +08:00
|
|
|
io::SeekFrom::Current(x) => io::SeekFrom::Current(x - (self.len as i64 - self.pos as i64)),
|
2017-10-09 20:59:52 +08:00
|
|
|
_ => pos,
|
|
|
|
};
|
2017-10-10 20:48:57 +08:00
|
|
|
self.pos = 0;
|
|
|
|
self.len = 0;
|
2017-10-09 20:59:52 +08:00
|
|
|
self.inner.seek(new_pos)
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<T: Read+Write+Seek> Drop for BufStream<T> {
|
|
|
|
fn drop(&mut self) {
|
2017-11-07 08:13:02 +08:00
|
|
|
match self.flush() {
|
|
|
|
Err(err) => error!("flush failed {}", err),
|
|
|
|
_ => {},
|
|
|
|
}
|
2017-10-09 20:59:52 +08:00
|
|
|
}
|
|
|
|
}
|