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

generator_light/
gen_context.rs
