use std::cmp; use std::iter; use std::io::{self, Read, BufRead}; pub struct BufReader { inner: R, buf: Vec, pos: usize, cap: usize, } const INIT_BUFFER_SIZE: usize = 4096; const MAX_BUFFER_SIZE: usize = 8192 + 4096 * 100; impl BufReader { #[inline] pub fn new(rdr: R) -> BufReader { BufReader::with_capacity(rdr, INIT_BUFFER_SIZE) } #[inline] pub fn with_capacity(rdr: R, cap: usize) -> BufReader { let mut buf = Vec::with_capacity(cap); buf.extend(iter::repeat(0).take(cap)); BufReader { inner: rdr, buf: buf, pos: 0, cap: 0, } } #[inline] pub fn get_ref(&self) -> &R { &self.inner } #[inline] pub fn get_mut(&mut self) -> &mut R { &mut self.inner } #[inline] pub fn get_buf(&self) -> &[u8] { if self.pos < self.cap { debug!("slicing {:?}", (self.pos, self.cap, self.buf.len())); &self.buf[self.pos..self.cap] } else { &[] } } #[inline] pub fn into_inner(self) -> R { self.inner } #[inline] pub fn read_into_buf(&mut self) -> io::Result { self.maybe_reserve(); let v = &mut self.buf; if self.cap < v.capacity() { let nread = try!(self.inner.read(&mut v[self.cap..])); self.cap += nread; Ok(nread) } else { Ok(0) } } #[inline] fn maybe_reserve(&mut self) { let cap = self.buf.capacity(); if self.cap == cap { self.buf.reserve(cmp::min(cap * 4, MAX_BUFFER_SIZE) - cap); let new = self.buf.capacity() - self.buf.len(); self.buf.extend(iter::repeat(0).take(new)); } } } impl Read for BufReader { fn read(&mut self, buf: &mut [u8]) -> io::Result { if self.cap == self.pos && buf.len() >= self.buf.len() { return self.inner.read(buf); } let nread = { let mut rem = try!(self.fill_buf()); try!(rem.read(buf)) }; self.consume(nread); Ok(nread) } } impl BufRead for BufReader { fn fill_buf(&mut self) -> io::Result<&[u8]> { if self.pos == self.cap { self.cap = try!(self.inner.read(&mut self.buf)); self.pos = 0; } Ok(&self.buf[self.pos..self.cap]) } #[inline] fn consume(&mut self, amt: usize) { self.pos = cmp::min(self.pos + amt, self.cap); if self.pos == self.cap { self.pos = 0; self.cap = 0; } } } #[cfg(test)] mod tests { use std::io::{self, Read, BufRead}; use super::BufReader; struct SlowRead(u8); impl Read for SlowRead { fn read(&mut self, buf: &mut [u8]) -> io::Result { let state = self.0; self.0 += 1; (&match state % 3 { 0 => b"foo", 1 => b"bar", _ => b"baz", }[..]).read(buf) } } #[test] fn test_consume_and_get_buf() { let mut rdr = BufReader::new(SlowRead(0)); rdr.read_into_buf().unwrap(); rdr.consume(1); assert_eq!(rdr.get_buf(), b"oo"); rdr.read_into_buf().unwrap(); rdr.read_into_buf().unwrap(); assert_eq!(rdr.get_buf(), b"oobarbaz"); rdr.consume(5); assert_eq!(rdr.get_buf(), b"baz"); rdr.consume(3); assert_eq!(rdr.get_buf(), b""); assert_eq!(rdr.pos, 0); assert_eq!(rdr.cap, 0); } }