src/future/try_future/try_flatten.rs

use core::pin::Pin;
use futures_core::future::{FusedFuture, Future, TryFuture};
use futures_core::stream::{FusedStream, Stream, TryStream};
#[cfg(feature = "sink")]
use futures_sink::Sink;
use futures_core::task::{Context, Poll};
use pin_project::pin_project;

#[pin_project(project = TryFlattenProj)]
#[derive(Debug)]
pub enum TryFlatten<Fut1, Fut2> {
    First(#[pin] Fut1),
    Second(#[pin] Fut2),
    Empty,
}

impl<Fut1, Fut2> TryFlatten<Fut1, Fut2> {
    pub(crate) fn new(future: Fut1) -> Self {
        TryFlatten::First(future)
    }
}

impl<Fut> FusedFuture for TryFlatten<Fut, Fut::Ok>
    where Fut: TryFuture,
          Fut::Ok: TryFuture<Error=Fut::Error>,
{
    fn is_terminated(&self) -> bool {
        match self {
            TryFlatten::Empty => true,
            _ => false,
        }
    }
}

impl<Fut> Future for TryFlatten<Fut, Fut::Ok>
    where Fut: TryFuture,
          Fut::Ok: TryFuture<Error=Fut::Error>,
{
    type Output = Result<<Fut::Ok as TryFuture>::Ok, Fut::Error>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        Poll::Ready(loop {
            match self.as_mut().project() {
                TryFlattenProj::First(f) => {
                    match ready!(f.try_poll(cx)) {
                        Ok(f) => self.set(TryFlatten::Second(f)),
                        Err(e) => {
                            self.set(TryFlatten::Empty);
                            break Err(e);
                        }
                    }
                },
                TryFlattenProj::Second(f) => {
                    let output = ready!(f.try_poll(cx));
                    self.set(TryFlatten::Empty);
                    break output;
                },
                TryFlattenProj::Empty => panic!("TryFlatten polled after completion"),
            }
        })
    }
}

impl<Fut> FusedStream for TryFlatten<Fut, Fut::Ok>
    where Fut: TryFuture,
          Fut::Ok: TryStream<Error=Fut::Error>,
{
    fn is_terminated(&self) -> bool {
        match self {
            TryFlatten::Empty => true,
            _ => false,
        }
    }
}

impl<Fut> Stream for TryFlatten<Fut, Fut::Ok>
    where Fut: TryFuture,
          Fut::Ok: TryStream<Error=Fut::Error>,
{
    type Item = Result<<Fut::Ok as TryStream>::Ok, Fut::Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        Poll::Ready(loop {
            match self.as_mut().project() {
                TryFlattenProj::First(f) => {
                    match ready!(f.try_poll(cx)) {
                        Ok(f) => self.set(TryFlatten::Second(f)),
                        Err(e) => {
                            self.set(TryFlatten::Empty);
                            break Some(Err(e));
                        }
                    }
                },
                TryFlattenProj::Second(f) => {
                    let output = ready!(f.try_poll_next(cx));
                    if output.is_none() {
                        self.set(TryFlatten::Empty);
                    }
                    break output;
                },
                TryFlattenProj::Empty => break None,
            }
        })
    }
}


#[cfg(feature = "sink")]
impl<Fut, Item> Sink<Item> for TryFlatten<Fut, Fut::Ok>
where
    Fut: TryFuture,
    Fut::Ok: Sink<Item, Error=Fut::Error>,
{
    type Error = Fut::Error;

    fn poll_ready(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(loop {
            match self.as_mut().project() {
                TryFlattenProj::First(f) => {
                    match ready!(f.try_poll(cx)) {
                        Ok(f) => self.set(TryFlatten::Second(f)),
                        Err(e) => {
                            self.set(TryFlatten::Empty);
                            break Err(e);
                        }
                    }
                },
                TryFlattenProj::Second(f) => {
                    break ready!(f.poll_ready(cx));
                },
                TryFlattenProj::Empty => panic!("poll_ready called after eof"),
            }
        })
    }

    fn start_send(self: Pin<&mut Self>, item: Item) -> Result<(), Self::Error> {
        match self.project() {
            TryFlattenProj::First(_) => panic!("poll_ready not called first"),
            TryFlattenProj::Second(f) => f.start_send(item),
            TryFlattenProj::Empty => panic!("start_send called after eof"),
        }
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        match self.project() {
            TryFlattenProj::First(_) => Poll::Ready(Ok(())),
            TryFlattenProj::Second(f) => f.poll_flush(cx),
            TryFlattenProj::Empty => panic!("poll_flush called after eof"),
        }
    }

    fn poll_close(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        let res = match self.as_mut().project() {
            TryFlattenProj::Second(f) => f.poll_close(cx),
            _ => Poll::Ready(Ok(())),
        };
        if res.is_ready() {
            self.set(TryFlatten::Empty);
        }
        res
    }
}