Apply review suggestions and fix tests

3 files changed, 132 insertions, 44 deletions

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();
 }