Add waker_fn (#18)

* Add waker_fn

* Add a waker_fn test

* Double sleep times

* More benches

* Prohibit recursive block_on calls

* Reformat code

3 files changed, 76 insertions, 21 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 5fe858a..1bd7ca8 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,13 +1,11 @@
 //! Task abstraction for building executors.
 //!
+//! # Spawning
+//!
 //! To spawn a future onto an executor, we first need to allocate it on the heap and keep some
 //! state alongside it. The state indicates whether the future is ready for polling, waiting to be
 //! woken up, or completed. Such a future is called a *task*.
 //!
-//! This crate helps with task allocation and polling its future to completion.
-//!
-//! # Spawning
-//!
 //! All executors have some kind of queue that holds runnable tasks:
 //!
 //! ```
@@ -89,13 +87,31 @@
 //! Task construction incurs a single allocation that holds its state, the schedule function, and
 //! the future or the result of the future if completed.
 //!
-//! The layout of a task is equivalent to 4 words followed by the schedule function, and then by a
-//! union of the future and its output.
+//! The layout of a task is equivalent to 4 `usize`s followed by the schedule function, and then by
+//! a union of the future and its output.
+//!
+//! # Waking
+//!
+//! The handy [`waker_fn`] constructor converts any function into a [`Waker`]. Every time it is
+//! woken, the function gets called:
+//!
+//! ```
+//! let waker = async_task::waker_fn(|| println!("Wake!"));
+//!
+//! // Prints "Wake!" twice.
+//! waker.wake_by_ref();
+//! waker.wake_by_ref();
+//! ```
+//!
+//! This is useful for implementing single-future executors like [`block_on`].
 //!
 //! [`spawn`]: fn.spawn.html
 //! [`spawn_local`]: fn.spawn_local.html
+//! [`waker_fn`]: fn.waker_fn.html
 //! [`Task`]: struct.Task.html
 //! [`JoinHandle`]: struct.JoinHandle.html
+//! [`Waker`]: https://doc.rust-lang.org/std/task/struct.Waker.html
+//! [`block_on`]: https://github.com/async-rs/async-task/blob/master/examples/block.rs
 
 #![no_std]
 #![warn(missing_docs, missing_debug_implementations, rust_2018_idioms)]
@@ -110,6 +126,8 @@ mod raw;
 mod state;
 mod task;
 mod utils;
+mod waker_fn;
 
 pub use crate::join_handle::JoinHandle;
 pub use crate::task::{spawn, spawn_local, Task};
+pub use crate::waker_fn::waker_fn;
diff --git a/src/raw.rs b/src/raw.rs
index 1440e7e..ed3ee97 100644
--- a/src/raw.rs
+++ b/src/raw.rs
@@ -15,9 +15,6 @@ use crate::Task;
 
 /// The vtable for a task.
 pub(crate) struct TaskVTable {
-    /// The raw waker vtable.
-    pub(crate) raw_waker_vtable: RawWakerVTable,
-
     /// Schedules the task.
     pub(crate) schedule: unsafe fn(*const ()),
 
@@ -101,6 +98,13 @@ where
     F: Future<Output = R> + 'static,
     S: Fn(Task<T>) + Send + Sync + 'static,
 {
+    const RAW_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
+        Self::clone_waker,
+        Self::wake,
+        Self::wake_by_ref,
+        Self::drop_waker,
+    );
+
     /// Allocates a task with the given `future` and `schedule` function.
     ///
     /// It is assumed that initially only the `Task` reference and the `JoinHandle` exist.
@@ -122,12 +126,6 @@ where
                 state: AtomicUsize::new(SCHEDULED | HANDLE | REFERENCE),
                 awaiter: UnsafeCell::new(None),
                 vtable: &TaskVTable {
-                    raw_waker_vtable: RawWakerVTable::new(
-                        Self::clone_waker,
-                        Self::wake,
-                        Self::wake_by_ref,
-                        Self::drop_waker,
-                    ),
                     schedule: Self::schedule,
                     drop_future: Self::drop_future,
                     get_output: Self::get_output,
@@ -335,7 +333,6 @@ where
     /// Clones a waker.
     unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
         let raw = Self::from_ptr(ptr);
-        let raw_waker_vtable = &(*raw.header).vtable.raw_waker_vtable;
 
         // Increment the reference count. With any kind of reference-counted data structure,
         // relaxed ordering is appropriate when incrementing the counter.
@@ -346,7 +343,7 @@ where
             abort();
         }
 
-        RawWaker::new(ptr, raw_waker_vtable)
+        RawWaker::new(ptr, &Self::RAW_WAKER_VTABLE)
     }
 
     /// Drops a waker.
@@ -464,10 +461,7 @@ where
         let raw = Self::from_ptr(ptr);
 
         // Create a context from the raw task pointer and the vtable inside the its header.
-        let waker = ManuallyDrop::new(Waker::from_raw(RawWaker::new(
-            ptr,
-            &(*raw.header).vtable.raw_waker_vtable,
-        )));
+        let waker = ManuallyDrop::new(Waker::from_raw(RawWaker::new(ptr, &Self::RAW_WAKER_VTABLE)));
         let cx = &mut Context::from_waker(&waker);
 
         let mut state = (*raw.header).state.load(Ordering::Acquire);
diff --git a/src/waker_fn.rs b/src/waker_fn.rs
new file mode 100644
index 0000000..37105f1
--- /dev/null
+++ b/src/waker_fn.rs
@@ -0,0 +1,43 @@
+use alloc::sync::Arc;
+use core::mem::{self, ManuallyDrop};
+use core::task::{RawWaker, RawWakerVTable, Waker};
+
+/// Creates a waker from a wake function.
+///
+/// The function gets called every time the waker is woken.
+pub fn waker_fn<F: Fn() + Send + Sync + 'static>(f: F) -> Waker {
+    let raw = Arc::into_raw(Arc::new(f)) as *const ();
+    let vtable = &Helper::<F>::VTABLE;
+    unsafe { Waker::from_raw(RawWaker::new(raw, vtable)) }
+}
+
+struct Helper<F>(F);
+
+impl<F: Fn() + Send + Sync + 'static> Helper<F> {
+    const VTABLE: RawWakerVTable = RawWakerVTable::new(
+        Self::clone_waker,
+        Self::wake,
+        Self::wake_by_ref,
+        Self::drop_waker,
+    );
+
+    unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
+        let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const F));
+        mem::forget(arc.clone());
+        RawWaker::new(ptr, &Self::VTABLE)
+    }
+
+    unsafe fn wake(ptr: *const ()) {
+        let arc = Arc::from_raw(ptr as *const F);
+        (arc)();
+    }
+
+    unsafe fn wake_by_ref(ptr: *const ()) {
+        let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const F));
+        (arc)();
+    }
+
+    unsafe fn drop_waker(ptr: *const ()) {
+        drop(Arc::from_raw(ptr as *const F));
+    }
+}