1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
//! A minimal async executor that lets you block on a future
//!
//! # Example
//!
//! ```ignore
//! let result = pollster::block_on(my_future);
//! ```

use std::{
    mem::forget,
    future::Future,
    pin::Pin,
    task::{Poll, Context, Waker, RawWaker, RawWakerVTable},
    sync::{Condvar, Mutex, Arc},
};

struct Signal {
    lock: Mutex<bool>,
    cond: Condvar,
}

impl Signal {
    fn wait(&self) {
        let mut wakeup = self.lock.lock().unwrap();
        if !*wakeup {
            // Signal was notified since the last wakeup, don't wait
            wakeup = self.cond.wait(wakeup).unwrap();
        }
        *wakeup = false;
    }

    fn notify(&self) {
        let mut wakeup = self.lock.lock().unwrap();
        *wakeup = true;
        self.cond.notify_one();
    }
}

static VTABLE: RawWakerVTable = unsafe { RawWakerVTable::new(
    |signal| {
        let arc = Arc::from_raw(signal);
        let waker = RawWaker::new(Arc::into_raw(arc.clone()) as *const _, &VTABLE);
        forget(arc);
        waker
    },
    // Notify by dropping the Arc (wake)
    |signal| Arc::from_raw(signal as *const Signal).notify(),
    // Notify without dropping the Arc (wake_by_ref)
    |signal| (&*(signal as *const Signal)).notify(),
    // Drop the Arc
    |signal| drop(Arc::from_raw(signal as *const Signal)),
) };

/// Block until the the future is ready.
pub fn block_on<F: Future>(mut fut: F) -> F::Output {
    let signal = Arc::new(Signal {
        lock: Mutex::new(false),
        cond: Condvar::new(),
    });

    let waker = unsafe { Waker::from_raw(RawWaker::new(Arc::into_raw(signal.clone()) as *const _, &VTABLE)) };

    loop {
        match unsafe { Pin::new_unchecked(&mut fut).poll(&mut Context::from_waker(&waker)) } {
            Poll::Pending => signal.wait(),
            Poll::Ready(item) => break item,
        }
    }
}