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