buf_redux/
policy.rs

Help
1// Copyright 2016-2018 Austin Bonander <austin.bonander@gmail.com>
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8//! Types which can be used to tune the behavior of `BufReader` and `BufWriter`.
9//!
10//! Some simple policies are provided for your convenience. You may prefer to create your own
11//! types and implement the traits for them instead.
12
13use super::Buffer;
14
15/// Flag for `ReaderPolicy` methods to signal whether or not `BufReader` should read into
16/// the buffer.
17///
18/// See `do_read!()` for a shorthand.
19#[derive(Copy, Clone, Debug)]
20pub struct DoRead(pub bool);
21
22/// Shorthand for `return DoRead(bool)` or `return DoRead(true)` (empty invocation)
23#[macro_export]
24macro_rules! do_read (
25    ($val:expr) => ( return $crate::policy::DoRead($val); );
26    () => ( do_read!(true); )
27);
28
29/// Default policy for both `BufReader` and `BufWriter` that reproduces the behaviors of their
30/// `std::io` counterparts:
31///
32/// * `BufReader`: only reads when the buffer is empty, does not resize or move data.
33/// * `BufWriter`: only flushes the buffer when there is not enough room for an incoming write.
34#[derive(Debug, Default)]
35pub struct StdPolicy;
36
37/// Trait that governs `BufReader`'s behavior.
38pub trait ReaderPolicy {
39    /// Consulted before attempting to read into the buffer.
40    ///
41    /// Return `DoRead(true)` to issue a read into the buffer before reading data out of it,
42    /// or `DoRead(false)` to read from the buffer as it is, even if it's empty.
43    /// `do_read!()` is provided as a shorthand.
44    ///
45    /// If there is no room in the buffer after this method is called,
46    /// the buffer will not be read into (so if the buffer is full but you want more data
47    /// you should call `.make_room()` or reserve more space). If there *is* room, `BufReader` will
48    /// attempt to read into the buffer. If successful (`Ok(x)` where `x > 0` is returned), this
49    /// method will be consulted again for another read attempt.
50    ///
51    /// By default, this implements `std::io::BufReader`'s behavior: only read into the buffer if
52    /// it is empty.
53    ///
54    /// ### Note
55    /// If the read will ignore the buffer entirely (if the buffer is empty and the amount to be
56    /// read matches or exceeds its capacity) or if `BufReader::read_into_buf()` was called to force
57    /// a read into the buffer manually, this method will not be called.
58    fn before_read(&mut self, buffer: &mut Buffer) -> DoRead { DoRead(buffer.len() == 0) }
59
60    /// Called after bytes are consumed from the buffer.
61    ///
62    /// Supplies the true amount consumed if the amount passed to `BufReader::consume`
63    /// was in excess.
64    ///
65    /// This is a no-op by default.
66    fn after_consume(&mut self, _buffer: &mut Buffer, _amt: usize) {}
67}
68
69/// Behavior of `std::io::BufReader`: the buffer will only be read into if it is empty.
70impl ReaderPolicy for StdPolicy {}
71
72/// A policy for [`BufReader`](::BufReader) which ensures there is at least the given number of
73/// bytes in  the buffer, failing this only if the reader is at EOF.
74///
75/// If the minimum buffer length is greater than the buffer capacity, it will be resized.
76///
77/// ### Example
78/// ```rust
79/// use buf_redux::BufReader;
80/// use buf_redux::policy::MinBuffered;
81/// use std::io::{BufRead, Cursor};
82/// 
83/// let data = (1 .. 16).collect::<Vec<u8>>();
84///
85/// // normally you should use `BufReader::new()` or give a capacity of several KiB or more
86/// let mut reader = BufReader::with_capacity(8, Cursor::new(data))
87///     // always at least 4 bytes in the buffer (or until the source is empty)
88///     .set_policy(MinBuffered(4)); // always at least 4 bytes in the buffer
89///
90/// // first buffer fill, same as `std::io::BufReader`
91/// assert_eq!(reader.fill_buf().unwrap(), &[1, 2, 3, 4, 5, 6, 7, 8]);
92/// reader.consume(3);
93///
94/// // enough data in the buffer, another read isn't done yet
95/// assert_eq!(reader.fill_buf().unwrap(), &[4, 5, 6, 7, 8]);
96/// reader.consume(4);
97///
98/// // `std::io::BufReader` would return `&[8]`
99/// assert_eq!(reader.fill_buf().unwrap(), &[8, 9, 10, 11, 12, 13, 14, 15]);
100/// reader.consume(5);
101///
102/// // no data left in the reader
103/// assert_eq!(reader.fill_buf().unwrap(), &[13, 14, 15]);
104/// ```
105#[derive(Debug)]
106pub struct MinBuffered(pub usize);
107
108impl MinBuffered {
109    /// Set the number of bytes to ensure are in the buffer.
110    pub fn set_min(&mut self, min: usize) {
111        self.0 = min;
112    }
113}
114
115impl ReaderPolicy for MinBuffered {
116    fn before_read(&mut self, buffer: &mut Buffer) -> DoRead {
117        // do nothing if we have enough data
118        if buffer.len() >= self.0 { do_read!(false) }
119
120        let cap = buffer.capacity();
121
122        // if there's enough room but some of it's stuck after the head
123        if buffer.usable_space() < self.0 && buffer.free_space() >= self.0 {
124            buffer.make_room();
125        } else if cap < self.0 {
126            buffer.reserve(self.0 - cap);
127        }
128
129        DoRead(true)
130    }
131}
132
133/// Flag for `WriterPolicy` methods to tell `BufWriter` how many bytes to flush to the
134/// underlying reader.
135///
136/// See `flush_amt!()` for a shorthand.
137#[derive(Copy, Clone, Debug)]
138pub struct FlushAmt(pub usize);
139
140/// Shorthand for `return FlushAmt(n)` or `return FlushAmt(0)` (empty invocation)
141#[macro_export]
142macro_rules! flush_amt (
143    ($n:expr) => ( return $crate::policy::FlushAmt($n); );
144    () => ( flush_amt!(0); )
145);
146
147/// A trait which tells `BufWriter` when to flush.
148pub trait WriterPolicy {
149    /// Return `FlushAmt(n > 0)` if the buffer should be flushed before reading into it.
150    /// If the returned amount is 0 or greater than the amount of buffered data, no flush is
151    /// performed.
152    ///
153    /// The buffer is provided, as well as `incoming` which is
154    /// the size of the buffer that will be written to the `BufWriter`.
155    ///
156    /// By default, flushes the buffer if the usable space is smaller than the incoming write.
157    fn before_write(&mut self, buf: &mut Buffer, incoming: usize) -> FlushAmt {
158        FlushAmt(if incoming > buf.usable_space() { buf.len() } else { 0 })
159    }
160
161    /// Return `true` if the buffer should be flushed after reading into it.
162    ///
163    /// `buf` references the updated buffer after the read.
164    ///
165    /// Default impl is a no-op.
166    fn after_write(&mut self, _buf: &Buffer) -> FlushAmt {
167        FlushAmt(0)
168    }
169}
170
171/// Default behavior of `std::io::BufWriter`: flush before a read into the buffer
172/// only if the incoming data is larger than the buffer's writable space.
173impl WriterPolicy for StdPolicy {}
174
175/// Flush the buffer if it contains at least the given number of bytes.
176#[derive(Debug, Default)]
177pub struct FlushAtLeast(pub usize);
178
179impl WriterPolicy for FlushAtLeast {
180    fn before_write(&mut self, buf: &mut Buffer, incoming: usize) -> FlushAmt {
181        ensure_capacity(buf, self.0);
182        FlushAmt(if incoming > buf.usable_space() { buf.len() } else { 0 })
183    }
184
185    fn after_write(&mut self, buf: &Buffer) -> FlushAmt {
186        FlushAmt(::std::cmp::max(buf.len(), self.0))
187    }
188}
189
190/// Only ever flush exactly the given number of bytes, until the writer is empty.
191#[derive(Debug, Default)]
192pub struct FlushExact(pub usize);
193
194impl WriterPolicy for FlushExact {
195    /// Flushes the buffer if there is not enough room to fit `incoming` bytes,
196    /// but only when the buffer contains at least `self.0` bytes.
197    ///
198    /// Otherwise, calls [`Buffer::make_room()`](::Buffer::make_room)
199    fn before_write(&mut self, buf: &mut Buffer, incoming: usize) -> FlushAmt {
200        ensure_capacity(buf, self.0);
201
202        // don't have enough room to fit the additional bytes but we can't flush,
203        // then make room for (at least some of) the incoming bytes.
204        if incoming > buf.usable_space() && buf.len() < self.0 {
205            buf.make_room();
206        }
207
208        FlushAmt(self.0)
209    }
210
211    /// Flushes the given amount if possible, nothing otherwise.
212    fn after_write(&mut self, _buf: &Buffer) -> FlushAmt {
213        FlushAmt(self.0)
214    }
215}
216
217/// Flush the buffer if it contains the given byte.
218///
219/// Only scans the buffer after reading. Searches from the end first.
220#[derive(Debug, Default)]
221pub struct FlushOn(pub u8);
222
223impl WriterPolicy for FlushOn {
224    fn after_write(&mut self, buf: &Buffer) -> FlushAmt {
225        // include the delimiter in the flush
226        FlushAmt(::memchr::memrchr(self.0, buf.buf()).map_or(0, |n| n + 1))
227    }
228}
229
230/// Flush the buffer if it contains a newline (`\n`).
231///
232/// Equivalent to `FlushOn(b'\n')`.
233#[derive(Debug, Default)]
234pub struct FlushOnNewline;
235
236impl WriterPolicy for FlushOnNewline {
237    fn after_write(&mut self, buf: &Buffer) -> FlushAmt {
238        FlushAmt(::memchr::memrchr(b'\n', buf.buf()).map_or(0, |n| n + 1))
239    }
240}
241
242fn ensure_capacity(buf: &mut Buffer, min_cap: usize) {
243    let cap = buf.capacity();
244
245    if cap < min_cap {
246        buf.reserve(min_cap - cap);
247    }
248}
249
250#[cfg(test)]
251mod test {
252    use {BufReader, BufWriter};
253    use policy::*;
254    use std::io::{BufRead, Cursor, Write};
255
256    #[test]
257    fn test_min_buffered() {
258        let min_buffered = 4;
259        let data = (0 .. 20).collect::<Vec<u8>>();
260        // create a reader with 0 capacity
261        let mut reader = BufReader::with_capacity(0, Cursor::new(data))
262            .set_policy(MinBuffered(min_buffered));
263
264        // policy reserves the required space in the buffer
265        assert_eq!(reader.fill_buf().unwrap(), &[0, 1, 2, 3][..]);
266        assert_eq!(reader.capacity(), min_buffered);
267
268        // double the size now that the buffer's full
269        reader.reserve(min_buffered);
270        assert_eq!(reader.capacity(), min_buffered * 2);
271
272        // we haven't consumed anything, the reader should have the same data
273        assert_eq!(reader.fill_buf().unwrap(), &[0, 1, 2, 3]);
274        reader.consume(2);
275        // policy read more data, `std::io::BufReader` doesn't do that
276        assert_eq!(reader.fill_buf().unwrap(), &[2, 3, 4, 5, 6, 7]);
277        reader.consume(4);
278        // policy made room and read more
279        assert_eq!(reader.fill_buf().unwrap(), &[6, 7, 8, 9, 10, 11, 12, 13]);
280        reader.consume(4);
281        assert_eq!(reader.fill_buf().unwrap(), &[10, 11, 12, 13]);
282        reader.consume(2);
283        assert_eq!(reader.fill_buf().unwrap(), &[12, 13, 14, 15, 16, 17, 18, 19]);
284        reader.consume(8);
285        assert_eq!(reader.fill_buf().unwrap(), &[])
286    }
287
288    #[test]
289    fn test_flush_at_least() {
290        let flush_min = 4;
291
292        let mut writer = BufWriter::with_capacity(0, vec![]).set_policy(FlushAtLeast(flush_min));
293        assert_eq!(writer.capacity(), 0);
294        assert_eq!(writer.write(&[1]).unwrap(), 1);
295        // policy reserved space for writing
296        assert_eq!(writer.capacity(), flush_min);
297        // one byte in buffer, we want to double our capacity
298        writer.reserve(flush_min * 2 - 1);
299        assert_eq!(writer.capacity(), flush_min * 2);
300
301        assert_eq!(writer.write(&[2, 3]).unwrap(), 2);
302        // no flush yet, only 3 bytes in buffer
303        assert_eq!(*writer.get_ref(), &[]);
304
305        assert_eq!(writer.write(&[4, 5, 6]).unwrap(), 3);
306        // flushed all
307        assert_eq!(*writer.get_ref(), &[1, 2, 3, 4, 5, 6]);
308
309        assert_eq!(writer.write(&[7, 8, 9]).unwrap(), 3);
310        // `.into_inner()` should flush always
311        assert_eq!(writer.into_inner().unwrap(), &[1, 2, 3, 4, 5, 6, 7, 8, 9]);
312    }
313
314    #[test]
315    fn test_flush_exact() {
316        let flush_exact = 4;
317
318        let mut writer = BufWriter::with_capacity(0, vec![]).set_policy(FlushExact(flush_exact));
319        assert_eq!(writer.capacity(), 0);
320        assert_eq!(writer.write(&[1]).unwrap(), 1);
321        // policy reserved space for writing
322        assert_eq!(writer.capacity(), flush_exact);
323        // one byte in buffer, we want to double our capacity
324        writer.reserve(flush_exact * 2 - 1);
325        assert_eq!(writer.capacity(), flush_exact * 2);
326
327        assert_eq!(writer.write(&[2, 3]).unwrap(), 2);
328        // no flush yet, only 3 bytes in buffer
329        assert_eq!(*writer.get_ref(), &[]);
330
331        assert_eq!(writer.write(&[4, 5, 6]).unwrap(), 3);
332        // flushed exactly 4 bytes
333        assert_eq!(*writer.get_ref(), &[1, 2, 3, 4]);
334
335        assert_eq!(writer.write(&[7, 8, 9, 10]).unwrap(), 4);
336        // flushed another 4 bytes
337        assert_eq!(*writer.get_ref(), &[1, 2, 3, 4, 5, 6, 7, 8]);
338        // `.into_inner()` should flush always
339        assert_eq!(writer.into_inner().unwrap(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
340    }
341
342    #[test]
343    fn test_flush_on() {
344        let mut writer = BufWriter::with_capacity(8, vec![]).set_policy(FlushOn(0));
345
346        assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
347        assert_eq!(*writer.get_ref(), &[]);
348
349        assert_eq!(writer.write(&[0, 4, 5]).unwrap(), 3);
350        assert_eq!(*writer.get_ref(), &[1, 2, 3, 0]);
351
352        assert_eq!(writer.write(&[6, 7, 8, 9, 10, 11, 12]).unwrap(), 7);
353        assert_eq!(*writer.get_ref(), &[1, 2, 3, 0, 4, 5]);
354
355        assert_eq!(writer.write(&[0]).unwrap(), 1);
356        assert_eq!(*writer.get_ref(), &[1, 2, 3, 0, 4, 5, 6, 7, 8, 9, 10, 11, 12, 0]);
357    }
358
359    #[test]
360    fn test_flush_on_newline() {
361        let mut writer = BufWriter::with_capacity(8, vec![]).set_policy(FlushOnNewline);
362
363        assert_eq!(writer.write(&[1, 2, 3]).unwrap(), 3);
364        assert_eq!(*writer.get_ref(), &[]);
365
366        assert_eq!(writer.write(&[b'\n', 4, 5]).unwrap(), 3);
367        assert_eq!(*writer.get_ref(), &[1, 2, 3, b'\n']);
368
369        assert_eq!(writer.write(&[6, 7, 8, 9, b'\n', 11, 12]).unwrap(), 7);
370        assert_eq!(*writer.get_ref(), &[1, 2, 3, b'\n', 4, 5, 6, 7, 8, 9, b'\n']);
371
372        assert_eq!(writer.write(&[b'\n']).unwrap(), 1);
373        assert_eq!(*writer.get_ref(), &[1, 2, 3, b'\n', 4, 5, 6, 7, 8, 9, b'\n', 11, 12, b'\n']);
374    }
375}
buf_redux/policy.rs

buf_redux/
policy.rs
