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

generator_light/
gen_context.rs
