Alternative `and_then` combinator

and_then2.rs

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

extern crate futures;

use futures::{Async, Future, IntoFuture, Poll};
use std::mem;

pub enum AndThen2<Original, Resulting, F> {
    First(F, Original),
    Second(Resulting),
    EmptyState,
}

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

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        loop {
            let previous_state = mem::replace(self, AndThen2::EmptyState);
            let new_state = match previous_state {
                AndThen2::First(f, mut original) => match original.poll()? {
                    Async::Ready(x) => AndThen2::Second(f(x).into_future()),
                    Async::NotReady => {
                        mem::replace(self, AndThen2::First(f, original));
                        return Ok(Async::NotReady);
                    }
                },
                AndThen2::Second(mut resulting) => match resulting.poll() {
                    Err(e) => return Ok(Async::Ready(Err(e))),
                    Ok(Async::Ready(res)) => return Ok(Async::Ready(Ok(res))),
                    Ok(Async::NotReady) => {
                        mem::replace(self, AndThen2::Second(resulting));
                        return Ok(Async::NotReady);
                    }
                },
                AndThen2::EmptyState => panic!("poll() called after completion"),
            };
            mem::replace(self, new_state);
        }
    }
}

pub trait FutureExt: Future {
    fn and_then2<F, B>(self, f: F) -> AndThen2<Self, B::Future, F>
    where
        B: IntoFuture,
        F: FnOnce(Self::Item) -> B,
        Self: Sized,
    {
        AndThen2::First(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());
    }
}