use std::marker::PhantomData;
use std::mem::ManuallyDrop;
use std::ptr::NonNull;
use std::task::{RawWaker, RawWakerVTable, Waker};
use std::{fmt, rc::Rc};
use crate::rt::task;
const NATIVE: NonNull<()> = NonNull::dangling();
const _: () = assert!(align_of::<RawWakerVTable>() > 1);
#[repr(C)]
pub struct LocalWaker {
data: *const (),
metadata: NonNull<()>,
_local: PhantomData<Rc<()>>,
}
#[cfg(target_pointer_width = "64")]
const _: () = {
assert!(size_of::<LocalWaker>() == 16);
assert!(size_of::<Option<LocalWaker>>() == 16);
};
impl LocalWaker {
#[inline]
pub fn from_waker(waker: &Waker) -> Self {
if let Some(ptr) = task::local_waker_ptr(waker) {
unsafe { task::clone_local_waker(ptr) };
Self::native(ptr)
} else {
Self::fallback(waker.clone())
}
}
#[inline]
pub fn wake(self) {
if self.is_native() {
let this = ManuallyDrop::new(self);
unsafe { task::wake_local_waker(this.data, true) };
} else {
self.into_waker().wake();
}
}
#[inline]
pub fn wake_by_ref(&self) {
if self.is_native() {
unsafe { task::wake_local_waker(self.data, false) };
} else {
self.borrowed_waker().wake_by_ref();
}
}
#[inline]
pub fn will_wake(&self, waker: &Waker) -> bool {
self.borrowed_waker().will_wake(waker)
}
#[inline]
pub fn into_waker(self) -> Waker {
let this = ManuallyDrop::new(self);
if this.is_native() {
unsafe { task::into_waker(this.data) }
} else {
unsafe { Waker::from_raw(this.raw_waker()) }
}
}
fn native(data: *const ()) -> Self {
Self {
data,
metadata: NATIVE,
_local: PhantomData,
}
}
fn fallback(waker: Waker) -> Self {
let waker = ManuallyDrop::new(waker);
Self {
data: waker.data(),
metadata: NonNull::from(waker.vtable()).cast(),
_local: PhantomData,
}
}
fn is_native(&self) -> bool {
self.metadata == NATIVE
}
fn raw_waker(&self) -> RawWaker {
if self.is_native() {
task::raw_waker(self.data)
} else {
RawWaker::new(self.data, unsafe {
self.metadata.cast::<RawWakerVTable>().as_ref()
})
}
}
fn borrowed_waker(&self) -> ManuallyDrop<Waker> {
ManuallyDrop::new(unsafe { Waker::from_raw(self.raw_waker()) })
}
}
impl From<Waker> for LocalWaker {
#[inline]
fn from(waker: Waker) -> Self {
if let Some(ptr) = task::local_waker_ptr(&waker) {
unsafe { task::clone_local_waker(ptr) };
drop(waker);
Self::native(ptr)
} else {
Self::fallback(waker)
}
}
}
impl Clone for LocalWaker {
#[inline]
fn clone(&self) -> Self {
if self.is_native() {
unsafe { task::clone_local_waker(self.data) };
Self::native(self.data)
} else {
Self::fallback((*self.borrowed_waker()).clone())
}
}
#[inline]
fn clone_from(&mut self, source: &Self) {
if self.data != source.data || self.metadata != source.metadata {
*self = source.clone();
}
}
}
impl Drop for LocalWaker {
#[inline]
fn drop(&mut self) {
if self.is_native() {
unsafe { task::drop_local_waker(self.data) };
} else {
drop(unsafe { Waker::from_raw(self.raw_waker()) });
}
}
}
impl fmt::Debug for LocalWaker {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("LocalWaker")
.field("data", &self.data)
.field("native", &self.is_native())
.finish()
}
}
#[cfg(all(test, not(ntex_loom)))]
#[path = "../tests/unit/local_waker.rs"]
mod tests;