use crate::worker;
use crossbeam_utils::sync::Parker;
use lightproc::proc_stack::ProcStack;
use std::cell::{Cell, UnsafeCell};
use std::future::Future;
use std::mem;
use std::mem::ManuallyDrop;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};
pub fn run<F, T>(future: F, stack: ProcStack) -> T
where
F: Future<Output = T>,
{
unsafe {
let out = &mut UnsafeCell::new(None);
let future = {
let out = out.get();
async move {
*out = Some(future.await);
}
};
pin_utils::pin_mut!(future);
let future = mem::transmute::<
Pin<&'_ mut dyn Future<Output = ()>>,
Pin<&'static mut dyn Future<Output = ()>>,
>(future);
worker::set_stack(&stack, || block(future));
match (*out.get()).take() {
Some(v) => v,
_ => unimplemented!(),
}
}
}
fn block<F, T>(f: F) -> T
where
F: Future<Output = T>,
{
thread_local! {
static CACHE: Cell<Option<Arc<Parker>>> = Cell::new(None);
}
pin_utils::pin_mut!(f);
CACHE.with(|cache| {
let arc_parker: Arc<Parker> = cache.take().unwrap_or_else(|| Arc::new(Parker::new()));
let ptr = (&*arc_parker as *const Parker) as *const ();
let vt = vtable();
let waker = unsafe { ManuallyDrop::new(Waker::from_raw(RawWaker::new(ptr, vt))) };
let cx = &mut Context::from_waker(&waker);
loop {
if let Poll::Ready(t) = f.as_mut().poll(cx) {
cache.set(Some(arc_parker));
return t;
}
arc_parker.park();
}
})
}
fn vtable() -> &'static RawWakerVTable {
unsafe fn clone_raw(ptr: *const ()) -> RawWaker {
#![allow(clippy::redundant_clone)]
let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const Parker));
mem::forget(arc.clone());
RawWaker::new(ptr, vtable())
}
unsafe fn wake_raw(ptr: *const ()) {
let arc = Arc::from_raw(ptr as *const Parker);
arc.unparker().unpark();
}
unsafe fn wake_by_ref_raw(ptr: *const ()) {
let arc = ManuallyDrop::new(Arc::from_raw(ptr as *const Parker));
arc.unparker().unpark();
}
unsafe fn drop_raw(ptr: *const ()) {
drop(Arc::from_raw(ptr as *const Parker))
}
&RawWakerVTable::new(clone_raw, wake_raw, wake_by_ref_raw, drop_raw)
}