future.rs

use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll};

pub struct ReturnFuture<T>(Option<T>);

impl<T> ReturnFuture<T> {
    pub fn new(t: T) -> Self {
        Self(Some(t))
    }
}

// Unpinを手で実装する
impl<T> Unpin for ReturnFuture<T> {}
// 自動で導出される場合は以下のようにパラメータにも `Unpin` が要求される。今回はそれは不要だった。
// impl<T: Unpin> Unpin for ReturnFuture<T> {}

impl<T> Future for ReturnFuture<T> {
    type Output = T;
    fn poll(self: Pin<&mut Self>, _: &mut Context) -> Poll<Self::Output> {
        Poll::Ready(
            self.get_mut()
                .0
                .take()
                .expect("A future should never be polled after it returns Ready"),
        )
    }
}

use std::sync::mpsc::{channel, Receiver};
use std::thread;

pub struct ThreadFuture<T> {
    rx: Receiver<T>,
    waker: Arc<Mutex<Option<Waker>>>,
}

impl<T> ThreadFuture<T>
where
    T: Send + 'static,
{
    pub fn start<F>(f: F) -> Self
    where
        F: FnOnce() -> T + Send + 'static,
    {
        let (tx, rx) = channel();
        let waker = Arc::new(Mutex::new(None::<Waker>));
        let w = waker.clone();
        thread::spawn(move || {
            tx.send(f()).unwrap();
            if let Some(waker) = &*w.lock().unwrap() {
                waker.wake_by_ref()
            }
        });

        Self { rx, waker }
    }
}

impl<T> Future for ThreadFuture<T>
where
    T: Send + Sync + 'static,
{
    type Output = T;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
        *self.waker.lock().unwrap() = Some(cx.waker().clone());
        match self.rx.try_recv() {
            Ok(t) => Poll::Ready(t),
            Err(_) => Poll::Pending,
        }
    }
}

use std::task::{RawWaker, RawWakerVTable, Waker};

#[derive(Debug, Clone)]
struct SpinWaker;

static SPIN_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
    SpinWaker::unsafe_clone,
    SpinWaker::unsafe_wake,
    SpinWaker::unsafe_wake_by_ref,
    SpinWaker::unsafe_drop,
);

impl SpinWaker {
    fn waker() -> Waker {
        // this is safe because waker's data and vtable is consistent
        unsafe { Waker::from_raw(Self::new().into_raw_waker()) }
    }

    fn new() -> Self {
        Self
    }

    unsafe fn into_raw_waker(self) -> RawWaker {
        let ptr = Box::into_raw(Box::new(self)) as *const ();
        RawWaker::new(ptr, &SPIN_WAKER_VTABLE)
    }

    unsafe fn unsafe_clone(this: *const ()) -> RawWaker {
        let ptr = this as *const Self;
        Box::new(ptr.as_ref().unwrap().clone()).into_raw_waker()
    }
    fn wake(self: Self) {}
    unsafe fn unsafe_wake(this: *const ()) {
        let ptr = this as *mut Self;
        Box::from_raw(ptr).wake()
    }

    fn wake_by_ref(&self) {
        Box::new(self.clone()).wake()
    }

    unsafe fn unsafe_wake_by_ref(this: *const ()) {
        let ptr = this as *const Self;
        ptr.as_ref().unwrap().wake_by_ref()
    }

    unsafe fn unsafe_drop(this: *const ()) {
        let ptr = this as *mut Self;
        Box::from_raw(ptr);
    }
}

pub struct SpinRunner;

use std::ops::{Deref, DerefMut};

impl SpinRunner {
    pub fn new() -> Self {
        Self
    }

    pub fn run<F>(&mut self, future: F) -> F::Output
    where
        F: Future,
    {
        let mut future = Pin::from(Box::new(future));
        self.run_pin(future.as_mut())
    }

