tokio/io/util/
read_buf.rs

Help
1use crate::io::AsyncRead;
2
3use bytes::BufMut;
4use pin_project_lite::pin_project;
5use std::future::Future;
6use std::io;
7use std::marker::PhantomPinned;
8use std::pin::Pin;
9use std::task::{ready, Context, Poll};
10
11pub(crate) fn read_buf<'a, R, B>(reader: &'a mut R, buf: &'a mut B) -> ReadBuf<'a, R, B>
12where
13    R: AsyncRead + Unpin + ?Sized,
14    B: BufMut + ?Sized,
15{
16    ReadBuf {
17        reader,
18        buf,
19        _pin: PhantomPinned,
20    }
21}
22
23pin_project! {
24    /// Future returned by [`read_buf`](crate::io::AsyncReadExt::read_buf).
25    #[derive(Debug)]
26    #[must_use = "futures do nothing unless you `.await` or poll them"]
27    pub struct ReadBuf<'a, R: ?Sized, B: ?Sized> {
28        reader: &'a mut R,
29        buf: &'a mut B,
30        #[pin]
31        _pin: PhantomPinned,
32    }
33}
34
35impl<R, B> Future for ReadBuf<'_, R, B>
36where
37    R: AsyncRead + Unpin + ?Sized,
38    B: BufMut + ?Sized,
39{
40    type Output = io::Result<usize>;
41
42    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<usize>> {
43        use crate::io::ReadBuf;
44
45        let me = self.project();
46
47        if !me.buf.has_remaining_mut() {
48            return Poll::Ready(Ok(0));
49        }
50
51        let n = {
52            let dst = me.buf.chunk_mut();
53            let dst = unsafe { dst.as_uninit_slice_mut() };
54            let mut buf = ReadBuf::uninit(dst);
55            let ptr = buf.filled().as_ptr();
56            ready!(Pin::new(me.reader).poll_read(cx, &mut buf)?);
57
58            // Ensure the pointer does not change from under us
59            assert_eq!(ptr, buf.filled().as_ptr());
60            buf.filled().len()
61        };
62
63        // Safety: This is guaranteed to be the number of initialized (and read)
64        // bytes due to the invariants provided by `ReadBuf::filled`.
65        unsafe {
66            me.buf.advance_mut(n);
67        }
68
69        Poll::Ready(Ok(n))
70    }
71}
tokio/io/util/read_buf.rs

tokio/io/util/
read_buf.rs
