dope-runtime 0.2.3

use std::{
    cell::{Cell, RefCell},
    future::Future,
    pin::Pin,
    rc::Rc,
    task::{Context, Poll, Waker},
};

pub struct JoinHandle<T> {
    kind: JoinHandleKind<T>,
}

enum JoinHandleKind<T> {
    Local(Rc<LocalJoinState<T>>),
}

struct LocalJoinState<T> {
    ready: Cell<bool>,
    value: RefCell<Option<T>>,
    waker: RefCell<Option<Waker>>,
}

impl<T: 'static> JoinHandle<T> {
    pub(crate) fn new_local() -> (Self, impl FnOnce(T) + 'static) {
        let state = Rc::new(LocalJoinState {
            ready: Cell::new(false),
            value: RefCell::new(None),
            waker: RefCell::new(None),
        });
        let state2 = state.clone();
        let complete = move |v: T| {
            *state2.value.borrow_mut() = Some(v);
            state2.ready.set(true);
            if let Some(w) = state2.waker.borrow_mut().take() {
                w.wake();
            }
        };

        (
            Self {
                kind: JoinHandleKind::Local(state),
            },
            complete,
        )
    }
}

impl<T: 'static> Future for JoinHandle<T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let me = self.get_mut();
        match &mut me.kind {
            JoinHandleKind::Local(state) => {
                if let Some(v) = state.value.borrow_mut().take() {
                    return Poll::Ready(v);
                }

                let mut slot = state.waker.borrow_mut();
                let replace = match slot.as_ref() {
                    Some(w) => !w.will_wake(cx.waker()),
                    None => true,
                };
                if replace {
                    *slot = Some(cx.waker().clone());
                }
                drop(slot);

                if let Some(v) = state.value.borrow_mut().take() {
                    return Poll::Ready(v);
                }
                Poll::Pending
            }
        }
    }
}