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