use std::io::IoSliceMut;
use std::mem;
use std::pin::Pin;
use std::str;
use cfg_if::cfg_if;
use futures::io::AsyncRead;
use crate::future::Future;
use crate::io;
use crate::task::{Context, Poll};
cfg_if! {
if #[cfg(feature = "docs")] {
#[doc(hidden)]
pub struct ImplFuture<'a, T>(std::marker::PhantomData<&'a T>);
macro_rules! ret {
($a:lifetime, $f:tt, $o:ty) => (ImplFuture<$a, $o>);
}
} else {
macro_rules! ret {
($a:lifetime, $f:tt, $o:ty) => ($f<$a, Self>);
}
}
}
pub trait Read {
fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> ret!('a, ReadFuture, io::Result<usize>)
where
Self: Unpin;
fn read_vectored<'a>(
&'a mut self,
bufs: &'a mut [IoSliceMut<'a>],
) -> ret!('a, ReadVectoredFuture, io::Result<usize>)
where
Self: Unpin,
{
ReadVectoredFuture { reader: self, bufs }
}
fn read_to_end<'a>(
&'a mut self,
buf: &'a mut Vec<u8>,
) -> ret!('a, ReadToEndFuture, io::Result<usize>)
where
Self: Unpin,
{
let start_len = buf.len();
ReadToEndFuture {
reader: self,
buf,
start_len,
}
}
fn read_to_string<'a>(
&'a mut self,
buf: &'a mut String,
) -> ret!('a, ReadToStringFuture, io::Result<usize>)
where
Self: Unpin,
{
let start_len = buf.len();
ReadToStringFuture {
reader: self,
bytes: unsafe { mem::replace(buf.as_mut_vec(), Vec::new()) },
buf,
start_len,
}
}
fn read_exact<'a>(&'a mut self, buf: &'a mut [u8]) -> ret!('a, ReadExactFuture, io::Result<()>)
where
Self: Unpin,
{
ReadExactFuture { reader: self, buf }
}
}
impl<T: AsyncRead + Unpin + ?Sized> Read for T {
fn read<'a>(&'a mut self, buf: &'a mut [u8]) -> ret!('a, ReadFuture, io::Result<usize>) {
ReadFuture { reader: self, buf }
}
}
#[doc(hidden)]
#[allow(missing_debug_implementations)]
pub struct ReadFuture<'a, T: Unpin + ?Sized> {
reader: &'a mut T,
buf: &'a mut [u8],
}
impl<T: AsyncRead + Unpin + ?Sized> Future for ReadFuture<'_, T> {
type Output = io::Result<usize>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let Self { reader, buf } = &mut *self;
Pin::new(reader).poll_read(cx, buf)
}
}
#[doc(hidden)]
#[allow(missing_debug_implementations)]
pub struct ReadVectoredFuture<'a, T: Unpin + ?Sized> {
reader: &'a mut T,
bufs: &'a mut [IoSliceMut<'a>],
}
impl<T: AsyncRead + Unpin + ?Sized> Future for ReadVectoredFuture<'_, T> {
type Output = io::Result<usize>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let Self { reader, bufs } = &mut *self;
Pin::new(reader).poll_read_vectored(cx, bufs)
}
}
#[doc(hidden)]
#[allow(missing_debug_implementations)]
pub struct ReadToEndFuture<'a, T: Unpin + ?Sized> {
reader: &'a mut T,
buf: &'a mut Vec<u8>,
start_len: usize,
}
impl<T: AsyncRead + Unpin + ?Sized> Future for ReadToEndFuture<'_, T> {
type Output = io::Result<usize>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let Self {
reader,
buf,
start_len,
} = &mut *self;
read_to_end_internal(Pin::new(reader), cx, buf, *start_len)
}
}
#[doc(hidden)]
#[allow(missing_debug_implementations)]
pub struct ReadToStringFuture<'a, T: Unpin + ?Sized> {
reader: &'a mut T,
buf: &'a mut String,
bytes: Vec<u8>,
start_len: usize,
}
impl<T: AsyncRead + Unpin + ?Sized> Future for ReadToStringFuture<'_, T> {
type Output = io::Result<usize>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let Self {
reader,
buf,
bytes,
start_len,
} = &mut *self;
let reader = Pin::new(reader);
let ret = futures::ready!(read_to_end_internal(reader, cx, bytes, *start_len));
if str::from_utf8(&bytes).is_err() {
Poll::Ready(ret.and_then(|_| {
Err(io::Error::new(
io::ErrorKind::InvalidData,
"stream did not contain valid UTF-8",
))
}))
} else {
debug_assert!(buf.is_empty());
mem::swap(unsafe { buf.as_mut_vec() }, bytes);
Poll::Ready(ret)
}
}
}
#[doc(hidden)]
#[allow(missing_debug_implementations)]
pub struct ReadExactFuture<'a, T: Unpin + ?Sized> {
reader: &'a mut T,
buf: &'a mut [u8],
}
impl<T: AsyncRead + Unpin + ?Sized> Future for ReadExactFuture<'_, T> {
type Output = io::Result<()>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let Self { reader, buf } = &mut *self;
while !buf.is_empty() {
let n = futures::ready!(Pin::new(&mut *reader).poll_read(cx, buf))?;
let (_, rest) = mem::replace(buf, &mut []).split_at_mut(n);
*buf = rest;
if n == 0 {
return Poll::Ready(Err(io::ErrorKind::UnexpectedEof.into()));
}
}
Poll::Ready(Ok(()))
}
}
pub fn read_to_end_internal<R: AsyncRead + ?Sized>(
mut rd: Pin<&mut R>,
cx: &mut Context<'_>,
buf: &mut Vec<u8>,
start_len: usize,
) -> Poll<io::Result<usize>> {
struct Guard<'a> {
buf: &'a mut Vec<u8>,
len: usize,
}
impl Drop for Guard<'_> {
fn drop(&mut self) {
unsafe {
self.buf.set_len(self.len);
}
}
}
let mut g = Guard {
len: buf.len(),
buf,
};
let ret;
loop {
if g.len == g.buf.len() {
unsafe {
g.buf.reserve(32);
let capacity = g.buf.capacity();
g.buf.set_len(capacity);
rd.initializer().initialize(&mut g.buf[g.len..]);
}
}
match futures::ready!(rd.as_mut().poll_read(cx, &mut g.buf[g.len..])) {
Ok(0) => {
ret = Poll::Ready(Ok(g.len - start_len));
break;
}
Ok(n) => g.len += n,
Err(e) => {
ret = Poll::Ready(Err(e));
break;
}
}
}
ret
}