1 use crate::loom::sync::Arc;
2
3 use std::marker::PhantomData;
4 use std::mem::ManuallyDrop;
5 use std::ops::Deref;
6 use std::task::{RawWaker, RawWakerVTable, Waker};
7
8 /// Simplified waking interface based on Arcs.
9 pub(crate) trait Wake: Send + Sync + Sized + 'static {
10 /// Wake by value.
wake(arc_self: Arc<Self>)11 fn wake(arc_self: Arc<Self>);
12
13 /// Wake by reference.
wake_by_ref(arc_self: &Arc<Self>)14 fn wake_by_ref(arc_self: &Arc<Self>);
15 }
16
17 /// A `Waker` that is only valid for a given lifetime.
18 #[derive(Debug)]
19 pub(crate) struct WakerRef<'a> {
20 waker: ManuallyDrop<Waker>,
21 _p: PhantomData<&'a ()>,
22 }
23
24 impl Deref for WakerRef<'_> {
25 type Target = Waker;
26
deref(&self) -> &Waker27 fn deref(&self) -> &Waker {
28 &self.waker
29 }
30 }
31
32 /// Creates a reference to a `Waker` from a reference to `Arc<impl Wake>`.
waker_ref<W: Wake>(wake: &Arc<W>) -> WakerRef<'_>33 pub(crate) fn waker_ref<W: Wake>(wake: &Arc<W>) -> WakerRef<'_> {
34 let ptr = Arc::as_ptr(wake).cast::<()>();
35
36 let waker = unsafe { Waker::from_raw(RawWaker::new(ptr, waker_vtable::<W>())) };
37
38 WakerRef {
39 waker: ManuallyDrop::new(waker),
40 _p: PhantomData,
41 }
42 }
43
waker_vtable<W: Wake>() -> &'static RawWakerVTable44 fn waker_vtable<W: Wake>() -> &'static RawWakerVTable {
45 &RawWakerVTable::new(
46 clone_arc_raw::<W>,
47 wake_arc_raw::<W>,
48 wake_by_ref_arc_raw::<W>,
49 drop_arc_raw::<W>,
50 )
51 }
52
clone_arc_raw<T: Wake>(data: *const ()) -> RawWaker53 unsafe fn clone_arc_raw<T: Wake>(data: *const ()) -> RawWaker {
54 Arc::<T>::increment_strong_count(data as *const T);
55 RawWaker::new(data, waker_vtable::<T>())
56 }
57
wake_arc_raw<T: Wake>(data: *const ())58 unsafe fn wake_arc_raw<T: Wake>(data: *const ()) {
59 let arc: Arc<T> = Arc::from_raw(data as *const T);
60 Wake::wake(arc);
61 }
62
63 // used by `waker_ref`
wake_by_ref_arc_raw<T: Wake>(data: *const ())64 unsafe fn wake_by_ref_arc_raw<T: Wake>(data: *const ()) {
65 // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
66 let arc = ManuallyDrop::new(Arc::<T>::from_raw(data.cast()));
67 Wake::wake_by_ref(&arc);
68 }
69
drop_arc_raw<T: Wake>(data: *const ())70 unsafe fn drop_arc_raw<T: Wake>(data: *const ()) {
71 drop(Arc::<T>::from_raw(data.cast()));
72 }
73