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::{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,
14 B: BufMut,
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, B> {
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,
38 B: BufMut,
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 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