generator_light/
gen_context.rs

1use core::cell::Cell;
2use core::marker::PhantomData;
3use core::pin::Pin;
4use core::task::{Poll, RawWaker, RawWakerVTable, Waker};
5
6use core::ptr::NonNull;
7
8enum YieldState<Y, R> {
9    Resume(R),
10    Yield(Option<Y>),
11}
12
13pub(crate) struct GeneratorContext<Y, R>(Cell<YieldState<Y, R>>);
14
15// pub struct Yielder<'a, Y, R> {
16//     state: NonNull<Cell<YieldState<Y, R>>>,
17//     brand: PhantomData<&'a mut ()>,
18// }
19
20pub struct Yielder<'a, Y, R> {
21    state: PhantomData<(Y, R)>,
22    brand: PhantomData<&'a mut ()>,
23}
24
25impl<Y, R> GeneratorContext<Y, R> {
26    pub(crate) const fn new_resumed(value: R) -> Self {
27        Self(Cell::new(YieldState::Resume(value)))
28    }
29
30    pub(crate) fn take_yielded(self) -> Option<Y> {
31        let YieldState::Yield(val) = self.0.into_inner() else {
32            unsafe { crate::core::hint::unreachable_unchecked() };
33        };
34        val
35    }
36}
37
38unsafe fn get_context<Y, R>(ctx: &mut core::task::Context<'_>) -> NonNull<Cell<YieldState<Y, R>>> {
39    unsafe { NonNull::new_unchecked(ctx.waker().data() as *mut ()) }.cast()
40}
41
42const YIELD_WAKER: RawWakerVTable = RawWakerVTable::new(
43    |_| panic!("Clone is not allowed for yielder"),
44    |_| {},
45    |_| {},
46    |_| {},
47);
48
49pub(crate) const unsafe fn make_yielder_waker<Y, R>(x: &GeneratorContext<Y, R>) -> Waker {
50    unsafe { Waker::from_raw(RawWaker::new((&raw const x.0).cast(), &YIELD_WAKER)) }
51}
52
53#[repr(transparent)]
54struct YieldFuture<'a, Y, R> {
55    state: Option<Y>,
56    _yileder: Yielder<'a, Y, R>,
57}
58
59impl<'a, Y, R> Future for YieldFuture<'a, Y, R> {
60    type Output = R;
61
62    #[inline(always)]
63    fn poll(self: Pin<&mut Self>, cx: &mut core::task::Context<'_>) -> Poll<Self::Output> {
64        // Safety:
65        //    YieldFuture may be created only by Yielder. Yeilder exists only during generator
66        //    execution. This means we are inside generator and cx has yilder_waker data inside
67        match unsafe {
68            get_context(cx)
69                .as_ref()
70                .replace(YieldState::Yield(self.get_unchecked_mut().state.take()))
71        } {
72            YieldState::Resume(r) => Poll::Ready(r),
73            _ => Poll::Pending,
74        }
75    }
76}
77
78impl<Y, R> Yielder<'_, Y, R> {
79    pub(crate) const fn new() -> Self {
80        Yielder {
81            state: PhantomData,
82            brand: PhantomData,
83        }
84    }
85
86    pub const fn yield_value<'a>(&'a mut self, value: Y) -> impl Future<Output = R> {
87        YieldFuture {
88            state: Some(value),
89            _yileder: Yielder::<'a, Y, R>::new(),
90        }
91    }
92
93    pub const fn suspend<'a>(&'a mut self) -> impl Future<Output = R> {
94        YieldFuture {
95            state: None,
96            _yileder: Yielder::<'a, Y, R>::new(),
97        }
98    }
99}
100
101#[macro_export]
102macro_rules! yield_ {
103    ($yielder:ident, $value:expr) => {
104        $yielder.yield_value($value).await
105    };
106}
107
108#[macro_export]
109macro_rules! suspend_ {
110    ($yielder:ident) => {
111        $yielder.suspend().await
112    };
113}
generator_light/gen_context.rs

generator_light/
gen_context.rs
