compio_io\compat/
waker_array.rs

1use std::{
2    mem::ManuallyDrop,
3    sync::Arc,
4    task::{RawWaker, RawWakerVTable, Wake, Waker},
5};
6
7/// A reference to an array of wakers, which can be used to create a waker that
8/// wakes all of them.
9pub struct WakerArrayRef<'a, const N: usize>([Option<&'a Waker>; N]);
10
11impl<'a, const N: usize> WakerArrayRef<'a, N> {
12    const VTABLE: &'static RawWakerVTable =
13        &RawWakerVTable::new(Self::clone, Self::wake, Self::wake_by_ref, Self::drop);
14
15    /// Creates a new [`WakerArrayRef`] from an array of optional waker
16    /// references.
17    pub fn new(wakers: [Option<&'a Waker>; N]) -> Self {
18        Self(wakers)
19    }
20
21    /// Creates a temporary waker that wakes all wakers in the array when woken.
22    pub fn with<F, R>(&self, f: F) -> R
23    where
24        F: FnOnce(&Waker) -> R,
25    {
26        let waker = unsafe { Waker::new(self as *const Self as *const (), Self::VTABLE) };
27        f(&waker)
28    }
29
30    fn wake_impl(&self) {
31        for waker in self.0.iter().flatten() {
32            waker.wake_by_ref();
33        }
34    }
35
36    fn to_owned(&self) -> WakerArray<N> {
37        WakerArray(self.0.map(|waker| waker.cloned()))
38    }
39
40    unsafe fn from_raw<'s>(ptr: *const ()) -> &'s Self {
41        unsafe { &*ptr.cast::<Self>() }
42    }
43
44    unsafe fn clone(ptr: *const ()) -> RawWaker {
45        let this = unsafe { Self::from_raw(ptr) };
46        let owned = this.to_owned();
47        let waker = ManuallyDrop::new(Waker::from(Arc::new(owned)));
48        RawWaker::new(waker.data(), waker.vtable())
49    }
50
51    unsafe fn wake(_: *const ()) {
52        unreachable!("WakerArrayRef will only be accessed with reference")
53    }
54
55    unsafe fn wake_by_ref(ptr: *const ()) {
56        unsafe { Self::from_raw(ptr) }.wake_impl();
57    }
58
59    unsafe fn drop(_: *const ()) {
60        // `WakerArrayRef` only contains reference, no need to drop.
61    }
62}
63
64struct WakerArray<const N: usize>([Option<Waker>; N]);
65
66impl<const N: usize> Wake for WakerArray<N> {
67    fn wake(self: Arc<Self>) {
68        self.wake_by_ref();
69    }
70
71    fn wake_by_ref(self: &Arc<Self>) {
72        for waker in self.0.iter().flatten() {
73            waker.wake_by_ref();
74        }
75    }
76}
compio_io\compat/waker_array.rs

compio_io\compat/
waker_array.rs
