and-then2-with-state_machine_future.rs

//! An alternative to a traditional `Future::and_then` combinator.

#[macro_use]
extern crate futures;

#[macro_use]
extern crate state_machine_future;

use futures::{Async, Future, IntoFuture, Poll};
use state_machine_future::RentToOwn;

#[derive(StateMachineFuture)]
pub enum StateMachine<Original, Resulting, ResultingFuture, Item, FirstError, SecondError, F>
where
    Original: Future<Error = FirstError>,
    Resulting: IntoFuture<Future = ResultingFuture, Item = Item, Error = SecondError>,
    ResultingFuture: Future<Item = Item, Error = SecondError>,
    F: FnOnce(Original::Item) -> Resulting,
{
    #[state_machine_future(start, transitions(Second))]
    First(F, Original),
    #[state_machine_future(transitions(Ready))]
    Second(ResultingFuture),
    #[state_machine_future(ready)]
    Ready(Result<Item, SecondError>),
    #[state_machine_future(error)]
    Error(FirstError),
}

impl<Original, Resulting, ResultingFuture, Item, FirstError, SecondError, F>
    PollStateMachine<Original, Resulting, ResultingFuture, Item, FirstError, SecondError, F>
    for StateMachine<Original, Resulting, ResultingFuture, Item, FirstError, SecondError, F>
where
    Original: Future<Error = FirstError>,
    Resulting: IntoFuture<Future = ResultingFuture, Item = Item, Error = SecondError>,
    ResultingFuture: Future<Item = Item, Error = SecondError>,
    F: FnOnce(Original::Item) -> Resulting,
{
    fn poll_first<'a>(
        data: &'a mut RentToOwn<
            'a,
            First<Original, Resulting, ResultingFuture, Item, FirstError, SecondError, F>,
        >,
    ) -> Poll<AfterFirst<ResultingFuture, Item, SecondError>, FirstError> {
        let x = try_ready!(data.1.poll());
        transition!(Second((data.take().0)(x).into_future()))
    }

    fn poll_second<'a>(
        state: &'a mut RentToOwn<'a, Second<ResultingFuture, Item, SecondError>>,
    ) -> Poll<AfterSecond<Item, SecondError>, FirstError> {
        match state.0.poll() {
            Ok(Async::Ready(x)) => transition!(Ready(Ok(x))),
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Err(e) => transition!(Ready(Err(e))),
        }
    }
}

pub struct AndThen2<Original, Resulting, F>(
    StateMachineFuture<
        Original,
        Resulting,
        Resulting::Future,
        Resulting::Item,
        Original::Error,
        Resulting::Error,
        F,
    >,
)
where
    Original: Future,
    Resulting: IntoFuture,
    F: FnOnce(Original::Item) -> Resulting;

impl<Original, Resulting, F> Future for AndThen2<Original, Resulting, F>
where
    Original: Future,
    Resulting: IntoFuture,
    F: FnOnce(Original::Item) -> Resulting,
{
    type Item = Result<Resulting::Item, Resulting::Error>;
    type Error = Original::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        self.0.poll()
    }
}

pub trait FutureExt: Future {
    fn and_then2<F, B>(self, f: F) -> AndThen2<Self, B, F>
    where
        B: IntoFuture,
        F: FnOnce(Self::Item) -> B,
        Self: Sized,
    {
        AndThen2(StateMachine::start(f, self))
    }
}

impl<T: Future> FutureExt for T {}

#[cfg(test)]
mod tests {
    use super::*;

    #[derive(Debug, PartialEq)]
    struct Err1;
    #[derive(Debug, PartialEq)]
    struct Err2;

    fn plus_one_ok(x: usize) -> impl IntoFuture<Item = usize, Error = Err2> {
        futures::future::ok(x + 1)
    }

    fn just_err(_: usize) -> impl IntoFuture<Item = usize, Error = Err2> {
        futures::future::err(Err2)
    }

    #[test]
    fn test_ok_ok() {
        let f1 = futures::future::ok::<_, Err1>(1);
        let combined = f1.and_then2(plus_one_ok);
        assert_eq!(Ok(Ok(2)), combined.wait());
    }

    #[test]
    fn test_ok_err() {
        let f1 = futures::future::ok::<_, Err1>(1);
        let combined = f1.and_then2(just_err);
        assert_eq!(Ok(Err(Err2)), combined.wait());
    }

    #[test]
    fn test_err_ok() {
        let f1 = futures::future::err(Err1);
        let combined = f1.and_then2(plus_one_ok);
        assert_eq!(Err(Err1), combined.wait());
    }

    #[test]
    fn test_err_err() {
        let f1 = futures::future::err(Err1);
        let combined = f1.and_then2(just_err);
        assert_eq!(Err(Err1), combined.wait());
    }
}