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}
72