#![feature(futures_api, async_await, await_macro)]

use core::pin::Pin;
use core::future::Future;
use core::task::{Waker, Poll, RawWaker, RawWakerVTable};
use core::ptr::null;

struct A;

impl Future for A {
    type Output = String;

    fn poll(self: Pin<&mut Self>, _waker: &Waker) -> Poll<Self::Output> {
        Poll::Ready("poll".to_string())
    }
}

struct B {
    a: A,
}

impl B {
    fn run<'a>(self: Pin<&'a mut Self>) -> impl Future<Output = String> + 'a {
        unsafe { Pin::new_unchecked(&mut Pin::get_unchecked_mut(self).a) }
    }
}

async fn foo_async(mut b: B) -> String {
    let mut b = unsafe { Pin::new_unchecked(&mut b) };
    await!(b.as_mut().run()) + &await!(b.as_mut().run())
}

fn foo(mut b: B) -> impl Future<Output = String> {
    let mut b = unsafe { Pin::new_unchecked(&mut b) };
    b.as_mut().run().then(move |s1| b.as_mut().run().then(move |s2| ready(s1 + &s2)))
}

fn main() {
    println!("result: {}", wait(foo(B { a: A })));
    println!("result: {}", wait(foo_async(B { a: A })));
}

// ------

fn wait<F>(mut future: F) -> F::Output where F: Future {
    let waker = noop_waker();
    let mut future = unsafe { Pin::new_unchecked(&mut future) };
    loop {
        if let Poll::Ready(output) = future.as_mut().poll(&waker) {
            return output;
        }
    }
}

const NOOP_WAKER_VTABLE: RawWakerVTable = RawWakerVTable {
    clone: noop_clone,
    drop: noop,
    wake: noop,
};

pub fn noop_waker() -> Waker {
    unsafe {
        Waker::new_unchecked(noop_raw_waker())
    }
}

unsafe fn noop_clone(_data: *const()) -> RawWaker {
    noop_raw_waker()
}

unsafe fn noop(_data: *const()) {
}

fn noop_raw_waker() -> RawWaker {
    RawWaker::new(null(), &NOOP_WAKER_VTABLE)
}

#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct Then<Fut1, Fut2, F> {
    chain: Chain<Fut1, Fut2, F>,
}

impl<Fut1, Fut2, F> Then<Fut1, Fut2, F>
    where Fut1: Future,
          Fut2: Future,
{
    fn new(future: Fut1, f: F) -> Then<Fut1, Fut2, F> {
        Then {
            chain: Chain::new(future, f),
        }
    }

    fn chain<'a>(self: Pin<&'a mut Self>) -> Pin<&'a mut Chain<Fut1, Fut2, F>> {
        unsafe { Pin::new_unchecked(&mut Pin::get_unchecked_mut(self).chain) }
    }
}

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

    fn poll(mut self: Pin<&mut Self>, waker: &Waker) -> Poll<Fut2::Output> {
        self.as_mut().chain().poll(waker, |output, f| f(output))
    }
}

trait FutureExt: Future {
    fn then<Fut, F>(self, f: F) -> Then<Self, Fut, F>
        where F: FnOnce(Self::Output) -> Fut,
              Fut: Future,
              Self: Sized,
    {
        Then::new(self, f)
}
}

impl<Fut: Future> FutureExt for Fut {
}
#[must_use = "futures do nothing unless polled"]
#[derive(Debug)]
pub(crate) enum Chain<Fut1, Fut2, Data> {
    First(Fut1, Option<Data>),
    Second(Fut2),
    Empty,
}

impl<Fut1, Fut2, Data> Chain<Fut1, Fut2, Data> {
    pub(crate)fn is_terminated(&self) -> bool {
        if let Chain::Empty = *self { true } else { false }
    }
}

impl<Fut1, Fut2, Data> Chain<Fut1, Fut2, Data>
    where Fut1: Future,
          Fut2: Future,
{
    pub(crate) fn new(fut1: Fut1, data: Data) -> Chain<Fut1, Fut2, Data> {
        Chain::First(fut1, Some(data))
    }

    pub(crate) fn poll<F>(
        self: Pin<&mut Self>,
        waker: &Waker,
        f: F,
    ) -> Poll<Fut2::Output>
        where F: FnOnce(Fut1::Output, Data) -> Fut2,
    {
        let mut f = Some(f);

        // Safe to call `get_unchecked_mut` because we won't move the futures.
        let this = unsafe { Pin::get_unchecked_mut(self) };

        loop {
            let (output, data) = match this {
                Chain::First(fut1, data) => {
                    match unsafe { Pin::new_unchecked(fut1) }.poll(waker) {
                        Poll::Pending => return Poll::Pending,
                        Poll::Ready(output) => (output, data.take().unwrap()),
                    }
                }
                Chain::Second(fut2) => {
                    return unsafe { Pin::new_unchecked(fut2) }.poll(waker);
                }
                Chain::Empty => unreachable!()
            };

            *this = Chain::Empty; // Drop fut1
            let fut2 = (f.take().unwrap())(output, data);
            *this = Chain::Second(fut2)
        }
    }
}

#[derive(Debug, Clone)]
#[must_use = "futures do nothing unless polled"]
pub struct Ready<T>(Option<T>);

impl<T> Unpin for Ready<T> {}

impl<T> Future for Ready<T> {
    type Output = T;

    #[inline]
    fn poll(mut self: Pin<&mut Self>, _waker: &Waker) -> Poll<T> {
        Poll::Ready(self.0.take().unwrap())
    }
}

pub fn ready<T>(t: T) -> Ready<T> {
    Ready(Some(t))
}