src/lib/fuchsia-async/src/temp.rs - fuchsia - Git at Google

 // Copyright 2018 The Fuchsia Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 //! Temporary Futures extensions used during the transition from 0.2 -> 0.3.
 //! These SHOULD NOT be used for new code. Instead, use the corresponding 0.3
 //! functions or `async_await`-based utilities.

 #![allow(missing_docs)]

 use {
     futures::{
         future::{Future, FutureExt},
         io::{self, AsyncRead, AsyncWrite},
         ready,
         stream::{Stream, StreamExt},
         task::{Context, Poll},
     },
     pin_utils::unsafe_pinned,
     std::{marker::Unpin, mem, pin::Pin},
 };

 pub trait TempStreamExt: Stream + Sized {
     fn first_elem(self) -> FirstElem<Self> {
         FirstElem { stream: self }
     }
     fn try_into_future<T, E>(self) -> TryIntoFuture<Self>
     where
         Self: Stream<Item = Result<T, E>> + Unpin,
     {
         TryIntoFuture { stream: Some(self) }
     }
 }

 impl<T: Stream + Sized> TempStreamExt for T {}

 pub struct FirstElem<St> {
     stream: St,
 }

 impl<St> FirstElem<St> {
     // Safety: `FirstElem` is `Unpin` iff `St` is `Unpin`.
     unsafe_pinned!(stream: St);
 }

 impl<St: Stream> Future for FirstElem<St> {
     type Output = Option<St::Item>;

     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         self.stream().poll_next(cx)
     }
 }

 pub struct TryIntoFuture<St> {
     stream: Option<St>,
 }

 impl<St> Unpin for TryIntoFuture<St> {}

 impl<T, E, St: Stream<Item = Result<T, E>> + Unpin> Future for TryIntoFuture<St> {
     type Output = Result<(Option<T>, St), E>;

     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let res = ready!(self.stream.as_mut().unwrap().poll_next_unpin(cx));
         Poll::Ready(match res {
             Some(Ok(elem)) => Ok((Some(elem), self.stream.take().unwrap())),
             None => Ok((None, self.stream.take().unwrap())),
             Some(Err(e)) => Err(e),
         })
     }
 }

 pub trait TempAsyncWriteExt: AsyncWrite + Sized {
     fn write_all<T: AsRef<[u8]>>(self, buf: T) -> WriteAll<Self, T> {
         write_all(self, buf)
     }
 }

 impl<T: AsyncWrite + Sized> TempAsyncWriteExt for T {}

 pub struct WriteAll<A, T> {
     state: WriteState<A, T>,
 }

 impl<A, T> Unpin for WriteAll<A, T> {}

 enum WriteState<A, T> {
     Writing { a: A, buf: T, pos: usize },
     Empty,
 }

 pub fn write_all<A, T>(a: A, buf: T) -> WriteAll<A, T>
 where
     A: AsyncWrite,
     T: AsRef<[u8]>,
 {
     WriteAll { state: WriteState::Writing { a: a, buf: buf, pos: 0 } }
 }

 fn zero_write() -> io::Error {
     io::Error::new(io::ErrorKind::WriteZero, "zero-length write")
 }

 impl<A, T> Future for WriteAll<A, T>
 where
     A: AsyncWrite + Unpin,
     T: AsRef<[u8]>,
 {
     type Output = io::Result<(A, T)>;

     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         match &mut self.state {
             WriteState::Writing { a, buf, pos } => {
                 let buf = buf.as_ref();
                 while *pos < buf.len() {
                     let n = ready!(Pin::new(&mut *a).poll_write(cx, &buf[*pos..]))?;
                     *pos += n;
                     if n == 0 {
                         return Poll::Ready(Err(zero_write()));
                     }
                 }
             }
             WriteState::Empty => panic!("poll a WriteAll after it's done"),
         }

         match mem::replace(&mut self.state, WriteState::Empty) {
             WriteState::Writing { a, buf, .. } => Poll::Ready(Ok((a, buf).into())),
             WriteState::Empty => panic!(),
         }
     }
 }

 pub trait TempAsyncReadExt: AsyncRead + Sized {
     fn read_to_end(self, buf: Vec<u8>) -> ReadToEnd<Self> {
         read_to_end(self, buf)
     }
 }

 impl<T: AsyncRead + Sized> TempAsyncReadExt for T {}

 pub struct ReadToEnd<A> {
     state: State<A>,
 }

 impl<A> Unpin for ReadToEnd<A> {}

 enum State<A> {
     Reading { a: A, buf: Vec<u8> },
     Empty,
 }

 pub fn read_to_end<A>(a: A, buf: Vec<u8>) -> ReadToEnd<A>
 where
     A: AsyncRead,
 {
     ReadToEnd { state: State::Reading { a, buf } }
 }

 struct Guard<'a> {
     buf: &'a mut Vec<u8>,
     len: usize,
 }

 impl<'a> Drop for Guard<'a> {
     fn drop(&mut self) {
         unsafe {
             self.buf.set_len(self.len);
         }
     }
 }

 // This uses an adaptive system to extend the vector when it fills. We want to
 // avoid paying to allocate and zero a huge chunk of memory if the reader only
 // has 4 bytes while still making large reads if the reader does have a ton
 // of data to return. Simply tacking on an extra DEFAULT_BUF_SIZE space every
 // time is 4,500 times (!) slower than this if the reader has a very small
 // amount of data to return.
 //
 // Because we're extending the buffer with uninitialized data for trusted
 // readers, we need to make sure to truncate that if any of this panics.
 fn read_to_end_internal<R: AsyncRead + Unpin>(
     r: &mut R,
     cx: &mut Context<'_>,
     buf: &mut Vec<u8>,
 ) -> Poll<io::Result<usize>> {
     let start_len = buf.len();
     let mut g = Guard { len: buf.len(), buf: buf };
     let ret;
     loop {
         if g.len == g.buf.len() {
             g.buf.reserve(32);
             let capacity = g.buf.capacity();
             // FIXME: don't zero when a sound `std::mem::freeze` or similar exists
             g.buf.resize(capacity, 0);
         }

         match ready!(Pin::new(&mut *r).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
 }

 impl<A> Future for ReadToEnd<A>
 where
     A: AsyncRead + Unpin,
 {
     type Output = io::Result<(A, Vec<u8>)>;

     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let this = &mut *self;
         match this.state {
             State::Reading { ref mut a, ref mut buf } => {
                 // If we get `Ok`, then we know the stream hit EOF and we're done. If we
                 // hit "would block" then all the read data so far is in our buffer, and
                 // otherwise we propagate errors
                 ready!(read_to_end_internal(a, cx, buf))?;
             }
             State::Empty => panic!("poll ReadToEnd after it's done"),
         }

         match mem::replace(&mut this.state, State::Empty) {
             State::Reading { a, buf } => Poll::Ready(Ok((a, buf).into())),
             State::Empty => unreachable!(),
         }
     }
 }

 pub trait TempFutureExt: Future + Sized {
     fn select<B>(self, b: B) -> Select<Self, B> {
         Select { a: self, b }
     }

     fn select_unpin<B>(self, b: B) -> SelectUnpin<Self, B>
     where
         Self: Unpin,
         B: Unpin,
     {
         SelectUnpin { a: Some(self), b: Some(b) }
     }
 }

 impl<T: Future + Sized> TempFutureExt for T {}

 pub enum Either<A, B> {
     Left(A),
     Right(B),
 }

 impl<A, B> Either<A, B> {
     pub fn either<T>(self, lf: impl FnOnce(A) -> T, rf: impl FnOnce(B) -> T) -> T {
         match self {
             Either::Left(a) => (lf)(a),
             Either::Right(b) => (rf)(b),
         }
     }
 }

 impl<A: Future, B: Future<Output = A::Output>> Future for Either<A, B> {
     type Output = A::Output;
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         unsafe {
             // Safety: neither child future is ever moved
             match Pin::get_unchecked_mut(self) {
                 Either::Left(a) => Pin::new_unchecked(a).poll(cx),
                 Either::Right(b) => Pin::new_unchecked(b).poll(cx),
             }
         }
     }
 }

 pub struct Select<A, B> {
     a: A,
     b: B,
 }

 impl<A, B> Select<A, B> {
     unsafe_pinned!(a: A);
     unsafe_pinned!(b: B);
 }

 impl<A: Future, B: Future> Future for Select<A, B> {
     type Output = Either<A::Output, B::Output>;
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         if let Poll::Ready(a) = self.as_mut().a().poll(cx) {
             return Poll::Ready(Either::Left(a));
         }
         if let Poll::Ready(b) = self.as_mut().b().poll(cx) {
             return Poll::Ready(Either::Right(b));
         }
         Poll::Pending
     }
 }

 pub struct SelectUnpin<A, B> {
     a: Option<A>,
     b: Option<B>,
 }

 impl<A: Future + Unpin, B: Future + Unpin> Future for SelectUnpin<A, B> {
     type Output = Either<(A::Output, B), (A, B::Output)>;
     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let this = &mut *self;
         if let Poll::Ready(a) = this.a.as_mut().unwrap().poll_unpin(cx) {
             return Poll::Ready(Either::Left((a, this.b.take().unwrap())));
         }
         if let Poll::Ready(b) = this.b.as_mut().unwrap().poll_unpin(cx) {
             return Poll::Ready(Either::Right((this.a.take().unwrap(), b)));
         }
         Poll::Pending
     }
 }
	// Copyright 2018 The Fuchsia Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	//! Temporary Futures extensions used during the transition from 0.2 -> 0.3.
	//! These SHOULD NOT be used for new code. Instead, use the corresponding 0.3
	//! functions or `async_await`-based utilities.

	#![allow(missing_docs)]

	use {
	futures::{
	future::{Future, FutureExt},
	io::{self, AsyncRead, AsyncWrite},
	ready,
	stream::{Stream, StreamExt},
	task::{Context, Poll},
	},
	pin_utils::unsafe_pinned,
	std::{marker::Unpin, mem, pin::Pin},
	};

	pub trait TempStreamExt: Stream + Sized {
	fn first_elem(self) -> FirstElem<Self> {
	FirstElem { stream: self }
	}
	fn try_into_future<T, E>(self) -> TryIntoFuture<Self>
	where
	Self: Stream<Item = Result<T, E>> + Unpin,
	{
	TryIntoFuture { stream: Some(self) }
	}
	}

	impl<T: Stream + Sized> TempStreamExt for T {}

	pub struct FirstElem<St> {
	stream: St,
	}

	impl<St> FirstElem<St> {
	// Safety: `FirstElem` is `Unpin` iff `St` is `Unpin`.
	unsafe_pinned!(stream: St);
	}

	impl<St: Stream> Future for FirstElem<St> {
	type Output = Option<St::Item>;

	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
	self.stream().poll_next(cx)
	}
	}

	pub struct TryIntoFuture<St> {
	stream: Option<St>,
	}

	impl<St> Unpin for TryIntoFuture<St> {}

	impl<T, E, St: Stream<Item = Result<T, E>> + Unpin> Future for TryIntoFuture<St> {
	type Output = Result<(Option<T>, St), E>;

	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
	let res = ready!(self.stream.as_mut().unwrap().poll_next_unpin(cx));
	Poll::Ready(match res {
	Some(Ok(elem)) => Ok((Some(elem), self.stream.take().unwrap())),
	None => Ok((None, self.stream.take().unwrap())),
	Some(Err(e)) => Err(e),
	})
	}
	}

	pub trait TempAsyncWriteExt: AsyncWrite + Sized {
	fn write_all<T: AsRef<[u8]>>(self, buf: T) -> WriteAll<Self, T> {
	write_all(self, buf)
	}
	}

	impl<T: AsyncWrite + Sized> TempAsyncWriteExt for T {}

	pub struct WriteAll<A, T> {
	state: WriteState<A, T>,
	}

	impl<A, T> Unpin for WriteAll<A, T> {}

	enum WriteState<A, T> {
	Writing { a: A, buf: T, pos: usize },
	Empty,
	}

	pub fn write_all<A, T>(a: A, buf: T) -> WriteAll<A, T>
	where
	A: AsyncWrite,
	T: AsRef<[u8]>,
	{
	WriteAll { state: WriteState::Writing { a: a, buf: buf, pos: 0 } }
	}

	fn zero_write() -> io::Error {
	io::Error::new(io::ErrorKind::WriteZero, "zero-length write")
	}

	impl<A, T> Future for WriteAll<A, T>
	where
	A: AsyncWrite + Unpin,
	T: AsRef<[u8]>,
	{
	type Output = io::Result<(A, T)>;

	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
	match &mut self.state {
	WriteState::Writing { a, buf, pos } => {
	let buf = buf.as_ref();
	while *pos < buf.len() {
	let n = ready!(Pin::new(&mut a).poll_write(cx, &buf[pos..]))?;
	*pos += n;
	if n == 0 {
	return Poll::Ready(Err(zero_write()));
	}
	}
	}
	WriteState::Empty => panic!("poll a WriteAll after it's done"),
	}

	match mem::replace(&mut self.state, WriteState::Empty) {
	WriteState::Writing { a, buf, .. } => Poll::Ready(Ok((a, buf).into())),
	WriteState::Empty => panic!(),
	}
	}
	}

	pub trait TempAsyncReadExt: AsyncRead + Sized {
	fn read_to_end(self, buf: Vec<u8>) -> ReadToEnd<Self> {
	read_to_end(self, buf)
	}
	}

	impl<T: AsyncRead + Sized> TempAsyncReadExt for T {}

	pub struct ReadToEnd<A> {
	state: State<A>,
	}

	impl<A> Unpin for ReadToEnd<A> {}

	enum State<A> {
	Reading { a: A, buf: Vec<u8> },
	Empty,
	}

	pub fn read_to_end<A>(a: A, buf: Vec<u8>) -> ReadToEnd<A>
	where
	A: AsyncRead,
	{
	ReadToEnd { state: State::Reading { a, buf } }
	}

	struct Guard<'a> {
	buf: &'a mut Vec<u8>,
	len: usize,
	}

	impl<'a> Drop for Guard<'a> {
	fn drop(&mut self) {
	unsafe {
	self.buf.set_len(self.len);
	}
	}
	}

	// This uses an adaptive system to extend the vector when it fills. We want to
	// avoid paying to allocate and zero a huge chunk of memory if the reader only
	// has 4 bytes while still making large reads if the reader does have a ton
	// of data to return. Simply tacking on an extra DEFAULT_BUF_SIZE space every
	// time is 4,500 times (!) slower than this if the reader has a very small
	// amount of data to return.
	//
	// Because we're extending the buffer with uninitialized data for trusted
	// readers, we need to make sure to truncate that if any of this panics.
	fn read_to_end_internal<R: AsyncRead + Unpin>(
	r: &mut R,
	cx: &mut Context<'_>,
	buf: &mut Vec<u8>,
	) -> Poll<io::Result<usize>> {
	let start_len = buf.len();
	let mut g = Guard { len: buf.len(), buf: buf };
	let ret;
	loop {
	if g.len == g.buf.len() {
	g.buf.reserve(32);
	let capacity = g.buf.capacity();
	// FIXME: don't zero when a sound `std::mem::freeze` or similar exists
	g.buf.resize(capacity, 0);
	}

	match ready!(Pin::new(&mut *r).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
	}

	impl<A> Future for ReadToEnd<A>
	where
	A: AsyncRead + Unpin,
	{
	type Output = io::Result<(A, Vec<u8>)>;

	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
	let this = &mut *self;
	match this.state {
	State::Reading { ref mut a, ref mut buf } => {
	// If we get `Ok`, then we know the stream hit EOF and we're done. If we
	// hit "would block" then all the read data so far is in our buffer, and
	// otherwise we propagate errors
	ready!(read_to_end_internal(a, cx, buf))?;
	}
	State::Empty => panic!("poll ReadToEnd after it's done"),
	}

	match mem::replace(&mut this.state, State::Empty) {
	State::Reading { a, buf } => Poll::Ready(Ok((a, buf).into())),
	State::Empty => unreachable!(),
	}
	}
	}

	pub trait TempFutureExt: Future + Sized {
	fn select<B>(self, b: B) -> Select<Self, B> {
	Select { a: self, b }
	}

	fn select_unpin<B>(self, b: B) -> SelectUnpin<Self, B>
	where
	Self: Unpin,
	B: Unpin,
	{
	SelectUnpin { a: Some(self), b: Some(b) }
	}
	}

	impl<T: Future + Sized> TempFutureExt for T {}

	pub enum Either<A, B> {
	Left(A),
	Right(B),
	}

	impl<A, B> Either<A, B> {
	pub fn either<T>(self, lf: impl FnOnce(A) -> T, rf: impl FnOnce(B) -> T) -> T {
	match self {
	Either::Left(a) => (lf)(a),
	Either::Right(b) => (rf)(b),
	}
	}
	}

	impl<A: Future, B: Future<Output = A::Output>> Future for Either<A, B> {
	type Output = A::Output;
	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
	unsafe {
	// Safety: neither child future is ever moved
	match Pin::get_unchecked_mut(self) {
	Either::Left(a) => Pin::new_unchecked(a).poll(cx),
	Either::Right(b) => Pin::new_unchecked(b).poll(cx),
	}
	}
	}
	}

	pub struct Select<A, B> {
	a: A,
	b: B,
	}

	impl<A, B> Select<A, B> {
	unsafe_pinned!(a: A);
	unsafe_pinned!(b: B);
	}

	impl<A: Future, B: Future> Future for Select<A, B> {
	type Output = Either<A::Output, B::Output>;
	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
	if let Poll::Ready(a) = self.as_mut().a().poll(cx) {
	return Poll::Ready(Either::Left(a));
	}
	if let Poll::Ready(b) = self.as_mut().b().poll(cx) {
	return Poll::Ready(Either::Right(b));
	}
	Poll::Pending
	}
	}

	pub struct SelectUnpin<A, B> {
	a: Option<A>,
	b: Option<B>,
	}

	impl<A: Future + Unpin, B: Future + Unpin> Future for SelectUnpin<A, B> {
	type Output = Either<(A::Output, B), (A, B::Output)>;
	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
	let this = &mut *self;
	if let Poll::Ready(a) = this.a.as_mut().unwrap().poll_unpin(cx) {
	return Poll::Ready(Either::Left((a, this.b.take().unwrap())));
	}
	if let Poll::Ready(b) = this.b.as_mut().unwrap().poll_unpin(cx) {
	return Poll::Ready(Either::Right((this.a.take().unwrap(), b)));
	}
	Poll::Pending
	}
	}