handle edge cases

Author: carllerche

async-stream-impl/src/lib.rs

@@ -13,6 +13,7 @@ struct AsyncStreamImpl {
 }
 
 struct Scrub {
+    is_xforming: bool,
     is_try: bool,
     yielder: syn::Ident,
     unit: Box<syn::Expr>,
@@ -39,6 +40,11 @@ impl Parse for AsyncStreamImpl {
 
 impl VisitMut for Scrub {
     fn visit_expr_mut(&mut self, i: &mut syn::Expr) {
+        if !self.is_xforming {
+            syn::visit_mut::visit_expr_mut(self, i);
+            return;
+        }
+
         match i {
             syn::Expr::Yield(yield_expr) => {
                 self.num_yield += 1;
@@ -71,9 +77,23 @@ impl VisitMut for Scrub {
                     }
                 };
             }
-            expr => syn::visit_mut::visit_expr_mut(self, expr),
+            syn::Expr::Closure(_) |
+                syn::Expr::Async(_) => {
+                let prev = self.is_xforming;
+                self.is_xforming = false;
+                syn::visit_mut::visit_expr_mut(self, i);
+                self.is_xforming = prev;
+            }
+            _ => syn::visit_mut::visit_expr_mut(self, i),
         }
     }
+
+    fn visit_item_mut(&mut self, i: &mut syn::Item) {
+        let prev = self.is_xforming;
+        self.is_xforming = false;
+        syn::visit_mut::visit_item_mut(self, i);
+        self.is_xforming = prev;
+    }
 }
 
 #[proc_macro_hack]
@@ -84,6 +104,7 @@ pub fn async_stream_impl(input: TokenStream) -> TokenStream {
     } = syn::parse_macro_input!(input as AsyncStreamImpl);
 
     let mut scrub = Scrub {
+        is_xforming: true,
         is_try: false,
         yielder,
         unit: syn::parse_quote!(()),
@@ -119,6 +140,7 @@ pub fn async_try_stream_impl(input: TokenStream) -> TokenStream {
     } = syn::parse_macro_input!(input as AsyncStreamImpl);
 
     let mut scrub = Scrub {
+        is_xforming: true,
         is_try: true,
         yielder,
         unit: syn::parse_quote!(()),

async-stream/Cargo.toml

@@ -13,3 +13,4 @@ proc-macro-hack = "0.5.8"
 tokio = "=0.2.0-alpha.1"
 tokio-test = "=0.2.0-alpha.1"
 futures-util-preview = "=0.3.0-alpha.17"
+trybuild = "1.0"

async-stream/tests/stream.rs

@@ -123,3 +123,28 @@ async fn borrow_self() {
 
     assert_eq!(Some("hello"), s.next().await);
 }
+
+#[tokio::test]
+async fn stream_in_stream() {
+    let s = stream! {
+        let s = stream! {
+            for i in 0..3 {
+                yield i;
+            }
+        };
+
+        pin_mut!(s);
+        while let Some(v) = s.next().await {
+            yield v;
+        }
+    };
+
+    let values: Vec<_> = s.collect().await;
+    assert_eq!(3, values.len());
+}
+
+#[test]
+fn test() {
+    let t = trybuild::TestCases::new();
+    t.compile_fail("tests/ui/*.rs");
+}

async-stream/tests/ui/yield_in_async.rs

@@ -0,0 +1,13 @@
+#![feature(async_await)]
+
+use async_stream::stream;
+
+fn main() {
+    stream! {
+        let f = async {
+            yield 123;
+        };
+
+        let v = f.await;
+    };
+}

async-stream/tests/ui/yield_in_async.stderr

@@ -0,0 +1,16 @@
+error[E0658]: yield syntax is experimental
+ --> $DIR/yield_in_async.rs:8:13
+  |
+8 |             yield 123;
+  |             ^^^^^^^^^
+  |
+  = note: for more information, see https://github.com/rust-lang/rust/issues/43122
+  = help: add `#![feature(generators)]` to the crate attributes to enable
+
+error[E0727]: `async` generators are not yet supported
+ --> $DIR/yield_in_async.rs:8:13
+  |
+8 |             yield 123;
+  |             ^^^^^^^^^
+
+For more information about this error, try `rustc --explain E0658`.

async-stream/tests/ui/yield_in_closure.rs

@@ -0,0 +1,13 @@
+#![feature(async_await)]
+
+use async_stream::stream;
+
+fn main() {
+    stream! {
+        Ok("value")
+            .and_then(|v| {
+                yield v;
+                Ok(())
+            });
+    };
+}

async-stream/tests/ui/yield_in_closure.stderr

@@ -0,0 +1,16 @@
+error[E0658]: yield syntax is experimental
+ --> $DIR/yield_in_closure.rs:9:17
+  |
+9 |                 yield v;
+  |                 ^^^^^^^
+  |
+  = note: for more information, see https://github.com/rust-lang/rust/issues/43122
+  = help: add `#![feature(generators)]` to the crate attributes to enable
+
+error[E0628]: generators cannot have explicit arguments
+ --> $DIR/yield_in_closure.rs:8:23
+  |
+8 |             .and_then(|v| {
+  |                       ^^^
+
+For more information about this error, try `rustc --explain E0658`.

async-stream/tests/ui/yield_in_nested_fn.rs

@@ -0,0 +1,11 @@
+#![feature(async_await)]
+
+use async_stream::stream;
+
+fn main() {
+    stream! {
+        fn foo() {
+            yield "hello";
+        }
+    };
+}

async-stream/tests/ui/yield_in_nested_fn.stderr

@@ -0,0 +1,16 @@
+error[E0658]: yield syntax is experimental
+ --> $DIR/yield_in_nested_fn.rs:8:13
+  |
+8 |             yield "hello";
+  |             ^^^^^^^^^^^^^
+  |
+  = note: for more information, see https://github.com/rust-lang/rust/issues/43122
+  = help: add `#![feature(generators)]` to the crate attributes to enable
+
+error[E0627]: yield statement outside of generator literal
+ --> $DIR/yield_in_nested_fn.rs:8:13
+  |
+8 |             yield "hello";
+  |             ^^^^^^^^^^^^^
+
+For more information about this error, try `rustc --explain E0658`.