1 use core::mem::MaybeUninit;
2 use core::ptr;
3 use std::task::Waker;
4 
5 const NUM_WAKERS: usize = 32;
6 
7 /// A list of wakers to be woken.
8 ///
9 /// # Invariants
10 ///
11 /// The first `curr` elements of `inner` are initialized.
12 pub(crate) struct WakeList {
13     inner: [MaybeUninit<Waker>; NUM_WAKERS],
14     curr: usize,
15 }
16 
17 impl WakeList {
new() -> Self18     pub(crate) fn new() -> Self {
19         const UNINIT_WAKER: MaybeUninit<Waker> = MaybeUninit::uninit();
20 
21         Self {
22             inner: [UNINIT_WAKER; NUM_WAKERS],
23             curr: 0,
24         }
25     }
26 
27     #[inline]
can_push(&self) -> bool28     pub(crate) fn can_push(&self) -> bool {
29         self.curr < NUM_WAKERS
30     }
31 
push(&mut self, val: Waker)32     pub(crate) fn push(&mut self, val: Waker) {
33         debug_assert!(self.can_push());
34 
35         self.inner[self.curr] = MaybeUninit::new(val);
36         self.curr += 1;
37     }
38 
wake_all(&mut self)39     pub(crate) fn wake_all(&mut self) {
40         struct DropGuard {
41             start: *mut Waker,
42             end: *mut Waker,
43         }
44 
45         impl Drop for DropGuard {
46             fn drop(&mut self) {
47                 // SAFETY: Both pointers are part of the same object, with `start <= end`.
48                 let len = unsafe { self.end.offset_from(self.start) } as usize;
49                 let slice = ptr::slice_from_raw_parts_mut(self.start, len);
50                 // SAFETY: All elements in `start..len` are initialized, so we can drop them.
51                 unsafe { ptr::drop_in_place(slice) };
52             }
53         }
54 
55         debug_assert!(self.curr <= NUM_WAKERS);
56 
57         let mut guard = {
58             let start = self.inner.as_mut_ptr().cast::<Waker>();
59             // SAFETY: The resulting pointer is in bounds or one after the length of the same object.
60             let end = unsafe { start.add(self.curr) };
61             // Transfer ownership of the wakers in `inner` to `DropGuard`.
62             self.curr = 0;
63             DropGuard { start, end }
64         };
65         while !ptr::eq(guard.start, guard.end) {
66             // SAFETY: `start` is always initialized if `start != end`.
67             let waker = unsafe { ptr::read(guard.start) };
68             // SAFETY: The resulting pointer is in bounds or one after the length of the same object.
69             guard.start = unsafe { guard.start.add(1) };
70             // If this panics, then `guard` will clean up the remaining wakers.
71             waker.wake();
72         }
73     }
74 }
75 
76 impl Drop for WakeList {
drop(&mut self)77     fn drop(&mut self) {
78         let slice =
79             ptr::slice_from_raw_parts_mut(self.inner.as_mut_ptr().cast::<Waker>(), self.curr);
80         // SAFETY: The first `curr` elements are initialized, so we can drop them.
81         unsafe { ptr::drop_in_place(slice) };
82     }
83 }
84