1 | use crate::io::AsyncBufRead; |
2 | |
3 | use pin_project_lite::pin_project; |
4 | use std::future::Future; |
5 | use std::io; |
6 | use std::marker::PhantomPinned; |
7 | use std::pin::Pin; |
8 | use std::task::{Context, Poll}; |
9 | |
10 | pin_project! { |
11 | /// Future for the [`fill_buf`](crate::io::AsyncBufReadExt::fill_buf) method. |
12 | #[derive(Debug)] |
13 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
14 | pub struct FillBuf<'a, R: ?Sized> { |
15 | reader: Option<&'a mut R>, |
16 | #[pin] |
17 | _pin: PhantomPinned, |
18 | } |
19 | } |
20 | |
21 | pub(crate) fn fill_buf<R>(reader: &mut R) -> FillBuf<'_, R> |
22 | where |
23 | R: AsyncBufRead + ?Sized + Unpin, |
24 | { |
25 | FillBuf { |
26 | reader: Some(reader), |
27 | _pin: PhantomPinned, |
28 | } |
29 | } |
30 | |
31 | impl<'a, R: AsyncBufRead + ?Sized + Unpin> Future for FillBuf<'a, R> { |
32 | type Output = io::Result<&'a [u8]>; |
33 | |
34 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
35 | let me = self.project(); |
36 | |
37 | let reader = me.reader.take().expect("Polled after completion." ); |
38 | match Pin::new(&mut *reader).poll_fill_buf(cx) { |
39 | Poll::Ready(Ok(slice)) => unsafe { |
40 | // Safety: This is necessary only due to a limitation in the |
41 | // borrow checker. Once Rust starts using the polonius borrow |
42 | // checker, this can be simplified. |
43 | // |
44 | // The safety of this transmute relies on the fact that the |
45 | // value of `reader` is `None` when we return in this branch. |
46 | // Otherwise the caller could poll us again after |
47 | // completion, and access the mutable reference while the |
48 | // returned immutable reference still exists. |
49 | let slice = std::mem::transmute::<&[u8], &'a [u8]>(slice); |
50 | Poll::Ready(Ok(slice)) |
51 | }, |
52 | Poll::Ready(Err(err)) => Poll::Ready(Err(err)), |
53 | Poll::Pending => { |
54 | *me.reader = Some(reader); |
55 | Poll::Pending |
56 | } |
57 | } |
58 | } |
59 | } |
60 | |