     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

    use std::marker::PhantomData;
    use std::mem::ManuallyDrop;
    use std::ops::Deref;
    use std::sync::Arc;
    use std::task::{RawWaker, RawWakerVTable, Waker};

    /// Simplfied waking interface based on Arcs
    pub(crate) trait Wake: Send + Sync {
        /// Wake by value
        fn wake(self: Arc<Self>);

        /// Wake by reference
        fn wake_by_ref(arc_self: &Arc<Self>);
    }

    /// A `Waker` that is only valid for a given lifetime.
    #[derive(Debug)]
    pub(crate) struct WakerRef<'a> {
        waker: ManuallyDrop<Waker>,
        _p: PhantomData<&'a ()>,
    }

    impl Deref for WakerRef<'_> {
        type Target = Waker;

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

    /// Creates a reference to a `Waker` from a reference to `Arc<impl Wake>`.
    pub(crate) fn waker_ref<W: Wake>(wake: &Arc<W>) -> WakerRef<'_> {
        let ptr = &**wake as *const _ as *const ();

        let waker = unsafe { Waker::from_raw(RawWaker::new(ptr, waker_vtable::<W>())) };

        WakerRef {
            waker: ManuallyDrop::new(waker),
            _p: PhantomData,
        }
    }

    fn waker_vtable<W: Wake>() -> &'static RawWakerVTable {
        &RawWakerVTable::new(
            clone_arc_raw::<W>,
            wake_arc_raw::<W>,
            wake_by_ref_arc_raw::<W>,
            drop_arc_raw::<W>,
        )
    }

    unsafe fn inc_ref_count<T: Wake>(data: *const ()) {
        // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
        let arc = ManuallyDrop::new(Arc::<T>::from_raw(data as *const T));

        // Now increase refcount, but don't drop new refcount either
        let arc_clone: ManuallyDrop<_> = arc.clone();

        // Drop explicitly to avoid clippy warnings
        drop(arc);
        drop(arc_clone);
    }

    unsafe fn clone_arc_raw<T: Wake>(data: *const ()) -> RawWaker {
        inc_ref_count::<T>(data);
        RawWaker::new(data, waker_vtable::<T>())
    }

    unsafe fn wake_arc_raw<T: Wake>(data: *const ()) {
        let arc: Arc<T> = Arc::from_raw(data as *const T);
        Wake::wake(arc);
    }

    // used by `waker_ref`
    unsafe fn wake_by_ref_arc_raw<T: Wake>(data: *const ()) {
        // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
        let arc = ManuallyDrop::new(Arc::<T>::from_raw(data as *const T));
        Wake::wake_by_ref(&arc);
    }

    unsafe fn drop_arc_raw<T: Wake>(data: *const ()) {
        drop(Arc::<T>::from_raw(data as *const T))
    }