use crate::async_engine::AsyncEngine;
use crate::async_engine::Phase;
use crate::tracy::ZoneName;
use crate::utils::numcell::NumCell;
use crate::utils::ptr_ext::MutPtrExt;
use crate::utils::ptr_ext::PtrExt;
use std::cell::Cell;
use std::cell::UnsafeCell;
use std::mem::ManuallyDrop;
use std::pin::Pin;
use std::ptr;
use std::rc::Rc;
use std::task::Context;
use std::task::Poll;
use std::task::RawWaker;
use std::task::RawWakerVTable;
use std::task::Waker;
#[must_use]
pub struct SpawnedFuture<T: 'static> {
vtable: &'static SpawnedFutureVtable<T>,
data: *mut u8,
}
impl<T> Future for SpawnedFuture<T> {
type Output = T;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
unsafe { (self.vtable.poll)(self.data, cx) }
}
}
impl<T> Drop for SpawnedFuture<T> {
fn drop(&mut self) {
unsafe {
(self.vtable.drop)(self.data);
}
}
}
struct SpawnedFutureVTableProxy<T, F>(T, F);
impl<T: 'static, F: Future<Output = T>> SpawnedFutureVTableProxy<T, F> {
const VTABLE: &'static SpawnedFutureVtable<T> = &SpawnedFutureVtable {
poll: Self::poll,
drop: Self::drop,
};
unsafe fn poll(data: *mut u8, ctx: &mut Context<'_>) -> Poll<T> {
unsafe {
let task = (data as *const Task<T, F>).deref();
if &task.state & COMPLETED == 0 {
task.waker.set(Some(ctx.waker().clone()));
Poll::Pending
} else if &task.state & EMPTIED == 0 {
task.state.or_assign(EMPTIED);
Poll::Ready(ptr::read(&*task.data.get().deref().result))
} else {
panic!("Future polled after it has already been emptied");
}
}
}
unsafe fn drop(data: *mut u8) {
unsafe {
{
let task = (data as *const Task<T, F>).deref();
task.state.or_assign(CANCELLED);
if &task.state & RUNNING == 0 {
task.drop_data();
}
}
Task::<T, F>::dec_ref_count(data as _);
}
}
}
struct SpawnedFutureVtable<T> {
poll: unsafe fn(data: *mut u8, ctx: &mut Context<'_>) -> Poll<T>,
drop: unsafe fn(data: *mut u8),
}
union TaskData<T, F: Future<Output = T>> {
result: ManuallyDrop<T>,
future: ManuallyDrop<F>,
}
const RUNNING: u32 = 1;
const RUN_AGAIN: u32 = 2;
const COMPLETED: u32 = 4;
const EMPTIED: u32 = 8;
const CANCELLED: u32 = 16;
struct Task<T, F: Future<Output = T>> {
ref_count: NumCell<u64>,
phase: Phase,
state: NumCell<u32>,
data: UnsafeCell<TaskData<T, F>>,
waker: Cell<Option<Waker>>,
queue: Rc<AsyncEngine>,
#[cfg_attr(not(feature = "tracy"), expect(unused))]
zone: ZoneName,
}
pub(super) struct Runnable {
data: *const u8,
run: unsafe fn(data: *const u8, run: bool),
}
impl Runnable {
pub(super) fn run(self) {
let slf = ManuallyDrop::new(self);
unsafe {
(slf.run)(slf.data, true);
}
}
}
impl Drop for Runnable {
fn drop(&mut self) {
unsafe {
(self.run)(self.data, false);
}
}
}
impl AsyncEngine {
pub(super) fn spawn_<T, F: Future<Output = T>>(
self: &Rc<Self>,
#[cfg_attr(not(feature = "tracy"), expect(unused))] name: &str,
phase: Phase,
f: F,
) -> SpawnedFuture<T> {
let f = Box::new(Task {
ref_count: NumCell::new(1),
phase,
state: NumCell::new(0),
data: UnsafeCell::new(TaskData {
future: ManuallyDrop::new(f),
}),
waker: Cell::new(None),
queue: self.clone(),
zone: create_zone_name!("task:{}", name),
});
unsafe {
f.schedule_run();
}
let f = Box::into_raw(f);
SpawnedFuture {
vtable: SpawnedFutureVTableProxy::<T, F>::VTABLE,
data: f as _,
}
}
}
impl<T, F: Future<Output = T>> Task<T, F> {
const VTABLE: &'static RawWakerVTable = &RawWakerVTable::new(
Self::waker_clone,
Self::waker_wake,
Self::waker_wake_by_ref,
Self::waker_drop,
);
unsafe fn run_proxy(data: *const u8, run: bool) {
unsafe {
let task = data as *const Self;
if run {
task.deref().run();
} else {
Self::task_runnable_dropped(task);
}
Self::dec_ref_count(task);
}
}
#[cold]
unsafe fn task_runnable_dropped(task: *const Self) {
unsafe {
let task = task.deref();
task.state.and_assign(!RUNNING);
if task.state.get() & CANCELLED != 0 {
task.drop_data();
}
}
}
unsafe fn dec_ref_count(slf: *const Self) {
unsafe {
if slf.deref().ref_count.fetch_sub(1) == 1 {
drop(Box::from_raw(slf as *mut Self));
}
}
}
unsafe fn inc_ref_count(&self) {
self.ref_count.fetch_add(1);
}
unsafe fn waker_clone(data: *const ()) -> RawWaker {
unsafe {
let task = (data as *mut Self).deref_mut();
task.inc_ref_count();
RawWaker::new(data, Self::VTABLE)
}
}
unsafe fn waker_wake(data: *const ()) {
unsafe {
Self::waker_wake_by_ref(data);
Self::waker_drop(data);
}
}
unsafe fn waker_wake_by_ref(data: *const ()) {
unsafe {
(data as *const Self).deref().schedule_run();
}
}
unsafe fn waker_drop(data: *const ()) {
unsafe { Self::dec_ref_count(data as _) }
}
unsafe fn schedule_run(&self) {
unsafe {
if &self.state & (COMPLETED | CANCELLED) == 0 {
if &self.state & RUNNING == 0 {
self.state.or_assign(RUNNING);
self.inc_ref_count();
let data = self as *const _ as _;
self.queue.push(
Runnable {
data,
run: Self::run_proxy,
},
self.phase,
);
} else {
self.state.or_assign(RUN_AGAIN);
}
}
}
}
unsafe fn run(&self) {
unsafe {
if &self.state & CANCELLED == 0 {
let data = self.data.get().deref_mut();
self.inc_ref_count();
let raw_waker = RawWaker::new(self as *const _ as _, Self::VTABLE);
let waker = Waker::from_raw(raw_waker);
let mut ctx = Context::from_waker(&waker);
let poll = {
dynamic_zone!(self.zone);
Pin::new_unchecked(&mut *data.future).poll(&mut ctx)
};
if let Poll::Ready(d) = poll {
ManuallyDrop::drop(&mut data.future);
ptr::write(&mut data.result, ManuallyDrop::new(d));
self.state.or_assign(COMPLETED);
if let Some(waker) = self.waker.take() {
waker.wake();
}
}
}
self.state.and_assign(!RUNNING);
if &self.state & CANCELLED != 0 {
self.drop_data();
} else if &self.state & RUN_AGAIN != 0 {
self.state.and_assign(!RUN_AGAIN);
self.schedule_run();
}
}
}
unsafe fn drop_data(&self) {
unsafe {
if &self.state & COMPLETED == 0 {
ManuallyDrop::drop(&mut self.data.get().deref_mut().future);
} else if &self.state & EMPTIED == 0 {
ManuallyDrop::drop(&mut self.data.get().deref_mut().result);
}
}
}
}