    pub fn run_pin<F>(&mut self, mut future: Pin<F>) -> <<F as Deref>::Target as Future>::Output
    where
        F: DerefMut,
        <F as Deref>::Target: Future,
    {
        let waker = SpinWaker::waker();
        let mut cx = Context::from_waker(&waker);
        // loopでpollするだけ
        let mut i = 0;
        loop {
            i += 1;
            match future.as_mut().poll(&mut cx) {
                Poll::Ready(ret) => {
                    println!("{}", i);
                    return ret;
                }
                Poll::Pending => continue,
            }
        }
    }
}

use std::sync::Condvar;

#[derive(Debug)]
struct CondWaker(Arc<(Mutex<()>, Condvar)>);

static COND_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
    CondWaker::unsafe_clone,
    CondWaker::unsafe_wake,
    CondWaker::unsafe_wake_by_ref,
    CondWaker::unsafe_drop,
);

impl CondWaker {
    fn waker(inner: Arc<(Mutex<()>, Condvar)>) -> Waker {
        // this is safe because waker's data and vtable is consistent
        unsafe { Waker::from_raw(Box::new(Self::new(inner)).into_raw_waker()) }
    }

    fn new(inner: Arc<(Mutex<()>, Condvar)>) -> Self {
        Self(inner)
    }

    unsafe fn into_raw_waker(self) -> RawWaker {
        let ptr = Arc::into_raw(self.0) as *const ();
        RawWaker::new(ptr, &COND_WAKER_VTABLE)
    }

    unsafe fn unsafe_clone(this: *const ()) -> RawWaker {
        let ptr = this as *const (Mutex<()>, Condvar);
        let arc = Arc::from_raw(ptr);
        let ret = Self::new(arc.clone()).into_raw_waker();
        std::mem::forget(arc);
        ret
    }
    fn wake(self) {
        (self.0).1.notify_all()
    }
    unsafe fn unsafe_wake(this: *const ()) {
        let ptr = this as *const (Mutex<()>, Condvar);
        Self::new(Arc::from_raw(ptr)).wake()
    }

    unsafe fn unsafe_wake_by_ref(this: *const ()) {
        let ptr = this as *const (Mutex<()>, Condvar);
        let arc = Arc::from_raw(ptr);
        let ret = Self::new(arc.clone());
        std::mem::forget(arc);
        ret.wake()
    }
    unsafe fn unsafe_drop(this: *const ()) {
        let ptr = this as *const (Mutex<()>, Condvar);
        Self::new(Arc::from_raw(ptr));
    }
}

struct CondRunner(Arc<(Mutex<()>, Condvar)>);

impl CondRunner {
    fn new() -> Self {
        Self(Arc::new((Mutex::new(()), Condvar::new())))
    }

    pub fn run<F>(&mut self, future: F) -> F::Output
    where
        F: Future,
    {
        let mut future = Pin::from(Box::new(future));
        self.run_pin(future.as_mut())
    }

    pub fn run_pin<F>(&mut self, mut future: Pin<F>) -> <<F as Deref>::Target as Future>::Output
    where
        F: DerefMut,
        <F as Deref>::Target: Future,
    {
        let waker = CondWaker::waker(self.0.clone());
        let mut cx = Context::from_waker(&waker);
        let mut i = 0;
        loop {
            i += 1;
            match future.as_mut().poll(&mut cx) {
                Poll::Ready(ret) => {
                    println!("{}", i);
                    return ret;
                }
                Poll::Pending => {
                    // 一回pollしてまだだったら条件変数の通知を待つ
                    let lock = (self.0).0.lock().unwrap();
                    *(self.0).1.wait(lock).unwrap();
                    // 条件が通知されたら再度pollする
                }
            };
        }
    }
}

fn fib(n: u64) -> u64 {
    if n < 2 {
        1
    } else {
        fib(n - 1) + fib(n - 2)
    }
}

fn main() {
    let mut runner = CondRunner::new();
    let future = ThreadFuture::start(|| fib(42));
    println!("cond runner");
    let ret = runner.run(future);
    println!("answer is {}", ret);

    let mut runner = SpinRunner::new();
    let future = ThreadFuture::start(|| fib(42));
    println!("spin runner");
    let ret = runner.run(future);
    println!("answer is {}", ret);
}