takeaway 0.1.0

//! Executing async code.

use alloc::sync::Arc;
use core::{
    sync::atomic::{AtomicU32, Ordering},
    task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
};

//----------- block_on() -------------------------------------------------------

/// Run a future to completion, blocking on the current thread.
///
/// This is a simple async runtime; it allows using `async` code in synchronous
/// contexts.  It has no facilities for efficient I/O, timeouts, or other async
/// primitives; use a fully-fledged runtime like [`tokio`] for that.  The only
/// facility `block_on()` provides is to suspend the thread when the future is
/// waiting for some external computation (i.e. another thread) to progress.
///
/// [`tokio`]: https://tokio.rs
///
/// # Implementation
///
/// When the future returns [`Poll::Pending`], it is assumed to be waiting on
/// another thread with a [`Waker`]; the current thread will be put to sleep
/// until the [`Waker`] is activated by that other thread.
///
/// This is a re-implementation of [`futures::executor::block_on()`] with a
/// smaller dependency footprint.  It does not require the Rust standard library
/// or any of the [`futures`] crates; it uses [`alloc`] for allocations and
/// [`atomic_wait`] for OS-specific thread blocking.
///
/// [`futures`]: https://docs.rs/futures
/// [`futures::executor::block_on()`]: https://docs.rs/futures-executor/0.3/futures_executor/fn.block_on.html
pub fn block_on<F: Future>(future: F) -> F::Output {
    let mut future = core::pin::pin!(future);
    let unparker = Unparker::default();
    let waker = unparker.clone().into();
    let mut context = Context::from_waker(&waker);

    loop {
        if let Poll::Ready(result) = future.as_mut().poll(&mut context) {
            break result;
        }

        unparker.park();
    }
}

//----------- Unparker ---------------------------------------------------------

/// A thread unparker.
#[derive(Clone, Debug)]
#[repr(transparent)]
struct Unparker {
    /// The raw value.
    ///
    /// This is either 0 (unblocked) or 1 (blocked).
    raw: Arc<AtomicU32>,
}

impl Default for Unparker {
    fn default() -> Self {
        Self {
            raw: Arc::new(AtomicU32::new(1)),
        }
    }
}

impl Unparker {
    /// Park the thread.
    fn park(&self) {
        let prev = self.raw.swap(1, Ordering::Relaxed);
        if prev == 1 {
            atomic_wait::wait(&self.raw, 1);
        }
    }
}

impl From<Unparker> for Waker {
    fn from(value: Unparker) -> Self {
        let data = Arc::into_raw(value.raw);
        let raw_waker =
            RawWaker::new(data.cast::<()>(), &Unparker::WAKER_VTABLE);
        unsafe { Waker::from_raw(raw_waker) }
    }
}

impl Unparker {
    const WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
        Self::waker_clone,
        Self::waker_wake,
        Self::waker_wake_by_ref,
        Self::waker_drop,
    );

    unsafe fn waker_clone(data: *const ()) -> RawWaker {
        let data = data.cast::<AtomicU32>();
        unsafe { Arc::increment_strong_count(data) };
        RawWaker::new(data.cast::<()>(), &Self::WAKER_VTABLE)
    }

    unsafe fn waker_wake(data: *const ()) {
        let data = data.cast::<AtomicU32>();
        let data = unsafe { Arc::from_raw(data) };

        let value = data.swap(0, Ordering::Relaxed);
        if value != 0 {
            // The thread was asleep; wake it up.
            atomic_wait::wake_one(&*data);
        }
    }

    unsafe fn waker_wake_by_ref(data: *const ()) {
        let data = data.cast::<AtomicU32>();
        let data = unsafe { &*data };

        let value = data.swap(0, Ordering::Relaxed);
        if value != 0 {
            // The thread was asleep; wake it up.
            atomic_wait::wake_one(data);
        }
    }

    unsafe fn waker_drop(data: *const ()) {
        let data = data.cast::<AtomicU32>();
        unsafe { Arc::decrement_strong_count(data) };
    }
}