Revert "Let the ? operator work natively in try_stream!." (#55)

This reverts commit 623c556.

Author: taiki-e

async-stream-impl/src/lib.rs

@@ -5,6 +5,8 @@ use syn::parse::Parser;
 use syn::visit_mut::VisitMut;
 
 struct Scrub<'a> {
+    /// Whether the stream is a try stream.
+    is_try: bool,
     /// The unit expression, `()`.
     unit: Box<syn::Expr>,
     has_yielded: bool,
@@ -22,8 +24,9 @@ fn parse_input(input: TokenStream) -> syn::Result<(TokenStream2, Vec<syn::Stmt>)
 }
 
 impl<'a> Scrub<'a> {
-    fn new(crate_path: &'a TokenStream2) -> Self {
+    fn new(is_try: bool, crate_path: &'a TokenStream2) -> Self {
         Self {
+            is_try,
             unit: syn::parse_quote!(()),
             has_yielded: false,
             crate_path,
@@ -41,7 +44,26 @@ impl VisitMut for Scrub<'_> {
 
                 // let ident = &self.yielder;
 
-                *i = syn::parse_quote! { __yield_tx.send(#value_expr).await };
+                *i = if self.is_try {
+                    syn::parse_quote! { __yield_tx.send(::core::result::Result::Ok(#value_expr)).await }
+                } else {
+                    syn::parse_quote! { __yield_tx.send(#value_expr).await }
+                };
+            }
+            syn::Expr::Try(try_expr) => {
+                syn::visit_mut::visit_expr_try_mut(self, try_expr);
+                // let ident = &self.yielder;
+                let e = &try_expr.expr;
+
+                *i = syn::parse_quote! {
+                    match #e {
+                        ::core::result::Result::Ok(v) => v,
+                        ::core::result::Result::Err(e) => {
+                            __yield_tx.send(::core::result::Result::Err(e.into())).await;
+                            return;
+                        }
+                    }
+                };
             }
             syn::Expr::Closure(_) | syn::Expr::Async(_) => {
                 // Don't transform inner closures or async blocks.
@@ -102,7 +124,7 @@ pub fn stream_inner(input: TokenStream) -> TokenStream {
         Err(e) => return e.to_compile_error().into(),
     };
 
-    let mut scrub = Scrub::new(&crate_path);
+    let mut scrub = Scrub::new(false, &crate_path);
 
     for mut stmt in &mut stmts {
         scrub.visit_stmt_mut(&mut stmt);
@@ -136,7 +158,7 @@ pub fn try_stream_inner(input: TokenStream) -> TokenStream {
         Err(e) => return e.to_compile_error().into(),
     };
 
-    let mut scrub = Scrub::new(&crate_path);
+    let mut scrub = Scrub::new(true, &crate_path);
 
     for mut stmt in &mut stmts {
         scrub.visit_stmt_mut(&mut stmt);
@@ -152,13 +174,9 @@ pub fn try_stream_inner(input: TokenStream) -> TokenStream {
 
     quote!({
         let (mut __yield_tx, __yield_rx) = #crate_path::yielder::pair();
-        #crate_path::AsyncTryStream::new(__yield_rx, async move {
+        #crate_path::AsyncStream::new(__yield_rx, async move {
             #dummy_yield
-            let () = {
-                #(#stmts)*
-            };
-            #[allow(unreachable_code)]
-            Ok(())
+            #(#stmts)*
         })
     })
     .into()

async-stream/src/async_stream.rs

@@ -75,70 +75,3 @@ where
         }
     }
 }
-
-#[doc(hidden)]
-#[derive(Debug)]
-pub struct AsyncTryStream<T, U> {
-    rx: Receiver<T>,
-    done: bool,
-    generator: U,
-}
-
-impl<T, U> AsyncTryStream<T, U> {
-    #[doc(hidden)]
-    pub fn new(rx: Receiver<T>, generator: U) -> AsyncTryStream<T, U> {
-        AsyncTryStream {
-            rx,
-            done: false,
-            generator,
-        }
-    }
-}
-
-impl<T, U, E> FusedStream for AsyncTryStream<T, U>
-where
-    U: Future<Output = Result<(), E>>,
-{
-    fn is_terminated(&self) -> bool {
-        self.done
-    }
-}
-
-impl<T, U, E> Stream for AsyncTryStream<T, U>
-where
-    U: Future<Output = Result<(), E>>,
-{
-    type Item = Result<T, E>;
-
-    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        unsafe {
-            let me = Pin::get_unchecked_mut(self);
-
-            if me.done {
-                return Poll::Ready(None);
-            }
-
-            let mut dst = None;
-            let res = {
-                let _enter = me.rx.enter(&mut dst);
-                Pin::new_unchecked(&mut me.generator).poll(cx)
-            };
-
-            me.done = res.is_ready();
-
-            if let Poll::Ready(Err(e)) = res {
-                return Poll::Ready(Some(Err(e)));
-            }
-
-            if let Some(val) = dst.take() {
-                return Poll::Ready(Some(Ok(val)));
-            }
-
-            if me.done {
-                Poll::Ready(None)
-            } else {
-                Poll::Pending
-            }
-        }
-    }
-}

async-stream/src/lib.rs

@@ -164,7 +164,7 @@ pub mod yielder;
 
 // Used by the macro, but not intended to be accessed publicly.
 #[doc(hidden)]
-pub use crate::async_stream::{AsyncStream, AsyncTryStream};
+pub use crate::async_stream::AsyncStream;
 
 #[doc(hidden)]
 pub use async_stream_impl;

async-stream/tests/try_stream.rs

@@ -1,7 +1,6 @@
 use async_stream::try_stream;
 
 use futures_core::stream::Stream;
-use futures_util::pin_mut;
 use futures_util::stream::StreamExt;
 
 #[tokio::test]
@@ -79,22 +78,3 @@ async fn multi_try() {
         values
     );
 }
-
-macro_rules! try_macro {
-    ($e:expr) => {
-        $e?
-    };
-}
-
-#[tokio::test]
-async fn try_in_macro() {
-    let s = try_stream! {
-        yield "hi";
-        try_macro!(Err("bye"));
-    };
-    pin_mut!(s);
-
-    assert_eq!(s.next().await, Some(Ok("hi")));
-    assert_eq!(s.next().await, Some(Err("bye")));
-    assert_eq!(s.next().await, None);
-}