      1
      2
      3
      4
      5
      6
      7
      8
      9
     10
     11
     12
     13
     14
     15
     16
     17
     18
     19
     20
     21
     22
     23
     24
     25
     26
     27
     28
     29
     30
     31
     32
     33
     34
     35
     36
     37
     38
     39
     40
     41
     42
     43
     44
     45
     46
     47
     48
     49
     50
     51
     52
     53
     54
     55
     56
     57
     58
     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101

    use crate::runtime::task::harness::Harness;
    use crate::runtime::task::{Header, Schedule};

    use std::future::Future;
    use std::marker::PhantomData;
    use std::mem::ManuallyDrop;
    use std::ops;
    use std::ptr::NonNull;
    use std::task::{RawWaker, RawWakerVTable, Waker};

    pub(super) struct WakerRef<'a, S: 'static> {
        waker: ManuallyDrop<Waker>,
        _p: PhantomData<(&'a Header, S)>,
    }

    /// Returns a `WakerRef` which avoids having to pre-emptively increase the
    /// refcount if there is no need to do so.
    pub(super) fn waker_ref<T, S>(header: &Header) -> WakerRef<'_, S>
    where
        T: Future,
        S: Schedule,
    {
        // `Waker::will_wake` uses the VTABLE pointer as part of the check. This
        // means that `will_wake` will always return false when using the current
        // task's waker. (discussion at rust-lang/rust#66281).
        //
        // To fix this, we use a single vtable. Since we pass in a reference at this
        // point and not an *owned* waker, we must ensure that `drop` is never
        // called on this waker instance. This is done by wrapping it with
        // `ManuallyDrop` and then never calling drop.
        let waker = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker::<T, S>(header))) };

        WakerRef {
            waker,
            _p: PhantomData,
        }
    }

    impl<S> ops::Deref for WakerRef<'_, S> {
        type Target = Waker;

        fn deref(&self) -> &Waker {
            &self.waker
        }
    }

    unsafe fn clone_waker<T, S>(ptr: *const ()) -> RawWaker
    where
        T: Future,
        S: Schedule,
    {
        let header = ptr as *const Header;
        (*header).state.ref_inc();
        raw_waker::<T, S>(header)
    }

    unsafe fn drop_waker<T, S>(ptr: *const ())
    where
        T: Future,
        S: Schedule,
    {
        let ptr = NonNull::new_unchecked(ptr as *mut Header);
        let harness = Harness::<T, S>::from_raw(ptr);
        harness.drop_reference();
    }

    unsafe fn wake_by_val<T, S>(ptr: *const ())
    where
        T: Future,
        S: Schedule,
    {
        let ptr = NonNull::new_unchecked(ptr as *mut Header);
        let harness = Harness::<T, S>::from_raw(ptr);
        harness.wake_by_val();
    }

    // Wake without consuming the waker
    unsafe fn wake_by_ref<T, S>(ptr: *const ())
    where
        T: Future,
        S: Schedule,
    {
        let ptr = NonNull::new_unchecked(ptr as *mut Header);
        let harness = Harness::<T, S>::from_raw(ptr);
        harness.wake_by_ref();
    }

    fn raw_waker<T, S>(header: *const Header) -> RawWaker
    where
        T: Future,
        S: Schedule,
    {
        let ptr = header as *const ();
        let vtable = &RawWakerVTable::new(
            clone_waker::<T, S>,
            wake_by_val::<T, S>,
            wake_by_ref::<T, S>,
            drop_waker::<T, S>,
        );
        RawWaker::new(ptr, vtable)
    }