|
| 1 | +use super::DEFAULT_BUF_SIZE; |
| 2 | +use crate::{AsyncBufRead, AsyncRead}; |
| 3 | +use futures_core::ready; |
| 4 | +use pin_utils::{unsafe_pinned, unsafe_unpinned}; |
| 5 | +use std::io::{self, Read}; |
| 6 | +use std::pin::Pin; |
| 7 | +use std::task::{Context, Poll}; |
| 8 | +use std::{cmp, fmt}; |
| 9 | + |
| 10 | +/// The `BufReader` struct adds buffering to any reader. |
| 11 | +/// |
| 12 | +/// It can be excessively inefficient to work directly with a [`AsyncRead`] |
| 13 | +/// instance. A `BufReader` performs large, infrequent reads on the underlying |
| 14 | +/// [`AsyncRead`] and maintains an in-memory buffer of the results. |
| 15 | +/// |
| 16 | +/// `BufReader` can improve the speed of programs that make *small* and |
| 17 | +/// *repeated* read calls to the same file or network socket. It does not |
| 18 | +/// help when reading very large amounts at once, or reading just one or a few |
| 19 | +/// times. It also provides no advantage when reading from a source that is |
| 20 | +/// already in memory, like a `Vec<u8>`. |
| 21 | +/// |
| 22 | +/// When the `BufReader` is dropped, the contents of its buffer will be |
| 23 | +/// discarded. Creating multiple instances of a `BufReader` on the same |
| 24 | +/// stream can cause data loss. |
| 25 | +/// |
| 26 | +/// [`AsyncRead`]: tokio_io::AsyncRead |
| 27 | +/// |
| 28 | +// TODO: Examples |
| 29 | +pub struct BufReader<R> { |
| 30 | + inner: R, |
| 31 | + buf: Box<[u8]>, |
| 32 | + pos: usize, |
| 33 | + cap: usize, |
| 34 | +} |
| 35 | + |
| 36 | +impl<R: AsyncRead> BufReader<R> { |
| 37 | + unsafe_pinned!(inner: R); |
| 38 | + unsafe_unpinned!(pos: usize); |
| 39 | + unsafe_unpinned!(cap: usize); |
| 40 | + |
| 41 | + /// Creates a new `BufReader` with a default buffer capacity. The default is currently 8 KB, |
| 42 | + /// but may change in the future. |
| 43 | + pub fn new(inner: R) -> Self { |
| 44 | + Self::with_capacity(DEFAULT_BUF_SIZE, inner) |
| 45 | + } |
| 46 | + |
| 47 | + /// Creates a new `BufReader` with the specified buffer capacity. |
| 48 | + pub fn with_capacity(capacity: usize, inner: R) -> Self { |
| 49 | + unsafe { |
| 50 | + let mut buffer = Vec::with_capacity(capacity); |
| 51 | + buffer.set_len(capacity); |
| 52 | + inner.prepare_uninitialized_buffer(&mut buffer); |
| 53 | + Self { |
| 54 | + inner, |
| 55 | + buf: buffer.into_boxed_slice(), |
| 56 | + pos: 0, |
| 57 | + cap: 0, |
| 58 | + } |
| 59 | + } |
| 60 | + } |
| 61 | + |
| 62 | + /// Gets a reference to the underlying reader. |
| 63 | + /// |
| 64 | + /// It is inadvisable to directly read from the underlying reader. |
| 65 | + pub fn get_ref(&self) -> &R { |
| 66 | + &self.inner |
| 67 | + } |
| 68 | + |
| 69 | + /// Gets a mutable reference to the underlying reader. |
| 70 | + /// |
| 71 | + /// It is inadvisable to directly read from the underlying reader. |
| 72 | + pub fn get_mut(&mut self) -> &mut R { |
| 73 | + &mut self.inner |
| 74 | + } |
| 75 | + |
| 76 | + /// Gets a pinned mutable reference to the underlying reader. |
| 77 | + /// |
| 78 | + /// It is inadvisable to directly read from the underlying reader. |
| 79 | + pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R> { |
| 80 | + self.inner() |
| 81 | + } |
| 82 | + |
| 83 | + /// Consumes this `BufWriter`, returning the underlying reader. |
| 84 | + /// |
| 85 | + /// Note that any leftover data in the internal buffer is lost. |
| 86 | + pub fn into_inner(self) -> R { |
| 87 | + self.inner |
| 88 | + } |
| 89 | + |
| 90 | + /// Returns a reference to the internally buffered data. |
| 91 | + /// |
| 92 | + /// Unlike `fill_buf`, this will not attempt to fill the buffer if it is empty. |
| 93 | + pub fn buffer(&self) -> &[u8] { |
| 94 | + &self.buf[self.pos..self.cap] |
| 95 | + } |
| 96 | + |
| 97 | + /// Invalidates all data in the internal buffer. |
| 98 | + #[inline] |
| 99 | + fn discard_buffer(mut self: Pin<&mut Self>) { |
| 100 | + *self.as_mut().pos() = 0; |
| 101 | + *self.cap() = 0; |
| 102 | + } |
| 103 | +} |
| 104 | + |
| 105 | +impl<R: AsyncRead> AsyncRead for BufReader<R> { |
| 106 | + fn poll_read( |
| 107 | + mut self: Pin<&mut Self>, |
| 108 | + cx: &mut Context<'_>, |
| 109 | + buf: &mut [u8], |
| 110 | + ) -> Poll<io::Result<usize>> { |
| 111 | + // If we don't have any buffered data and we're doing a massive read |
| 112 | + // (larger than our internal buffer), bypass our internal buffer |
| 113 | + // entirely. |
| 114 | + if self.pos == self.cap && buf.len() >= self.buf.len() { |
| 115 | + let res = ready!(self.as_mut().inner().poll_read(cx, buf)); |
| 116 | + self.discard_buffer(); |
| 117 | + return Poll::Ready(res); |
| 118 | + } |
| 119 | + let mut rem = ready!(self.as_mut().poll_fill_buf(cx))?; |
| 120 | + let nread = rem.read(buf)?; |
| 121 | + self.consume(nread); |
| 122 | + Poll::Ready(Ok(nread)) |
| 123 | + } |
| 124 | + |
| 125 | + // we can't skip unconditionally because of the large buffer case in read. |
| 126 | + unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool { |
| 127 | + self.inner.prepare_uninitialized_buffer(buf) |
| 128 | + } |
| 129 | +} |
| 130 | + |
| 131 | +impl<R: AsyncRead> AsyncBufRead for BufReader<R> { |
| 132 | + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> { |
| 133 | + let Self { |
| 134 | + inner, |
| 135 | + buf, |
| 136 | + cap, |
| 137 | + pos, |
| 138 | + } = unsafe { self.get_unchecked_mut() }; |
| 139 | + let mut inner = unsafe { Pin::new_unchecked(inner) }; |
| 140 | + |
| 141 | + // If we've reached the end of our internal buffer then we need to fetch |
| 142 | + // some more data from the underlying reader. |
| 143 | + // Branch using `>=` instead of the more correct `==` |
| 144 | + // to tell the compiler that the pos..cap slice is always valid. |
| 145 | + if *pos >= *cap { |
| 146 | + debug_assert!(*pos == *cap); |
| 147 | + *cap = ready!(inner.as_mut().poll_read(cx, buf))?; |
| 148 | + *pos = 0; |
| 149 | + } |
| 150 | + Poll::Ready(Ok(&buf[*pos..*cap])) |
| 151 | + } |
| 152 | + |
| 153 | + fn consume(mut self: Pin<&mut Self>, amt: usize) { |
| 154 | + *self.as_mut().pos() = cmp::min(self.pos + amt, self.cap); |
| 155 | + } |
| 156 | +} |
| 157 | + |
| 158 | +impl<R: AsyncRead + fmt::Debug> fmt::Debug for BufReader<R> { |
| 159 | + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
| 160 | + f.debug_struct("BufReader") |
| 161 | + .field("reader", &self.inner) |
| 162 | + .field( |
| 163 | + "buffer", |
| 164 | + &format_args!("{}/{}", self.cap - self.pos, self.buf.len()), |
| 165 | + ) |
| 166 | + .finish() |
| 167 | + } |
| 168 | +} |
0 commit comments