1use crate::loom::sync::Arc;
2
3use std::marker::PhantomData;
4use std::mem::ManuallyDrop;
5use std::ops::Deref;
6use std::task::{RawWaker, RawWakerVTable, Waker};
7
8/// Simplified waking interface based on Arcs.
9pub(crate) trait Wake: Send + Sync + Sized + 'static {
10 /// Wake by value.
11 fn wake(arc_self: Arc<Self>);
12
13 /// Wake by reference.
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)]
19pub(crate) struct WakerRef<'a> {
20 waker: ManuallyDrop<Waker>,
21 _p: PhantomData<&'a ()>,
22}
23
24impl Deref for WakerRef<'_> {
25 type Target = Waker;
26
27 fn deref(&self) -> &Waker {
28 &self.waker
29 }
30}
31
32/// Creates a reference to a `Waker` from a reference to `Arc<impl Wake>`.
33pub(crate) fn waker_ref<W: Wake>(wake: &Arc<W>) -> WakerRef<'_> {
34 let ptr: *const () = Arc::as_ptr(this:wake) as *const ();
35
36 let waker: Waker = unsafe { Waker::from_raw(waker:RawWaker::new(data:ptr, vtable:waker_vtable::<W>())) };
37
38 WakerRef {
39 waker: ManuallyDrop::new(waker),
40 _p: PhantomData,
41 }
42}
43
44fn 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
53unsafe fn inc_ref_count<T: Wake>(data: *const ()) {
54 // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
55 let arc: ManuallyDrop> = ManuallyDrop::new(Arc::<T>::from_raw(ptr:data as *const T));
56
57 // Now increase refcount, but don't drop new refcount either
58 let _arc_clone: ManuallyDrop<_> = arc.clone();
59}
60
61unsafe fn clone_arc_raw<T: Wake>(data: *const ()) -> RawWaker {
62 inc_ref_count::<T>(data);
63 RawWaker::new(data, vtable:waker_vtable::<T>())
64}
65
66unsafe fn wake_arc_raw<T: Wake>(data: *const ()) {
67 let arc: Arc<T> = Arc::from_raw(ptr:data as *const T);
68 Wake::wake(arc_self:arc);
69}
70
71// used by `waker_ref`
72unsafe fn wake_by_ref_arc_raw<T: Wake>(data: *const ()) {
73 // Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
74 let arc: ManuallyDrop> = ManuallyDrop::new(Arc::<T>::from_raw(ptr:data as *const T));
75 Wake::wake_by_ref(&arc);
76}
77
78unsafe fn drop_arc_raw<T: Wake>(data: *const ()) {
79 drop(Arc::<T>::from_raw(ptr:data as *const T))
80}
81