Expand 'yield' in internal macro calls. (#57)

Author: SergioBenitez

async-stream-impl/src/lib.rs

@@ -1,7 +1,7 @@
 use proc_macro::TokenStream;
 use proc_macro2::{Group, TokenStream as TokenStream2, TokenTree};
 use quote::quote;
-use syn::parse::Parser;
+use syn::parse::{Parse, ParseStream, Parser, Result};
 use syn::visit_mut::VisitMut;
 
 struct Scrub<'a> {
@@ -34,6 +34,80 @@ impl<'a> Scrub<'a> {
     }
 }
 
+struct Partial<T>(T, TokenStream2);
+
+impl<T: Parse> Parse for Partial<T> {
+    fn parse(input: ParseStream) -> Result<Self> {
+        Ok(Partial(input.parse()?, input.parse()?))
+    }
+}
+
+fn visit_token_stream_impl(
+    visitor: &mut Scrub<'_>,
+    tokens: TokenStream2,
+    modified: &mut bool,
+    out: &mut TokenStream2,
+) {
+    use quote::ToTokens;
+    use quote::TokenStreamExt;
+
+    let mut tokens = tokens.into_iter().peekable();
+    while let Some(tt) = tokens.next() {
+        match tt {
+            TokenTree::Ident(i) if i == "yield" => {
+                let stream = std::iter::once(TokenTree::Ident(i)).chain(tokens).collect();
+                match syn::parse2(stream) {
+                    Ok(Partial(yield_expr, rest)) => {
+                        let mut expr = syn::Expr::Yield(yield_expr);
+                        visitor.visit_expr_mut(&mut expr);
+                        expr.to_tokens(out);
+                        *modified = true;
+                        tokens = rest.into_iter().peekable();
+                    }
+                    Err(e) => {
+                        out.append_all(&mut e.to_compile_error().into_iter());
+                        *modified = true;
+                        return;
+                    }
+                }
+            }
+            TokenTree::Ident(i) if i == "stream" || i == "try_stream" => {
+                out.append(TokenTree::Ident(i));
+                match tokens.peek() {
+                    Some(TokenTree::Punct(p)) if p.as_char() == '!' => {
+                        out.extend(tokens.next()); // !
+                        if let Some(TokenTree::Group(_)) = tokens.peek() {
+                            out.extend(tokens.next()); // { .. } or [ .. ] or ( .. )
+                        }
+                    }
+                    _ => {}
+                }
+            }
+            TokenTree::Group(group) => {
+                let mut content = group.stream();
+                *modified |= visitor.visit_token_stream(&mut content);
+                let mut new = Group::new(group.delimiter(), content);
+                new.set_span(group.span());
+                out.append(new);
+            }
+            other => out.append(other),
+        }
+    }
+}
+
+impl Scrub<'_> {
+    fn visit_token_stream(&mut self, tokens: &mut TokenStream2) -> bool {
+        let (mut out, mut modified) = (TokenStream2::new(), false);
+        visit_token_stream_impl(self, tokens.clone(), &mut modified, &mut out);
+
+        if modified {
+            *tokens = out;
+        }
+
+        modified
+    }
+}
+
 impl VisitMut for Scrub<'_> {
     fn visit_expr_mut(&mut self, i: &mut syn::Expr) {
         match i {
@@ -109,8 +183,20 @@ impl VisitMut for Scrub<'_> {
         }
     }
 
-    fn visit_item_mut(&mut self, _: &mut syn::Item) {
-        // Don't transform inner items.
+    fn visit_macro_mut(&mut self, mac: &mut syn::Macro) {
+        let mac_ident = mac.path.segments.last().map(|p| &p.ident);
+        if mac_ident.map_or(false, |i| i == "stream" || i == "try_stream") {
+            return;
+        }
+
+        self.visit_token_stream(&mut mac.tokens);
+    }
+
+    fn visit_item_mut(&mut self, i: &mut syn::Item) {
+        // Recurse into macros but otherwise don't transform inner items.
+        if let syn::Item::Macro(i) = i {
+            self.visit_macro_mut(&mut i.mac);
+        }
     }
 }
 

async-stream/tests/stream.rs

@@ -80,6 +80,42 @@ async fn yield_multi_value() {
     assert_eq!("dizzy", values[2]);
 }
 
+#[tokio::test]
+async fn unit_yield_in_select() {
+    use tokio::select;
+
+    async fn do_stuff_async() {}
+
+    let s = stream! {
+        select! {
+            _ = do_stuff_async() => yield,
+            else => yield,
+        }
+    };
+
+    let values: Vec<_> = s.collect().await;
+    assert_eq!(values.len(), 1);
+}
+
+#[tokio::test]
+async fn yield_with_select() {
+    use tokio::select;
+
+    async fn do_stuff_async() {}
+    async fn more_async_work() {}
+
+    let s = stream! {
+        select! {
+            _ = do_stuff_async() => yield "hey",
+            _ = more_async_work() => yield "hey",
+            else => yield "hey",
+        }
+    };
+
+    let values: Vec<_> = s.collect().await;
+    assert_eq!(values, vec!["hey"]);
+}
+
 #[tokio::test]
 async fn return_stream() {
     fn build_stream() -> impl Stream<Item = u32> {
@@ -172,6 +208,27 @@ async fn yield_non_unpin_value() {
     assert_eq!(s, vec![0, 1, 2]);
 }
 
+#[test]
+fn inner_try_stream() {
+    use async_stream::try_stream;
+    use tokio::select;
+
+    async fn do_stuff_async() {}
+
+    let _ = stream! {
+        select! {
+            _ = do_stuff_async() => {
+                let another_s = try_stream! {
+                    yield;
+                };
+                let _: Result<(), ()> = Box::pin(another_s).next().await.unwrap();
+            },
+            else => {},
+        }
+        yield
+    };
+}
+
 #[test]
 fn test() {
     let t = trybuild::TestCases::new();

async-stream/tests/ui/yield_bad_expr_in_macro.rs

@@ -0,0 +1,11 @@
+use async_stream::stream;
+
+fn main() {
+    async fn work() {}
+
+    stream! {
+        tokio::select! {
+            _ = work() => yield fn f() {},
+        }
+    };
+}

async-stream/tests/ui/yield_bad_expr_in_macro.stderr

@@ -0,0 +1,5 @@
+error: expected expression
+ --> $DIR/yield_bad_expr_in_macro.rs:8:33
+  |
+8 |             _ = work() => yield fn f() {},
+  |                                 ^^