rt: remove `unsafe` from shell runtime. (#2333)

Since the original shell runtime was implemented, utilities have been
added to encapsulate `unsafe`. The shell runtime is now able to use
those utilities and not include its own `unsafe` code.

1 files changed, 22 insertions, 43 deletions

diff --git a/tokio/src/runtime/shell.rs b/tokio/src/runtime/shell.rs
index efe77f20..294f2a16 100644
--- a/tokio/src/runtime/shell.rs
+++ b/tokio/src/runtime/shell.rs
@@ -1,49 +1,43 @@
 #![allow(clippy::redundant_clone)]
 
-use crate::park::Park;
+use crate::park::{Park, Unpark};
 use crate::runtime::enter;
 use crate::runtime::time;
+use crate::util::{waker_ref, Wake};
 
 use std::future::Future;
-use std::mem::ManuallyDrop;
-use std::pin::Pin;
 use std::sync::Arc;
+use std::task::Context;
 use std::task::Poll::Ready;
-use std::task::{Context, RawWaker, RawWakerVTable, Waker};
 
 #[derive(Debug)]
 pub(super) struct Shell {
     driver: time::Driver,
 
     /// TODO: don't store this
-    waker: Waker,
+    unpark: Arc<Handle>,
 }
 
-type Handle = <time::Driver as Park>::Unpark;
+#[derive(Debug)]
+struct Handle(<time::Driver as Park>::Unpark);
 
 impl Shell {
     pub(super) fn new(driver: time::Driver) -> Shell {
-        // Make sure we don't mess up types (as we do casts later)
-        let unpark: Arc<Handle> = Arc::new(driver.unpark());
-
-        let raw_waker = RawWaker::new(
-            Arc::into_raw(unpark) as *const Handle as *const (),
-            &RawWakerVTable::new(clone_waker, wake, wake_by_ref, drop_waker),
-        );
+        let unpark = Arc::new(Handle(driver.unpark()));
 
-        let waker = unsafe { Waker::from_raw(raw_waker) };
-
-        Shell { driver, waker }
+        Shell { driver, unpark }
     }
 
-    pub(super) fn block_on<F>(&mut self, mut f: F) -> F::Output
+    pub(super) fn block_on<F>(&mut self, f: F) -> F::Output
     where
         F: Future,
     {
         let _e = enter();
 
-        let mut f = unsafe { Pin::new_unchecked(&mut f) };
-        let mut cx = Context::from_waker(&self.waker);
+        pin!(f);
+
+        let waker = waker_ref(&self.unpark);
+        let mut cx = Context::from_waker(&waker);
 
         loop {
             if let Ready(v) = crate::coop::budget(|| f.as_mut().poll(&mut cx)) {
@@ -55,29 +49,14 @@ impl Shell {
     }
 }
 
-fn clone_waker(ptr: *const ()) -> RawWaker {
-    let w1 = unsafe { ManuallyDrop::new(Arc::from_raw(ptr as *const Handle)) };
-    let _w2 = ManuallyDrop::new(w1.clone());
-
-    RawWaker::new(
-        ptr,
-        &RawWakerVTable::new(clone_waker, wake, wake_by_ref, drop_waker),
-    )
-}
-
-fn wake(ptr: *const ()) {
-    use crate::park::Unpark;
-    let unpark = unsafe { Arc::from_raw(ptr as *const Handle) };
-    (unpark).unpark()
-}
-
-fn wake_by_ref(ptr: *const ()) {
-    use crate::park::Unpark;
-
-    let unpark = ptr as *const Handle;
-    unsafe { (*unpark).unpark() }
-}
+impl Wake for Handle {
+    /// Wake by value
+    fn wake(self: Arc<Self>) {
+        Wake::wake_by_ref(&self);
+    }
 
-fn drop_waker(ptr: *const ()) {
-    let _ = unsafe { Arc::from_raw(ptr as *const Handle) };
+    /// Wake by reference
+    fn wake_by_ref(arc_self: &Arc<Self>) {
+        arc_self.0.unpark();
+    }
 }