use std::{
async_iter::AsyncIterator,
future::Future,
ops::{Generator, GeneratorState},
pin::Pin,
ptr::NonNull,
task::{Context, Poll},
};
use crate::GeneratorImpl;
#[doc(hidden)]
pub struct UnsafeContextRef(NonNull<Context<'static>>);
impl UnsafeContextRef {
#[doc(hidden)]
pub unsafe fn get_context<'a, 'b>(&mut self) -> &'a mut Context<'b> {
std::mem::transmute(self.0)
}
}
impl<'a> From<&mut Context<'a>> for UnsafeContextRef {
fn from(cx: &mut Context<'a>) -> Self {
Self(unsafe { std::mem::transmute(cx) })
}
}
unsafe impl Send for UnsafeContextRef {}
impl<G> GeneratorImpl<G> {
#[doc(hidden)]
pub unsafe fn new_async<Y, R>(generator: G) -> impl AsyncGenerator<Y, R>
where
G: Generator<UnsafeContextRef, Yield = Poll<Y>, Return = R>,
{
Self { generator }
}
}
pub trait AsyncGenerator<Y, R>: AsyncIterator<Item = Y> + Future<Output = R> {
fn poll_resume(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<GeneratorState<Y, R>>;
#[doc(hidden)]
fn into_async_generator(self) -> Self
where
Self: Sized,
{
self
}
}
impl<Y, G> AsyncIterator for GeneratorImpl<G>
where
G: Generator<UnsafeContextRef, Yield = Poll<Y>>,
{
type Item = Y;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
match self.project_generator().resume(cx.into()) {
GeneratorState::Yielded(p) => p.map(Some),
GeneratorState::Complete(_) => Poll::Ready(None),
}
}
}
impl<R, G> Future for GeneratorImpl<G>
where
G: Generator<UnsafeContextRef, Return = R>,
{
type Output = R;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.project_generator().resume(cx.into()) {
GeneratorState::Yielded(_) => Poll::Pending,
GeneratorState::Complete(r) => Poll::Ready(r),
}
}
}
impl<Y, R, G> AsyncGenerator<Y, R> for GeneratorImpl<G>
where
G: Generator<UnsafeContextRef, Yield = Poll<Y>, Return = R>,
{
fn poll_resume(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<GeneratorState<Y, G::Return>> {
match self.project_generator().resume(cx.into()) {
GeneratorState::Yielded(p) => p.map(GeneratorState::Yielded),
GeneratorState::Complete(r) => Poll::Ready(GeneratorState::Complete(r)),
}
}
}