Added support for defining macros in Rust within the macro system

Also fixed a lot of bugs

async-fn-stream/src/lib.rs

@@ -1,114 +1,44 @@
-use std::cell::Cell;
-use std::future::poll_fn;
 use std::marker::PhantomData;
-use std::pin::Pin;
-use std::ptr;
-use std::task::{Context, Poll};
 
-use futures::future::LocalBoxFuture;
-use futures::{FutureExt, Stream};
-
-type YieldSlot<'a, T> = &'a Cell<Option<T>>;
+use futures::channel::mpsc;
+use futures::stream::{PollNext, select_with_strategy};
+use futures::{FutureExt, SinkExt, Stream, StreamExt};
 
 /// Handle that allows you to emit values on a stream. If you drop
 /// this, the stream will end and you will not be polled again.
-pub struct StreamCtx<'a, T>(&'a Cell<Option<T>>, PhantomData<&'a ()>);
+pub struct StreamCtx<'a, T>(mpsc::Sender<T>, PhantomData<&'a ()>);
 impl<T> StreamCtx<'_, T> {
-	pub fn emit(&mut self, value: T) -> impl Future<Output = ()> {
-		assert!(self.0.replace(Some(value)).is_none(), "Leftover value in stream");
-		let mut state = Poll::Pending;
-		poll_fn(move |_| std::mem::replace(&mut state, Poll::Ready(())))
+	pub async fn emit(&mut self, value: T) {
+		(self.0.send(value).await)
+			.expect("Dropped a stream receiver without dropping the driving closure");
 	}
 }
 
-enum FnOrFut<'a, T, O> {
-	Fn(Option<Box<dyn FnOnce(YieldSlot<'a, T>) -> LocalBoxFuture<'a, O> + 'a>>),
-	Fut(LocalBoxFuture<'a, O>),
-}
-
-struct AsyncFnStream<'a, T> {
-	driver: FnOrFut<'a, T, ()>,
-	output: Cell<Option<T>>,
-}
-impl<'a, T> Stream for AsyncFnStream<'a, T> {
-	type Item = T;
-	fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-		unsafe {
-			let self_mut = self.get_unchecked_mut();
-			let fut = match &mut self_mut.driver {
-				FnOrFut::Fut(fut) => fut,
-				FnOrFut::Fn(f) => {
-					// safety: the cell is held inline in self, which is pinned.
-					let cell = ptr::from_ref(&self_mut.output).as_ref().unwrap();
-					let fut = f.take().unwrap()(cell);
-					self_mut.driver = FnOrFut::Fut(fut);
-					return Pin::new_unchecked(self_mut).poll_next(cx);
-				},
-			};
-			match fut.as_mut().poll(cx) {
-				Poll::Ready(()) => Poll::Ready(None),
-				Poll::Pending => match self_mut.output.replace(None) {
-					None => Poll::Pending,
-					Some(t) => Poll::Ready(Some(t)),
-				},
-			}
-		}
-	}
-}
-
-struct AsyncFnTryStream<'a, T, E> {
-	driver: FnOrFut<'a, T, Result<StreamCtx<'a, T>, E>>,
-	output: Cell<Option<T>>,
-}
-impl<'a, T, E> Stream for AsyncFnTryStream<'a, T, E> {
-	type Item = Result<T, E>;
-	fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-		unsafe {
-			let self_mut = self.get_unchecked_mut();
-			let fut = match &mut self_mut.driver {
-				FnOrFut::Fut(fut) => fut,
-				FnOrFut::Fn(f) => {
-					// safety: the cell is held inline in self, which is pinned.
-					let cell = ptr::from_ref(&self_mut.output).as_ref().unwrap();
-					let fut = f.take().unwrap()(cell);
-					self_mut.driver = FnOrFut::Fut(fut);
-					return Pin::new_unchecked(self_mut).poll_next(cx);
-				},
-			};
-			match fut.as_mut().poll(cx) {
-				Poll::Ready(Ok(_)) => Poll::Ready(None),
-				Poll::Ready(Err(ex)) => Poll::Ready(Some(Err(ex))),
-				Poll::Pending => match self_mut.output.replace(None) {
-					None => Poll::Pending,
-					Some(t) => Poll::Ready(Some(Ok(t))),
-				},
-			}
-		}
-	}
-}
+fn left_strat(_: &mut ()) -> PollNext { PollNext::Left }
 
 /// Create a stream from an async function acting as a coroutine
 pub fn stream<'a, T: 'a>(
 	f: impl for<'b> AsyncFnOnce(StreamCtx<'b, T>) + 'a,
 ) -> impl Stream<Item = T> + 'a {
-	AsyncFnStream {
-		output: Cell::new(None),
-		driver: FnOrFut::Fn(Some(Box::new(|t| {
-			async { f(StreamCtx(t, PhantomData)).await }.boxed_local()
-		}))),
-	}
+	let (send, recv) = mpsc::channel::<T>(1);
+	let fut = async { f(StreamCtx(send, PhantomData)).await };
+	// use options to ensure that the stream is driven to exhaustion
+	select_with_strategy(fut.into_stream().map(|()| None), recv.map(|t| Some(t)), left_strat)
+		.filter_map(async |opt| opt)
 }
 
 /// Create a stream of result from a fallible function.
 pub fn try_stream<'a, T: 'a, E: 'a>(
 	f: impl for<'b> AsyncFnOnce(StreamCtx<'b, T>) -> Result<StreamCtx<'b, T>, E> + 'a,
 ) -> impl Stream<Item = Result<T, E>> + 'a {
-	AsyncFnTryStream {
-		output: Cell::new(None),
-		driver: FnOrFut::Fn(Some(Box::new(|t| {
-			async { f(StreamCtx(t, PhantomData)).await }.boxed_local()
-		}))),
-	}
+	let (send, recv) = mpsc::channel::<T>(1);
+	let fut = async { f(StreamCtx(send, PhantomData)).await };
+	select_with_strategy(
+		fut.into_stream().map(|res| if let Err(e) = res { Some(Err(e)) } else { None }),
+		recv.map(|t| Some(Ok(t))),
+		left_strat,
+	)
+	.filter_map(async |opt| opt)
 }
 
 #[cfg(test)]