use std::task::Waker;
use super::task::TaskId;
pub(crate) struct TaskWaker;
#[cfg(target_pointer_width = "64")]
mod raw_task_waker {
use super::TaskId;
use std::ptr::NonNull;
use std::task::{RawWaker, RawWakerVTable, Waker};
#[inline]
fn encode(task_id: TaskId) -> *const () {
NonNull::<u8>::dangling()
.as_ptr()
.wrapping_add(task_id.0 as usize)
.cast::<()>()
}
#[inline]
fn decode(data: *const ()) -> TaskId {
let base = NonNull::<u8>::dangling().as_ptr() as usize;
TaskId((data as usize).wrapping_sub(base) as u64)
}
unsafe fn clone_waker(data: *const ()) -> RawWaker {
RawWaker::new(data, &VTABLE)
}
unsafe fn wake_waker(data: *const ()) {
super::super::wake_task(decode(data));
}
unsafe fn wake_by_ref_waker(data: *const ()) {
super::super::wake_task(decode(data));
}
unsafe fn drop_waker(_data: *const ()) {}
static VTABLE: RawWakerVTable =
RawWakerVTable::new(clone_waker, wake_waker, wake_by_ref_waker, drop_waker);
#[inline]
pub(super) fn into_waker(task_id: TaskId) -> Waker {
unsafe { Waker::from_raw(RawWaker::new(encode(task_id), &VTABLE)) }
}
}
#[cfg(not(target_pointer_width = "64"))]
mod raw_task_waker {
use super::TaskId;
use std::sync::Arc;
use std::task::{Wake, Waker};
struct ArcTaskWaker(TaskId);
impl Wake for ArcTaskWaker {
fn wake(self: Arc<Self>) {
super::super::wake_task(self.0);
}
fn wake_by_ref(self: &Arc<Self>) {
super::super::wake_task(self.0);
}
}
pub(super) fn into_waker(task_id: TaskId) -> Waker {
Waker::from(Arc::new(ArcTaskWaker(task_id)))
}
}
impl TaskWaker {
#[inline]
pub(crate) fn into_waker(task_id: TaskId) -> Waker {
raw_task_waker::into_waker(task_id)
}
}
pub(crate) struct RootWaker;
thread_local! {
static ROOT_WOKEN: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
mod raw_root_waker {
use super::ROOT_WOKEN;
use std::ptr::NonNull;
use std::task::{RawWaker, RawWakerVTable, Waker};
unsafe fn clone_root(data: *const ()) -> RawWaker {
RawWaker::new(data, &VTABLE)
}
unsafe fn wake_root(_data: *const ()) {
ROOT_WOKEN.with(|w| w.set(true));
}
unsafe fn drop_root(_data: *const ()) {}
static VTABLE: RawWakerVTable =
RawWakerVTable::new(clone_root, wake_root, wake_root, drop_root);
#[inline]
pub(super) fn into_waker() -> Waker {
unsafe {
Waker::from_raw(RawWaker::new(
NonNull::<u8>::dangling().as_ptr().cast(),
&VTABLE,
))
}
}
}
impl RootWaker {
#[inline]
pub(crate) fn into_waker() -> Waker {
raw_root_waker::into_waker()
}
#[inline]
pub(crate) fn is_woken() -> bool {
ROOT_WOKEN.with(|w| w.get())
}
#[inline]
pub(crate) fn clear() {
ROOT_WOKEN.with(|w| w.set(false));
}
}