Apply review suggestions and fix tests

4 files changed, 163 insertions, 78 deletions

diff --git a/src/libcore/task/wake.rs b/src/libcore/task/wake.rs
index adaca504345..fe2de61c594 100644
--- a/src/libcore/task/wake.rs
+++ b/src/libcore/task/wake.rs
@@ -5,14 +5,14 @@
 use fmt;
 use marker::Unpin;
 
-/// A `RawWaker` allows the implementor of a task executor to create a `Waker`
+/// A `RawWaker` allows the implementor of a task executor to create a [`Waker`]
 /// which provides customized wakeup behavior.
 ///
 /// [vtable]: https://en.wikipedia.org/wiki/Virtual_method_table
 ///
 /// It consists of a data pointer and a [virtual function pointer table (vtable)][vtable] that
 /// customizes the behavior of the `RawWaker`.
-#[derive(PartialEq)]
+#[derive(PartialEq, Debug)]
 pub struct RawWaker {
     /// A data pointer, which can be used to store arbitrary data as required
     /// by the executor. This could be e.g. a type-erased pointer to an `Arc`
@@ -24,55 +24,41 @@ pub struct RawWaker {
     pub vtable: &'static RawWakerVTable,
 }
 
-impl fmt::Debug for RawWaker {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        f.debug_struct("RawWaker")
-            .finish()
-    }
-}
-
 /// A virtual function pointer table (vtable) that specifies the behavior
 /// of a [`RawWaker`].
 ///
 /// The pointer passed to all functions inside the vtable is the `data` pointer
-/// from the enclosing `RawWaker` object.
-#[derive(PartialEq, Copy, Clone)]
+/// from the enclosing [`RawWaker`] object.
+#[derive(PartialEq, Copy, Clone, Debug)]
 pub struct RawWakerVTable {
-    /// This function will be called when the `RawWaker` gets cloned, e.g. when
-    /// the `Waker` in which the `RawWaker` is stored gets cloned.
+    /// This function will be called when the [`RawWaker`] gets cloned, e.g. when
+    /// the [`Waker`] in which the [`RawWaker`] is stored gets cloned.
     ///
     /// The implementation of this function must retain all resources that are
-    /// required for this additional instance of a `RawWaker` and associated
-    /// task. Calling `wake` on the resulting `RawWaker` should result in a wakeup
-    /// of the same task that would have been awoken by the original `RawWaker`.
+    /// required for this additional instance of a [`RawWaker`] and associated
+    /// task. Calling `wake` on the resulting [`RawWaker`] should result in a wakeup
+    /// of the same task that would have been awoken by the original [`RawWaker`].
     pub clone: unsafe fn(*const ()) -> RawWaker,
 
-    /// This function will be called when `wake` is called on the `Waker`.
-    /// It must wake up the task associated with this `RawWaker`.
+    /// This function will be called when `wake` is called on the [`Waker`].
+    /// It must wake up the task associated with this [`RawWaker`].
     pub wake: unsafe fn(*const ()),
 
-    /// This function gets called when a `RawWaker` gets dropped.
+    /// This function gets called when a [`RawWaker`] gets dropped.
     ///
     /// The implementation of this function must make sure to release any
-    /// resources that are associated with this instance of a `RawWaker` and
+    /// resources that are associated with this instance of a [`RawWaker`] and
     /// associated task.
-    pub drop_fn: unsafe fn(*const ()),
-}
-
-impl fmt::Debug for RawWakerVTable {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        f.debug_struct("RawWakerVTable")
-            .finish()
-    }
+    pub drop: unsafe fn(*const ()),
 }
 
 /// A `Waker` is a handle for waking up a task by notifying its executor that it
 /// is ready to be run.
 ///
-/// This handle encapsulates a `RawWaker` instance, which defines the
+/// This handle encapsulates a [`RawWaker`] instance, which defines the
 /// executor-specific wakeup behavior.
 ///
