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

generator_light/
gen_context.rs
