future-and-then.rs

use async_std::task::{sleep, spawn};
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use std::time::Duration;

enum AndThen<Fut1, Fut2, F> {
    First(Fut1, F),
    Second(Fut2),
}

impl<Fut1, Fut2, F> Future for AndThen<Fut1, Fut2, F>
where
    Fut1: Future,
    Fut2: Future,
    F: Fn(Fut1::Output) -> Fut2,
{
    type Output = Fut2::Output;

    fn poll(self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Self::Output> {
        let state = unsafe { self.get_unchecked_mut() };
        loop {
            match state {
                AndThen::First(fut1, f) => match unsafe { Pin::new_unchecked(fut1) }.poll(ctx) {
                    Poll::Pending => {
                        return Poll::Pending;
                    }
                    Poll::Ready(x) => {
                        *state = AndThen::Second(f(x));
                    }
                },
                AndThen::Second(fut2) => {
                    return unsafe { Pin::new_unchecked(fut2) }.poll(ctx);
                }
            }
        }
    }
}

trait FutureAndThen: Future {
    fn and_then<F: Fn(Self::Output) -> Fut2, Fut2: Future>(self, f: F) -> AndThen<Self, Fut2, F>
    where
        Self: Sized,
    {
        AndThen::First(self, f)
    }
}

impl<Fut: Future> FutureAndThen for Fut {}

fn sleepus() -> impl Future<Output = ()> {
    println!("Sleepus 1");
    sleep(Duration::from_millis(500)).and_then(|()| {
        println!("Sleepus 2");
        sleep(Duration::from_millis(500)).and_then(|()| {
            println!("Sleepus 3");
            sleep(Duration::from_millis(500)).and_then(|()| {
                println!("Sleepus 4");
                async_std::future::ready(())
            })
        })
    })
}

async fn interruptus() {
    for i in 1..=5 {
        println!("Interruptus {}", i);
        sleep(Duration::from_millis(1000)).await;
    }
}

#[async_std::main]
async fn main() {
    let sleepus = spawn(sleepus());
    interruptus().await;

    sleepus.await;
}