-/// Implements `Clone`, `Send`, and `Sync`.
+/// Implements [`Clone`], [`Send`], and [`Sync`].
 #[repr(transparent)]
 pub struct Waker {
     waker: RawWaker,
@@ -87,6 +73,10 @@ impl Waker {
     pub fn wake(&self) {
         // The actual wakeup call is delegated through a virtual function call
         // to the implementation which is defined by the executor.
+
+        // SAFETY: This is safe because `Waker::new_unchecked` is the only way
+        // to initialize `wake` and `data` requiring the user to acknowledge
+        // that the contract of `RawWaker` is upheld.
         unsafe { (self.waker.vtable.wake)(self.waker.data) }
     }
 
@@ -101,10 +91,11 @@ impl Waker {
         self.waker == other.waker
     }
 
-    /// Creates a new `Waker` from `RawWaker`.
+    /// Creates a new `Waker` from [`RawWaker`].
     ///
-    /// The method cannot check whether `RawWaker` fulfills the required API
-    /// contract to make it usable for `Waker` and is therefore unsafe.
+    /// The behavior of the returned `Waker` is undefined if the contract defined
+    /// in [RawWaker]'s documentation is not upheld. Therefore this method is
+    /// unsafe.
     pub unsafe fn new_unchecked(waker: RawWaker) -> Waker {
         Waker {
             waker,
@@ -115,6 +106,9 @@ impl Waker {
 impl Clone for Waker {
     fn clone(&self) -> Self {
         Waker {
+            // SAFETY: This is safe because `Waker::new_unchecked` is the only way
+            // to initialize `clone` and `data` requiring the user to acknowledge
+            // that the contract of [`RawWaker`] is upheld.
             waker: unsafe { (self.waker.vtable.clone)(self.waker.data) },
         }
     }
@@ -122,7 +116,10 @@ impl Clone for Waker {
 
 impl Drop for Waker {
     fn drop(&mut self) {
-        unsafe { (self.waker.vtable.drop_fn)(self.waker.data) }
+        // SAFETY: This is safe because `Waker::new_unchecked` is the only way
+        // to initialize `drop` and `data` requiring the user to acknowledge
+        // that the contract of `RawWaker` is upheld.
+        unsafe { (self.waker.vtable.drop)(self.waker.data) }
     }
 }
 
diff --git a/src/test/compile-fail/must_use-in-stdlib-traits.rs b/src/test/compile-fail/must_use-in-stdlib-traits.rs
index 7e446fdaeaf..b4f07ab3321 100644
--- a/src/test/compile-fail/must_use-in-stdlib-traits.rs
+++ b/src/test/compile-fail/must_use-in-stdlib-traits.rs
@@ -4,7 +4,7 @@
 use std::iter::Iterator;
 use std::future::Future;
 
-use std::task::{Poll, LocalWaker};
+use std::task::{Poll, Waker};
 use std::pin::Pin;
 use std::unimplemented;
 
@@ -13,7 +13,7 @@ struct MyFuture;
 impl Future for MyFuture {
    type Output = u32;
 
-   fn poll(self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<u32> {
+   fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll<u32> {
       Poll::Pending
    }
 }
diff --git a/src/test/run-pass/async-await.rs b/src/test/run-pass/async-await.rs
index 552f3821573..2def62a6ba6 100644
--- a/src/test/run-pass/async-await.rs
+++ b/src/test/run-pass/async-await.rs
@@ -9,17 +9,70 @@ use std::sync::{
     atomic::{self, AtomicUsize},
 };
 use std::task::{
-    LocalWaker, Poll, Wake,
-    local_waker_from_nonlocal,
+    Poll, Waker, RawWaker, RawWakerVTable,
 };
 
+macro_rules! waker_vtable {
+    ($ty:ident) => {
+        &RawWakerVTable {
+            clone: clone_arc_raw::<$ty>,
+            drop: drop_arc_raw::<$ty>,
+            wake: wake_arc_raw::<$ty>,
+        }
+    };
+}
+
+pub trait ArcWake {
+    fn wake(arc_self: &Arc<Self>);
+
+    fn into_waker(wake: Arc<Self>) -> Waker where Self: Sized
+    {
+        let ptr = Arc::into_raw(wake) as *const();
+
+        unsafe {
+            Waker::new_unchecked(RawWaker{
+                data: ptr,
+                vtable: waker_vtable!(Self),
+            })
+        }
+    }
+}
+
+unsafe fn increase_refcount<T: ArcWake>(data: *const()) {
+    // Retain Arc by creating a copy
+    let arc: Arc<T> = Arc::from_raw(data as *const T);
+    let arc_clone = arc.clone();
+    // Forget the Arcs again, so that the refcount isn't decrased
+    let _ = Arc::into_raw(arc);
+    let _ = Arc::into_raw(arc_clone);
+}
+
+unsafe fn clone_arc_raw<T: ArcWake>(data: *const()) -> RawWaker {
+    increase_refcount::<T>(data);
+    RawWaker {
+        data: data,
+        vtable: waker_vtable!(T),
+    }
+}
+
+unsafe fn drop_arc_raw<T: ArcWake>(data: *const()) {
+    // Drop Arc
+    let _: Arc<T> = Arc::from_raw(data as *const T);
+}
+
+unsafe fn wake_arc_raw<T: ArcWake>(data: *const()) {
+    let arc: Arc<T> = Arc::from_raw(data as *const T);
+    ArcWake::wake(&arc);
+    let _ = Arc::into_raw(arc);
+}
+
 struct Counter {
     wakes: AtomicUsize,
 }
 
-impl Wake for Counter {
-    fn wake(this: &Arc<Self>) {
-        this.wakes.fetch_add(1, atomic::Ordering::SeqCst);
+impl ArcWake for Counter {
+    fn wake(arc_self: &Arc<Self>) {
+        arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
     }
 }
 
@@ -29,11 +82,11 @@ fn wake_and_yield_once() -> WakeOnceThenComplete { WakeOnceThenComplete(false) }
 
 impl Future for WakeOnceThenComplete {
     type Output = ();
-    fn poll(mut self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<()> {
+    fn poll(mut self: Pin<&mut Self>, waker: &Waker) -> Poll<()> {
         if self.0 {
             Poll::Ready(())
         } else {
-            lw.wake();
+            waker.wake();
             self.0 = true;
             Poll::Pending
         }
@@ -130,7 +183,7 @@ where
 {
     let mut fut = Box::pin(f(9));
     let counter = Arc::new(Counter { wakes: AtomicUsize::new(0) });
-    let waker = local_waker_from_nonlocal(counter.clone());
+    let waker = ArcWake::into_waker(counter.clone());
     assert_eq!(0, counter.wakes.load(atomic::Ordering::SeqCst));
     assert_eq!(Poll::Pending, fut.as_mut().poll(&waker));
     assert_eq!(1, counter.wakes.load(atomic::Ordering::SeqCst));
diff --git a/src/test/run-pass/futures-api.rs b/src/test/run-pass/futures-api.rs
index e3023521100..058d09e2420 100644
--- a/src/test/run-pass/futures-api.rs
+++ b/src/test/run-pass/futures-api.rs
@@ -3,28 +3,75 @@
 
 use std::future::Future;
 use std::pin::Pin;
-use std::rc::Rc;
 use std::sync::{
     Arc,
     atomic::{self, AtomicUsize},
 };
 use std::task::{
-    Poll, Wake, Waker, LocalWaker,
-    local_waker, local_waker_from_nonlocal,
+    Poll, Waker, RawWaker, RawWakerVTable,
 };
 
-struct Counter {
-    local_wakes: AtomicUsize,
-    nonlocal_wakes: AtomicUsize,
+macro_rules! waker_vtable {
+    ($ty:ident) => {
+        &RawWakerVTable {
+            clone: clone_arc_raw::<$ty>,
+            drop: drop_arc_raw::<$ty>,
+            wake: wake_arc_raw::<$ty>,
+        }
+    };
 }
 
-impl Wake for Counter {
-    fn wake(this: &Arc<Self>) {
-        this.nonlocal_wakes.fetch_add(1, atomic::Ordering::SeqCst);
+pub trait ArcWake {
+    fn wake(arc_self: &Arc<Self>);
+
+    fn into_waker(wake: Arc<Self>) -> Waker where Self: Sized
+    {
+        let ptr = Arc::into_raw(wake) as *const();
+
+        unsafe {
+            Waker::new_unchecked(RawWaker{
+                data: ptr,
+                vtable: waker_vtable!(Self),
+            })
+        }
     }
+}
+
+unsafe fn increase_refcount<T: ArcWake>(data: *const()) {
+    // Retain Arc by creating a copy
+    let arc: Arc<T> = Arc::from_raw(data as *const T);
+    let arc_clone = arc.clone();
+    // Forget the Arcs again, so that the refcount isn't decrased
+    let _ = Arc::into_raw(arc);
+    let _ = Arc::into_raw(arc_clone);
+}
 
-    unsafe fn wake_local(this: &Arc<Self>) {
-        this.local_wakes.fetch_add(1, atomic::Ordering::SeqCst);
+unsafe fn clone_arc_raw<T: ArcWake>(data: *const()) -> RawWaker {
+    increase_refcount::<T>(data);
+    RawWaker {
+        data: data,
+        vtable: waker_vtable!(T),
+    }
+}
+
+unsafe fn drop_arc_raw<T: ArcWake>(data: *const()) {
+    // Drop Arc
+    let _: Arc<T> = Arc::from_raw(data as *const T);
+}
+
+unsafe fn wake_arc_raw<T: ArcWake>(data: *const()) {
+    let arc: Arc<T> = Arc::from_raw(data as *const T);
+    ArcWake::wake(&arc);
+    let _ = Arc::into_raw(arc);
+}
+
+struct Counter {
+    wakes: AtomicUsize,
+}
+
+impl Wake for Counter {
+    fn wake(arc_self: &Arc<Self>) {
+        arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
     }
 }
 
@@ -32,40 +79,28 @@ struct MyFuture;
 
 impl Future for MyFuture {
     type Output = ();
-    fn poll(self: Pin<&mut Self>, lw: &LocalWaker) -> Poll<Self::Output> {
-        // Wake once locally
-        lw.wake();
-        // Wake twice non-locally
-        let waker = lw.clone().into_waker();
+    fn poll(self: Pin<&mut Self>, waker: &Waker) -> Poll<Self::Output> {
+        // Wake twice
         waker.wake();
         waker.wake();
         Poll::Ready(())
     }
 }
 
-fn test_local_waker() {
+fn test_waker() {
     let counter = Arc::new(Counter {
-        local_wakes: AtomicUsize::new(0),
-        nonlocal_wakes: AtomicUsize::new(0),
+        wakes: AtomicUsize::new(0),
     });
-    let waker = unsafe { local_waker(counter.clone()) };
-    assert_eq!(Poll::Ready(()), Pin::new(&mut MyFuture).poll(&waker));
-    assert_eq!(1, counter.local_wakes.load(atomic::Ordering::SeqCst));
-    assert_eq!(2, counter.nonlocal_wakes.load(atomic::Ordering::SeqCst));
-}
+    let waker = ArcWake::into_waker(counter.clone());
+    assert_eq!(2, Arc::strong_count(&counter));
 
-fn test_local_as_nonlocal_waker() {
-    let counter = Arc::new(Counter {
-        local_wakes: AtomicUsize::new(0),
-        nonlocal_wakes: AtomicUsize::new(0),
-    });
-    let waker: LocalWaker = local_waker_from_nonlocal(counter.clone());
     assert_eq!(Poll::Ready(()), Pin::new(&mut MyFuture).poll(&waker));
-    assert_eq!(0, counter.local_wakes.load(atomic::Ordering::SeqCst));
-    assert_eq!(3, counter.nonlocal_wakes.load(atomic::Ordering::SeqCst));
+    assert_eq!(2, counter.wakes.load(atomic::Ordering::SeqCst));
+
+    drop(waker);
+    assert_eq!(1, Arc::strong_count(&counter));
 }
 
 fn main() {
-    test_local_waker();
-    test_local_as_nonlocal_waker();
+    test_waker();
 }