nuclei/io/util/
buffered.rs

1//! Buffering around a reader or a writer: [`BufReader`], [`BufWriter`], and
2//! both at once, [`BufStream`].
3
4use core::pin::Pin;
5use core::task::{Context, Poll, ready};
6use std::io::{self, IoSlice, SeekFrom};
7
8use futures_io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite};
9
10/// The buffer the wrappers make when no capacity is given.
11const DEFAULT_CAPACITY: usize = 8 * 1024;
12
13/// A reader with a buffer in front: small reads take from the buffer, which
14/// is filled by large reads of the reader. Writes, if the reader is also a
15/// writer, go straight through.
16#[derive(Debug)]
17pub struct BufReader<R> {
18    inner: R,
19    buf: Box<[u8]>,
20    pos: usize,
21    cap: usize,
22}
23
24impl<R> BufReader<R> {
25    /// A reader with a buffer of 8 KiB in front of `inner`.
26    pub fn new(inner: R) -> Self {
27        Self::with_capacity(DEFAULT_CAPACITY, inner)
28    }
29
30    /// A reader with a buffer of `capacity` bytes in front of `inner`.
31    pub fn with_capacity(capacity: usize, inner: R) -> Self {
32        Self {
33            inner,
34            buf: vec![0; capacity].into_boxed_slice(),
35            pos: 0,
36            cap: 0,
37        }
38    }
39
40    /// The reader inside.
41    pub fn get_ref(&self) -> &R {
42        &self.inner
43    }
44
45    /// The reader inside, to change (reading it directly skips the bytes
46    /// buffered).
47    pub fn get_mut(&mut self) -> &mut R {
48        &mut self.inner
49    }
50
51    /// The reader inside, pinned.
52    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut R>
53    where
54        R: Unpin,
55    {
56        Pin::new(&mut self.get_mut().inner)
57    }
58
59    /// The reader inside, taken back; the bytes buffered are lost.
60    pub fn into_inner(self) -> R {
61        self.inner
62    }
63
64    /// The bytes buffered and not read yet.
65    pub fn buffer(&self) -> &[u8] {
66        &self.buf[self.pos..self.cap]
67    }
68
69    fn discard_buffer(&mut self) {
70        self.pos = 0;
71        self.cap = 0;
72    }
73}
74
75impl<R: AsyncRead + Unpin> AsyncRead for BufReader<R> {
76    fn poll_read(
77        self: Pin<&mut Self>,
78        cx: &mut Context<'_>,
79        out: &mut [u8],
80    ) -> Poll<io::Result<usize>> {
81        let this = self.get_mut();
82        // An empty buffer and a read at least its size: the buffer would only
83        // add a copy.
84        if this.pos == this.cap && out.len() >= this.buf.len() {
85            let n = ready!(Pin::new(&mut this.inner).poll_read(cx, out))?;
86            this.discard_buffer();
87            return Poll::Ready(Ok(n));
88        }
89        let available = ready!(Pin::new(&mut *this).poll_fill_buf(cx))?;
90        let n = available.len().min(out.len());
91        out[..n].copy_from_slice(&available[..n]);
92        Pin::new(this).consume(n);
93        Poll::Ready(Ok(n))
94    }
95}
96
97impl<R: AsyncRead + Unpin> AsyncBufRead for BufReader<R> {
98    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
99        let this = self.get_mut();
100        if this.pos >= this.cap {
101            let n = ready!(Pin::new(&mut this.inner).poll_read(cx, &mut this.buf))?;
102            this.cap = n;
103            this.pos = 0;
104        }
105        Poll::Ready(Ok(&this.buf[this.pos..this.cap]))
106    }
107
108    fn consume(self: Pin<&mut Self>, amount: usize) {
109        let this = self.get_mut();
110        this.pos = (this.pos + amount).min(this.cap);
111    }
112}
113
114impl<R: AsyncRead + AsyncSeek + Unpin> AsyncSeek for BufReader<R> {
115    /// Seeks the reader; a seek relative to the current position counts
116    /// from the position of the next byte read through the buffer. The
117    /// buffer is emptied.
118    fn poll_seek(
119        self: Pin<&mut Self>,
120        cx: &mut Context<'_>,
121        pos: SeekFrom,
122    ) -> Poll<io::Result<u64>> {
123        let this = self.get_mut();
124        let pos = match pos {
125            SeekFrom::Current(n) => {
126                let buffered = (this.cap - this.pos) as i64;
127                match n.checked_sub(buffered) {
128                    Some(n) => SeekFrom::Current(n),
129                    None => {
130                        // Past `i64`'s range: two seeks, the buffered part first.
131                        ready!(
132                            Pin::new(&mut this.inner).poll_seek(cx, SeekFrom::Current(-buffered))
133                        )?;
134                        this.discard_buffer();
135                        SeekFrom::Current(n)
136                    }
137                }
138            }
139            other => other,
140        };
141        let at = ready!(Pin::new(&mut this.inner).poll_seek(cx, pos))?;
142        this.discard_buffer();
143        Poll::Ready(Ok(at))
144    }
145}
146
147impl<R: AsyncWrite + Unpin> AsyncWrite for BufReader<R> {
148    fn poll_write(
149        self: Pin<&mut Self>,
150        cx: &mut Context<'_>,
151        buf: &[u8],
152    ) -> Poll<io::Result<usize>> {
153        Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
154    }
155
156    fn poll_write_vectored(
157        self: Pin<&mut Self>,
158        cx: &mut Context<'_>,
159        bufs: &[IoSlice<'_>],
160    ) -> Poll<io::Result<usize>> {
161        Pin::new(&mut self.get_mut().inner).poll_write_vectored(cx, bufs)
162    }
163
164    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
165        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
166    }
167
168    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
169        Pin::new(&mut self.get_mut().inner).poll_close(cx)
170    }
171}
172
173/// A writer with a buffer in front: small writes gather in the buffer, which
174/// goes to the writer when full, on a flush, or on a close. Reads, if the
175/// writer is also a reader, go straight through. Dropping it loses what is
176/// still buffered: flush or close it first.
177#[derive(Debug)]
178pub struct BufWriter<W> {
179    inner: W,
180    buf: Vec<u8>,
181    written: usize,
182}
183
184impl<W> BufWriter<W> {
185    /// A writer with a buffer of 8 KiB in front of `inner`.
186    pub fn new(inner: W) -> Self {
187        Self::with_capacity(DEFAULT_CAPACITY, inner)
188    }
189
190    /// A writer with a buffer of `capacity` bytes in front of `inner`.
191    pub fn with_capacity(capacity: usize, inner: W) -> Self {
192        Self {
193            inner,
194            buf: Vec::with_capacity(capacity),
195            written: 0,
196        }
197    }
198
199    /// The writer inside.
200    pub fn get_ref(&self) -> &W {
201        &self.inner
202    }
203
204    /// The writer inside, to change (writing it directly puts the bytes
205    /// ahead of what is buffered).
206    pub fn get_mut(&mut self) -> &mut W {
207        &mut self.inner
208    }
209
210    /// The writer inside, pinned.
211    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut W>
212    where
213        W: Unpin,
214    {
215        Pin::new(&mut self.get_mut().inner)
216    }
217
218    /// The writer inside, taken back; the bytes buffered are lost.
219    pub fn into_inner(self) -> W {
220        self.inner
221    }
222
223    /// The bytes buffered and not written to the writer yet.
224    pub fn buffer(&self) -> &[u8] {
225        &self.buf[self.written..]
226    }
227}
228
229impl<W: AsyncWrite + Unpin> BufWriter<W> {
230    /// Writes the buffer out to the writer, all of it.
231    fn flush_buf(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
232        while self.written < self.buf.len() {
233            let n = ready!(Pin::new(&mut self.inner).poll_write(cx, &self.buf[self.written..]))?;
234            if n == 0 {
235                return Poll::Ready(Err(io::Error::new(
236                    io::ErrorKind::WriteZero,
237                    "the writer took no more of the buffered bytes",
238                )));
239            }
240            self.written += n;
241        }
242        self.buf.clear();
243        self.written = 0;
244        Poll::Ready(Ok(()))
245    }
246}
247
248impl<W: AsyncWrite + Unpin> AsyncWrite for BufWriter<W> {
249    fn poll_write(
250        self: Pin<&mut Self>,
251        cx: &mut Context<'_>,
252        buf: &[u8],
253    ) -> Poll<io::Result<usize>> {
254        let this = self.get_mut();
255        if this.buf.len() + buf.len() > this.buf.capacity() {
256            ready!(this.flush_buf(cx))?;
257        }
258        if buf.len() >= this.buf.capacity() {
259            // As large as the buffer: straight to the writer.
260            return Pin::new(&mut this.inner).poll_write(cx, buf);
261        }
262        this.buf.extend_from_slice(buf);
263        Poll::Ready(Ok(buf.len()))
264    }
265
266    fn poll_write_vectored(
267        self: Pin<&mut Self>,
268        cx: &mut Context<'_>,
269        bufs: &[IoSlice<'_>],
270    ) -> Poll<io::Result<usize>> {
271        let this = self.get_mut();
272        let total: usize = bufs.iter().map(|b| b.len()).sum();
273        if this.buf.len() + total > this.buf.capacity() {
274            ready!(this.flush_buf(cx))?;
275        }
276        if total >= this.buf.capacity() {
277            return Pin::new(&mut this.inner).poll_write_vectored(cx, bufs);
278        }
279        for b in bufs {
280            this.buf.extend_from_slice(b);
281        }
282        Poll::Ready(Ok(total))
283    }
284
285    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
286        let this = self.get_mut();
287        ready!(this.flush_buf(cx))?;
288        Pin::new(&mut this.inner).poll_flush(cx)
289    }
290
291    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
292        let this = self.get_mut();
293        ready!(this.flush_buf(cx))?;
294        Pin::new(&mut this.inner).poll_close(cx)
295    }
296}
297
298impl<W: AsyncRead + Unpin> AsyncRead for BufWriter<W> {
299    fn poll_read(
300        self: Pin<&mut Self>,
301        cx: &mut Context<'_>,
302        buf: &mut [u8],
303    ) -> Poll<io::Result<usize>> {
304        Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
305    }
306}
307
308impl<W: AsyncBufRead + Unpin> AsyncBufRead for BufWriter<W> {
309    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
310        Pin::new(&mut self.get_mut().inner).poll_fill_buf(cx)
311    }
312
313    fn consume(self: Pin<&mut Self>, amount: usize) {
314        Pin::new(&mut self.get_mut().inner).consume(amount);
315    }
316}
317
318impl<W: AsyncWrite + AsyncSeek + Unpin> AsyncSeek for BufWriter<W> {
319    /// Writes the buffer out, then seeks the writer.
320    fn poll_seek(
321        self: Pin<&mut Self>,
322        cx: &mut Context<'_>,
323        pos: SeekFrom,
324    ) -> Poll<io::Result<u64>> {
325        let this = self.get_mut();
326        ready!(this.flush_buf(cx))?;
327        Pin::new(&mut this.inner).poll_seek(cx, pos)
328    }
329}
330
331/// A stream with a read buffer and a write buffer in front: a
332/// [`BufReader`] around a [`BufWriter`].
333#[derive(Debug)]
334pub struct BufStream<S> {
335    inner: BufReader<BufWriter<S>>,
336}
337
338impl<S> BufStream<S> {
339    /// Buffers of 8 KiB each way in front of `stream`.
340    pub fn new(stream: S) -> Self {
341        Self {
342            inner: BufReader::new(BufWriter::new(stream)),
343        }
344    }
345
346    /// A read buffer of `read` bytes and a write buffer of `write` bytes in
347    /// front of `stream`.
348    pub fn with_capacity(read: usize, write: usize, stream: S) -> Self {
349        Self {
350            inner: BufReader::with_capacity(read, BufWriter::with_capacity(write, stream)),
351        }
352    }
353
354    /// The stream inside.
355    pub fn get_ref(&self) -> &S {
356        self.inner.get_ref().get_ref()
357    }
358
359    /// The stream inside, to change.
360    pub fn get_mut(&mut self) -> &mut S {
361        self.inner.get_mut().get_mut()
362    }
363
364    /// The stream inside, taken back; the bytes buffered either way are lost.
365    pub fn into_inner(self) -> S {
366        self.inner.into_inner().into_inner()
367    }
368}
369
370impl<S: AsyncRead + AsyncWrite + Unpin> AsyncRead for BufStream<S> {
371    fn poll_read(
372        self: Pin<&mut Self>,
373        cx: &mut Context<'_>,
374        buf: &mut [u8],
375    ) -> Poll<io::Result<usize>> {
376        Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
377    }
378}
379
380impl<S: AsyncRead + AsyncWrite + Unpin> AsyncBufRead for BufStream<S> {
381    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
382        Pin::new(&mut self.get_mut().inner).poll_fill_buf(cx)
383    }
384
385    fn consume(self: Pin<&mut Self>, amount: usize) {
386        Pin::new(&mut self.get_mut().inner).consume(amount);
387    }
388}
389
390impl<S: AsyncRead + AsyncWrite + Unpin> AsyncWrite for BufStream<S> {
391    fn poll_write(
392        self: Pin<&mut Self>,
393        cx: &mut Context<'_>,
394        buf: &[u8],
395    ) -> Poll<io::Result<usize>> {
396        Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
397    }
398
399    fn poll_write_vectored(
400        self: Pin<&mut Self>,
401        cx: &mut Context<'_>,
402        bufs: &[IoSlice<'_>],
403    ) -> Poll<io::Result<usize>> {
404        Pin::new(&mut self.get_mut().inner).poll_write_vectored(cx, bufs)
405    }
406
407    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
408        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
409    }
410
411    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
412        Pin::new(&mut self.get_mut().inner).poll_close(cx)
413    }
414}
nuclei/io/util/buffered.rs

nuclei/io/util/
buffered.rs
