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
13 use 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)]
20 pub struct DoRead(pub bool);
21
22 /// Shorthand for `return DoRead(bool)` or `return DoRead(true)` (empty invocation)
23 #[macro_export]
24 macro_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)]
35 pub struct StdPolicy;
36
37 /// Trait that governs `BufReader`'s behavior.
38 pub 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.
before_read(&mut self, buffer: &mut Buffer) -> DoRead58 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.
after_consume(&mut self, _buffer: &mut Buffer, _amt: usize)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.
70 impl 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)]
106 pub struct MinBuffered(pub usize);
107
108 impl MinBuffered {
109 /// Set the number of bytes to ensure are in the buffer.
set_min(&mut self, min: usize)110 pub fn set_min(&mut self, min: usize) {
111 self.0 = min;
112 }
113 }
114
115 impl ReaderPolicy for MinBuffered {
before_read(&mut self, buffer: &mut Buffer) -> DoRead116 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)]
138 pub struct FlushAmt(pub usize);
139
140 /// Shorthand for `return FlushAmt(n)` or `return FlushAmt(0)` (empty invocation)
141 #[macro_export]
142 macro_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.
148 pub 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.
before_write(&mut self, buf: &mut Buffer, incoming: usize) -> FlushAmt157 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.
after_write(&mut self, _buf: &Buffer) -> FlushAmt166 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.
173 impl WriterPolicy for StdPolicy {}
174
175 /// Flush the buffer if it contains at least the given number of bytes.
176 #[derive(Debug, Default)]
177 pub struct FlushAtLeast(pub usize);
178
179 impl WriterPolicy for FlushAtLeast {
before_write(&mut self, buf: &mut Buffer, incoming: usize) -> FlushAmt180 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
after_write(&mut self, buf: &Buffer) -> FlushAmt185 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)]
192 pub struct FlushExact(pub usize);
193
194 impl 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)
before_write(&mut self, buf: &mut Buffer, incoming: usize) -> FlushAmt199 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.
after_write(&mut self, _buf: &Buffer) -> FlushAmt212 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)]
221 pub struct FlushOn(pub u8);
222
223 impl WriterPolicy for FlushOn {
after_write(&mut self, buf: &Buffer) -> FlushAmt224 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)]
234 pub struct FlushOnNewline;
235
236 impl WriterPolicy for FlushOnNewline {
after_write(&mut self, buf: &Buffer) -> FlushAmt237 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
ensure_capacity(buf: &mut Buffer, min_cap: usize)242 fn 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)]
251 mod test {
252 use {BufReader, BufWriter};
253 use policy::*;
254 use std::io::{BufRead, Cursor, Write};
255
256 #[test]
test_min_buffered()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]
test_flush_at_least()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]
test_flush_exact()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]
test_flush_on()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]
test_flush_on_newline()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 }
376