// Copyright 2016-2018 Austin Bonander <austin.bonander@gmail.com> | |
// | |
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or | |
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license | |
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your | |
// option. This file may not be copied, modified, or distributed | |
// except according to those terms. | |
//! Types which can be used to tune the behavior of `BufReader` and `BufWriter`. | |
//! | |
//! Some simple policies are provided for your convenience. You may prefer to create your own | |
//! types and implement the traits for them instead. | |
use super::Buffer; | |
/// Flag for `ReaderPolicy` methods to signal whether or not `BufReader` should read into | |
/// the buffer. | |
/// | |
/// See `do_read!()` for a shorthand. | |
#[derive(Copy, Clone, Debug)] | |
pub struct DoRead(pub bool); | |
/// Shorthand for `return DoRead(bool)` or `return DoRead(true)` (empty invocation) | |
#[macro_export] | |
macro_rules! do_read ( | |
($val:expr) => ( return $crate::policy::DoRead($val); ); | |
() => ( do_read!(true); ) | |
); | |
/// Default policy for both `BufReader` and `BufWriter` that reproduces the behaviors of their | |
/// `std::io` counterparts: | |
/// | |
/// * `BufReader`: only reads when the buffer is empty, does not resize or move data. | |
/// * `BufWriter`: only flushes the buffer when there is not enough room for an incoming write. | |
#[derive(Debug, Default)] | |
pub struct StdPolicy; | |
/// Trait that governs `BufReader`'s behavior. | |
pub trait ReaderPolicy { | |
/// Consulted before attempting to read into the buffer. | |
/// | |
/// Return `DoRead(true)` to issue a read into the buffer before reading data out of it, | |
/// or `DoRead(false)` to read from the buffer as it is, even if it's empty. | |
/// `do_read!()` is provided as a shorthand. | |
/// | |
/// If there is no room in the buffer after this method is called, | |
/// the buffer will not be read into (so if the buffer is full but you want more data | |
/// you should call `.make_room()` or reserve more space). If there *is* room, `BufReader` will | |
/// attempt to read into the buffer. If successful (`Ok(x)` where `x > 0` is returned), this | |
/// method will be consulted again for another read attempt. | |
/// | |
/// By default, this implements `std::io::BufReader`'s behavior: only read into the buffer if | |
/// it is empty. | |
/// | |
/// ### Note | |
/// If the read will ignore the buffer entirely (if the buffer is empty and the amount to be | |
/// read matches or exceeds its capacity) or if `BufReader::read_into_buf()` was called to force | |
/// a read into the buffer manually, this method will not be called. | |
fn before_read(&mut self, buffer: &mut Buffer) -> DoRead { DoRead(buffer.len() == 0) } | |
/// Called after bytes are consumed from the buffer. | |
/// | |
/// Supplies the true amount consumed if the amount passed to `BufReader::consume` | |
/// was in excess. | |
/// | |
/// This is a no-op by default. | |
fn after_consume(&mut self, _buffer: &mut Buffer, _amt: usize) {} | |
} | |
/// Behavior of `std::io::BufReader`: the buffer will only be read into if it is empty. | |
impl ReaderPolicy for StdPolicy {} | |
/// A policy for [`BufReader`](::BufReader) which ensures there is at least the given number of | |
/// bytes in the buffer, failing this only if the reader is at EOF. | |
/// | |
/// If the minimum buffer length is greater than the buffer capacity, it will be resized. | |
/// | |
/// ### Example | |
/// ```rust | |
/// use buf_redux::BufReader; | |
/// use buf_redux::policy::MinBuffered; | |
/// use std::io::{BufRead, Cursor}; | |
/// | |
/// let data = (1 .. 16).collect::<Vec<u8>>(); | |
/// | |
/// // normally you should use `BufReader::new()` or give a capacity of several KiB or more | |
/// let mut reader = BufReader::with_capacity(8, Cursor::new(data)) | |
/// // always at least 4 bytes in the buffer (or until the source is empty) | |
/// .set_policy(MinBuffered(4)); // always at least 4 bytes in the buffer | |
/// | |
/// // first buffer fill, same as `std::io::BufReader` | |
/// assert_eq!(reader.fill_buf().unwrap(), &[1, 2, 3, 4, 5, 6, 7, 8]); | |
/// reader.consume(3); | |
/// | |
/// // enough data in the buffer, another read isn't done yet | |
/// assert_eq!(reader.fill_buf().unwrap(), &[4, 5, 6, 7, 8]); | |
/// reader.consume(4); | |
/// | |
/// // `std::io::BufReader` would return `&[8]` | |
/// assert_eq!(reader.fill_buf().unwrap(), &[8, 9, 10, 11, 12, 13, 14, 15]); | |
/// reader.consume(5); | |
/// | |
/// // no data left in the reader | |
/// assert_eq!(reader.fill_buf().unwrap(), &[13, 14, 15]); | |
/// ``` | |
#[derive(Debug)] | |
pub struct MinBuffered(pub usize); | |
impl MinBuffered { | |
/// Set the number of bytes to ensure are in the buffer. | |
pub fn set_min(&mut self, min: usize) { | |
self.0 = min; | |
} | |
} | |
impl ReaderPolicy for MinBuffered { | |
fn before_read(&mut self, buffer: &mut Buffer) -> DoRead { | |
// do nothing if we have enough data | |
if buffer.len() >= self.0 { do_read!(false) } | |
let cap = buffer.capacity(); | |
// if there's enough room but some of it's stuck after the head | |
if buffer.usable_space() < self.0 && buffer.free_space() >= self.0 { | |
buffer.make_room(); | |
} else if cap < self.0 { | |
buffer.reserve(self.0 - cap); | |
} | |
DoRead(true) | |
} | |
} | |
/// Flag for `WriterPolicy` methods to tell `BufWriter` how many bytes to flush to the | |
/// underlying reader. | |
/// | |
/// See `flush_amt!()` for a shorthand. | |
#[derive(Copy, Clone, Debug)] | |
pub struct FlushAmt(pub usize); | |
/// Shorthand for `return FlushAmt(n)` or `return FlushAmt(0)` (empty invocation) | |
#[macro_export] | |
macro_rules! flush_amt ( | |
($n:expr) => ( return $crate::policy::FlushAmt($n); ); | |
() => ( flush_amt!(0); ) | |
); | |
/// A trait which tells `BufWriter` when to flush. | |
pub trait WriterPolicy { | |
/// Return `FlushAmt(n > 0)` if the buffer should be flushed before reading into it. | |
/// If the returned amount is 0 or greater than the amount of buffered data, no flush is | |
/// performed. | |
/// | |
/// The buffer is provided, as well as `incoming` which is | |
/// the size of the buffer that will be written to the `BufWriter`. | |
/// | |
/// By default, flushes the buffer if the usable space is smaller than the incoming write. | |
fn before_write(&mut self, buf: &mut Buffer, incoming: usize) -> FlushAmt { | |
FlushAmt(if incoming > buf.usable_space() { buf.len() } else { 0 }) | |
} | |
/// Return `true` if the buffer should be flushed after reading into it. | |
/// | |
/// `buf` references the updated buffer after the read. | |
/// | |
/// Default impl is a no-op. | |
fn after_write(&mut self, _buf: &Buffer) -> FlushAmt { | |
FlushAmt(0) | |
} | |
} | |
/// Default behavior of `std::io::BufWriter`: flush before a read into the buffer | |
/// only if the incoming data is larger than the buffer's writable space. | |
impl WriterPolicy for StdPolicy {} | |
/// Flush the buffer if it contains at least the given number of bytes. | |
#[derive(Debug, Default)] | |
pub struct FlushAtLeast(pub usize); | |
impl WriterPolicy for FlushAtLeast { | |
fn before_write(&mut self, buf: &mut Buffer, incoming: usize) -> FlushAmt { | |
ensure_capacity(buf, self.0); | |
FlushAmt(if incoming > buf.usable_space() { buf.len() } else { 0 }) | |
} | |
fn after_write(&mut self, buf: &Buffer) -> FlushAmt { | |
FlushAmt(::std::cmp::max(buf.len(), self.0)) | |
} | |
} | |
/// Only ever flush exactly the given number of bytes, until the writer is empty. | |
#[derive(Debug, Default)] | |
pub struct FlushExact(pub usize); | |
impl WriterPolicy for FlushExact { | |
/// Flushes the buffer if there is not enough room to fit `incoming` bytes, | |
/// but only when the buffer contains at least `self.0` bytes. | |
/// | |
/// Otherwise, calls [`Buffer::make_room()`](::Buffer::make_room) | |
fn before_write(&mut self, buf: &mut Buffer, incoming: usize) -> FlushAmt { | |
ensure_capacity(buf, self.0); | |
// don't have enough room to fit the additional bytes but we can't flush, | |
// then make room for (at least some of) the incoming bytes. | |
if incoming > buf.usable_space() && buf.len() < self.0 { | |
buf.make_room(); | |
} | |
FlushAmt(self.0) | |
} | |
/// Flushes the given amount if possible, nothing otherwise. | |
fn after_write(&mut self, _buf: &Buffer) -> FlushAmt { | |
FlushAmt(self.0) | |
} | |
} | |
/// Flush the buffer if it contains the given byte. | |
/// | |
/// Only scans the buffer after reading. Searches from the end first. | |
#[derive(Debug, Default)] | |
pub struct FlushOn(pub u8); | |
impl WriterPolicy for FlushOn { | |
fn after_write(&mut self, buf: &Buffer) -> FlushAmt { | |
// include the delimiter in the flush | |
FlushAmt(::memchr::memrchr(self.0, buf.buf()).map_or(0, |n| n + 1)) | |
} | |
} | |
/// Flush the buffer if it contains a newline (`\n`). | |
/// | |
/// Equivalent to `FlushOn(b'\n')`. | |
#[derive(Debug, Default)] | |
pub struct FlushOnNewline; | |
impl WriterPolicy for FlushOnNewline { | |
fn after_write(&mut self, buf: &Buffer) -> FlushAmt { | |
FlushAmt(::memchr::memrchr(b'\n', buf.buf()).map_or(0, |n| n + 1)) | |
} | |
} | |
fn ensure_capacity(buf: &mut Buffer, min_cap: usize) { | |
let cap = buf.capacity(); | |
if cap < min_cap { | |
buf.reserve(min_cap - cap); | |
} | |
} | |
#[cfg(test)] | |
mod test { | |
use {BufReader, BufWriter}; | |
use policy::*; | |
use std::io::{BufRead, Cursor, Write}; | |
#[test] | |
fn test_min_buffered() { | |
let min_buffered = 4; | |
let data = (0 .. 20).collect::<Vec<u8>>(); | |
// create a reader with 0 capacity | |
let mut reader = BufReader::with_capacity(0, Cursor::new(data)) | |
.set_policy(MinBuffered(min_buffered)); | |
// policy reserves the required space in the buffer | |
assert_eq!(reader.fill_buf().unwrap(), &[0, 1, 2, 3][..]); | |
assert_eq!(reader.capacity(), min_buffered); | |
// double the size now that the buffer's full | |
reader.reserve(min_buffered); | |
assert_eq!(reader.capacity(), min_buffered * 2); | |
// we haven't consumed anything, the reader should have the same data | |
assert_eq!(reader.fill_buf().unwrap(), &[0, 1, 2, 3]); | |
reader.consume(2); | |
// policy read more data, `std::io::BufReader` doesn't do that | |
assert_eq!(reader.fill_buf().unwrap(), &[2, 3, 4, 5, 6, 7]); | |
reader.consume(4); | |
// policy made room and read more | |
assert_eq!(reader.fill_buf().unwrap(), &[6, 7, 8, 9, 10, 11, 12, 13]); | |
reader.consume(4); | |
assert_eq!(reader.fill_buf().unwrap(), &[10, 11, 12, 13]); | |
reader.consume(2); | |
assert_eq!(reader.fill_buf().unwrap(), &[12, 13, 14, 15, 16, 17, 18, 19]); | |
reader.consume(8); | |
assert_eq!(reader.fill_buf().unwrap(), &[]) | |
} | |
#[test] | |
fn test_flush_at_least() { | |
let flush_min = 4; | |
let mut writer = BufWriter::with_capacity(0, vec![]).set_policy(FlushAtLeast(flush_min)); | |
assert_eq!(writer.capacity(), 0); | |
assert_eq!(writer.write(&[1]).unwrap(), 1); | |
// policy reserved space for writing | |
assert_eq!(writer.capacity(), flush_min); | |
// one byte in buffer, we want to double our capacity | |
writer.reserve(flush_min * 2 - 1); | |
assert_eq!(writer.capacity(), flush_min * 2); | |
assert_eq!(writer.write(&[2, 3]).unwrap(), 2); | |
// no flush yet, only 3 bytes in buffer | |
assert_eq!(*writer.get_ref(), &[]); | |
assert_eq!(writer.write(&[4, 5, 6]).unwrap(), 3); | |
// flushed all | |
assert_eq!(*writer.get_ref(), &[1, 2, 3, 4, 5, 6]); | |
assert_eq!(writer.write(&[7, 8, 9]).unwrap(), 3); | |
// `.into_inner()` should flush always | |
assert_eq!(writer.into_inner().unwrap(), &[1, 2, 3, 4, 5, 6, 7, 8, 9]); | |
} | |
#[test] | |
fn test_flush_exact() { | |
let flush_exact = 4; | |
let mut writer = BufWriter::with_capacity(0, vec![]).set_policy(FlushExact(flush_exact)); | |
assert_eq!(writer.capacity(), 0); | |
assert_eq!(writer.write(&[1]).unwrap(), 1); | |
// policy reserved space for writing | |
assert_eq!(writer.capacity(), flush_exact); | |
// one byte in buffer, we want to double our capacity | |
writer.reserve(flush_exact * 2 - 1); | |
assert_eq!(writer.capacity(), flush_exact * 2); | |
assert_eq!(writer.write(&[2, 3]).unwrap(), 2); | |
// no flush yet, only 3 bytes in buffer | |
assert_eq!(*writer.get_ref(), &[]); | |
assert_eq!(writer.write(&[4, 5, 6]).unwrap(), 3); | |
// flushed exactly 4 bytes | |
assert_eq!(*writer.get_ref(), &[1, 2, 3, 4]); | |
assert_eq!(writer.write(&[7, 8, 9, 10]).unwrap(), 4); | |
// flushed another 4 bytes | |
assert_eq!(*writer.get_ref(), &[1, 2, 3, 4, 5, 6, 7, 8]); | |
// `.into_inner()` should flush always | |
assert_eq!(writer.into_inner().unwrap(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]); | |
} | |
#[test] | |
fn test_flush_on() { | |
let mut writer = BufWriter::with_capacity(8, vec![]).set_policy(FlushOn(0)); | |
assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3); | |
assert_eq!(*writer.get_ref(), &[]); | |
assert_eq!(writer.write(&[0, 4, 5]).unwrap(), 3); | |
assert_eq!(*writer.get_ref(), &[1, 2, 3, 0]); | |
assert_eq!(writer.write(&[6, 7, 8, 9, 10, 11, 12]).unwrap(), 7); | |
assert_eq!(*writer.get_ref(), &[1, 2, 3, 0, 4, 5]); | |
assert_eq!(writer.write(&[0]).unwrap(), 1); | |
assert_eq!(*writer.get_ref(), &[1, 2, 3, 0, 4, 5, 6, 7, 8, 9, 10, 11, 12, 0]); | |
} | |
#[test] | |
fn test_flush_on_newline() { | |
let mut writer = BufWriter::with_capacity(8, vec![]).set_policy(FlushOnNewline); | |
assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3); | |
assert_eq!(*writer.get_ref(), &[]); | |
assert_eq!(writer.write(&[b'\n', 4, 5]).unwrap(), 3); | |
assert_eq!(*writer.get_ref(), &[1, 2, 3, b'\n']); | |
assert_eq!(writer.write(&[6, 7, 8, 9, b'\n', 11, 12]).unwrap(), 7); | |
assert_eq!(*writer.get_ref(), &[1, 2, 3, b'\n', 4, 5, 6, 7, 8, 9, b'\n']); | |
assert_eq!(writer.write(&[b'\n']).unwrap(), 1); | |
assert_eq!(*writer.get_ref(), &[1, 2, 3, b'\n', 4, 5, 6, 7, 8, 9, b'\n', 11, 12, b'\n']); | |
} | |
} |