Skip to main content

corosensei_awaiter_trait/
lib.rs

1#![no_std]
2
3use core::{
4    future::Future,
5    mem::{MaybeUninit, replace, transmute},
6    pin::Pin,
7    task::{Context, Poll, Waker},
8};
9
10use awaiter_trait::r#dyn::DynAwaiter;
11use corosensei::{Coroutine, Yielder};
12
13use spin::Mutex;
14struct CoroImpl<T, Stack: corosensei::stack::Stack + Unpin> {
15    cor: Coroutine<Waker, (), T, Stack>,
16}
17impl<T, Stack: corosensei::stack::Stack + Unpin> Future for CoroImpl<T, Stack> {
18    type Output = T;
19
20    fn poll(
21        self: core::pin::Pin<&mut Self>,
22        cx: &mut core::task::Context<'_>,
23    ) -> core::task::Poll<Self::Output> {
24        let s = self.get_mut();
25        match s.cor.resume(cx.waker().clone()) {
26            corosensei::CoroutineResult::Yield(_) => Poll::Pending,
27            corosensei::CoroutineResult::Return(a) => Poll::Ready(a),
28        }
29    }
30}
31// #[repr(transparent)]
32struct Awaiter<'a> {
33    y: &'a Yielder<Waker, ()>,
34    w: spin::Mutex<Waker>,
35}
36impl awaiter_trait::Awaiter for Awaiter<'_> {
37    fn r#await<T>(&self, mut f: Pin<&mut (dyn Future<Output = T> + '_)>) -> T {
38        // let mut waker = Waker::noop().clone();
39        loop {
40            let s = self.y.suspend(());
41            let waker = {
42                let mut lock = self.w.lock();
43                replace(&mut *lock, s)
44            };
45            match f.as_mut().poll(&mut Context::from_waker(&waker)) {
46                Poll::Ready(r) => return r,
47                Poll::Pending => {}
48            };
49        }
50    }
51}
52awaiter_trait::autoimpl!(<>Awaiter<'_> as Awaiter);
53pub struct Stacc<'a, Stack: corosensei::stack::Stack + Unpin> {
54    pub via: &'a (dyn Fn() -> Stack + 'a),
55}
56pub trait UPS: corosensei::stack::Stack + Unpin {}
57impl<T: corosensei::stack::Stack + Unpin + ?Sized> UPS for T {}
58awaiter_trait::autoimpl!(<Stack: UPS>Stacc<'_,Stack> as Coroutine);
59impl<Stack: corosensei::stack::Stack + Unpin> awaiter_trait::Coroutine for Stacc<'_, Stack> {
60    fn exec<T>(
61        &self,
62        f: impl FnOnce(&(dyn awaiter_trait::r#dyn::DynAwaiter + '_)) -> T,
63    ) -> impl Future<Output = T> {
64        async move {
65            let mut t = MaybeUninit::uninit();
66            let mut f = match &mut t {
67                mut t => match Some(f) {
68                    mut a => move |b: &(dyn DynAwaiter + '_)| {
69                        t.write(unsafe { a.take().unwrap_unchecked() }(b));
70                    },
71                },
72            };
73            let f: &mut (dyn FnMut(&(dyn awaiter_trait::r#dyn::DynAwaiter + '_)) -> () + '_) =
74                &mut f;
75            let mut f: &mut (dyn FnMut(&(dyn awaiter_trait::r#dyn::DynAwaiter + '_)) -> ()) =
76                unsafe { transmute(f) };
77            CoroImpl {
78                cor: Coroutine::with_stack((self.via)(), move |a, _| {
79                    f(&Awaiter {
80                        y: a,
81                        w: Mutex::new(Waker::noop().clone()),
82                    })
83                }),
84            }
85            .await;
86            unsafe { t.assume_init() }
87        }
88    }
89}