use std::io::{BufRead, Error as IoError, ErrorKind as IoErrorKind, Read, Result as IoResult};
use std::ops::Bound;
use std::ops::RangeBounds;
pub struct GreedyBufReader<R> {
inner: R,
buf: Vec<u8>,
consumed: usize,
}
impl<R> GreedyBufReader<R>
where
R: Read,
{
pub fn new(src: R) -> Self {
GreedyBufReader {
inner: src,
buf: Vec::new(),
consumed: 0,
}
}
pub fn with_capacity(src: R, capacity: usize) -> Self {
GreedyBufReader {
inner: src,
buf: Vec::with_capacity(capacity),
consumed: 0,
}
}
pub fn into_inner(self) -> R {
self.inner
}
pub fn into_buffer(self) -> Vec<u8> {
self.buf
}
pub fn get(&mut self, index: usize) -> IoResult<u8> {
if let Some(v) = self.buf.get(index) {
Ok(*v)
} else {
self.prefetch_up_to(index + 1)?;
self.buf
.get(index)
.cloned()
.ok_or_else(|| IoError::new(IoErrorKind::Other, "Index out of bounds"))
}
}
pub fn slice<T>(&mut self, range: T) -> IoResult<&[u8]>
where
T: Clone,
T: RangeBounds<usize>,
{
let end = range.end_bound();
let e = match end {
Bound::Unbounded => {
unimplemented!("Unbounded end is currently not supported");
}
Bound::Excluded(&e) => e,
Bound::Included(&e) => e + 1,
};
let b = match range.start_bound() {
Bound::Unbounded => 0,
Bound::Excluded(&b) | Bound::Included(&b) => b,
};
self.prefetch_up_to(e)?;
if b > e || e > self.buf.len() {
Err(IoError::new(IoErrorKind::Other, "Index out of bounds"))
} else {
Ok(&self.buf[b..e])
}
}
pub fn clear(&mut self) {
if self.consumed < self.buf.len() {
self.buf = self.buf[self.consumed..].to_vec();
} else {
self.buf = Vec::new();
}
self.consumed = 0;
}
pub fn shrink_to_fit(&mut self) {
self.buf.shrink_to_fit()
}
fn reserve_up_to(&mut self, index: usize) {
let mut new_size = 16;
while new_size < index || new_size < self.buf.capacity() {
new_size *= 2;
}
let additional = new_size - self.buf.capacity();
if additional > 0 {
self.buf.reserve(additional);
}
}
fn data_to_read(&self) -> &[u8] {
&self.buf[self.consumed..]
}
fn prefetch_up_to(&mut self, i: usize) -> IoResult<()> {
self.reserve_up_to(i);
let mut l = 0;
while self.buf.len() <= i {
let b = self.fill_buf()?;
if b.len() == l {
break;
} else {
l = b.len();
}
}
Ok(())
}
}
impl<R> Read for GreedyBufReader<R>
where
R: Read,
{
fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
let mut to_read = self.data_to_read();
if to_read.is_empty() {
self.fill_buf()?;
to_read = self.data_to_read();
}
let len = usize::min(to_read.len(), buf.len());
&mut buf[..len].copy_from_slice(&self.buf[self.consumed..self.consumed + len]);
self.consume(len);
Ok(len)
}
}
impl<R> BufRead for GreedyBufReader<R>
where
R: Read,
{
fn fill_buf(&mut self) -> IoResult<&[u8]> {
if self.buf.capacity() == self.consumed {
self.reserve_up_to(self.buf.capacity() + 16);
}
let b = self.buf.len();
unsafe {
self.buf.set_len(self.buf.capacity());
}
let buf = &mut self.buf[b..];
match self.inner.read(buf) {
Ok(o) => {
self.buf.truncate(b + o);
Ok(&self.buf[self.consumed..])
}
Err(e) => Err(e),
}
}
fn consume(&mut self, amt: usize) {
self.consumed += amt;
}
}
#[cfg(test)]
mod tests {
use super::GreedyBufReader;
use std::io::Read;
#[test]
fn smoke_test() {
let data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 50];
let mut read = GreedyBufReader::new(&data[..]);
let mut o = Vec::new();
read.read_to_end(&mut o).unwrap();
assert_eq!(o, &data);
}
#[test]
fn test_get() {
let data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 50];
let mut read = GreedyBufReader::new(&data[..]);
assert_eq!(read.get(1).unwrap(), 2);
assert_eq!(read.get(2).unwrap(), 3);
assert_eq!(read.get(16).unwrap(), 50);
assert_eq!(read.get(10).unwrap(), 11);
assert!(read.get(17).is_err());
}
#[test]
fn test_slice() {
let data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 50];
let mut read = GreedyBufReader::new(&data[..]);
assert_eq!(read.slice(0..0).unwrap(), &[]);
assert_eq!(read.slice(1..2).unwrap(), &[2]);
assert_eq!(read.slice(..=5).unwrap(), &[1, 2, 3, 4, 5, 6]);
assert_eq!(read.slice(14..=16).unwrap(), &[15, 16, 50]);
assert_eq!(read.slice(10..12).unwrap(), &[11, 12]);
assert!(read.slice(7..18).is_err());
assert!(read.slice(6..5).is_err());
}
#[test]
fn arbitrary_get_infinite() {
const B: u8 = 0x33;
let mut read = GreedyBufReader::new(std::io::repeat(B));
assert_eq!(read.get(4).unwrap(), B);
assert_eq!(read.get(13).unwrap(), B);
assert_eq!(read.get(24389).unwrap(), B);
assert_eq!(read.get(156).unwrap(), B);
assert_eq!(read.get(9006).unwrap(), B);
assert_eq!(read.get(2019).unwrap(), B);
assert_eq!(read.get(100000).unwrap(), B);
}
#[test]
fn test_clear() {
let data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 50];
let mut read = GreedyBufReader::new(&data[..]);
assert_eq!(read.get(0).unwrap(), 1);
assert_eq!(read.get(8).unwrap(), 9);
assert_eq!(read.get(16).unwrap(), 50);
let mut chunk = [0; 8];
read.read_exact(&mut chunk).unwrap();
assert_eq!(chunk, [1, 2, 3, 4, 5, 6, 7, 8]);
assert_eq!(read.get(0).unwrap(), 1);
assert_eq!(read.get(8).unwrap(), 9);
assert_eq!(read.get(16).unwrap(), 50);
read.clear();
assert_eq!(read.get(0).unwrap(), 9);
assert_eq!(read.get(8).unwrap(), 50);
assert!(read.get(16).is_err());
read.read_exact(&mut chunk).unwrap();
assert_eq!(chunk, [9, 10, 11, 12, 13, 14, 15, 16]);
assert_eq!(read.get(0).unwrap(), 9);
assert_eq!(read.get(8).unwrap(), 50);
assert!(read.get(16).is_err());
}
}