Add for_await to process streams using a for loop (#2)

Author: taiki-e

async-stream-impl/src/lib.rs

@@ -1,7 +1,7 @@
 extern crate proc_macro;
 
 use proc_macro::{TokenStream, TokenTree};
-use proc_macro2::Span;
+use proc_macro2::{Group, Span, TokenStream as TokenStream2, TokenTree as TokenTree2};
 use quote::quote;
 use syn::visit_mut::VisitMut;
 
@@ -19,7 +19,7 @@ struct AsyncStreamEnumHack {
 }
 
 impl AsyncStreamEnumHack {
-    fn parse(input: TokenStream) -> Self {
+    fn parse(input: TokenStream) -> syn::Result<Self> {
         macro_rules! n {
             ($i:ident) => {
                 $i.next().unwrap()
@@ -44,14 +44,15 @@ impl AsyncStreamEnumHack {
         n!(braces); // !
 
         let inner = n!(braces);
-        let syn::Block { stmts, .. } = syn::parse(inner.clone().into()).unwrap();
+        let inner = replace_for_await(TokenStream2::from(TokenStream::from(inner)));
+        let syn::Block { stmts, .. } = syn::parse2(inner.clone())?;
 
         let macro_ident = syn::Ident::new(
             &format!("stream_{}", count_bangs(inner.into())),
             Span::call_site(),
         );
 
-        AsyncStreamEnumHack { stmts, macro_ident }
+        Ok(AsyncStreamEnumHack { stmts, macro_ident })
     }
 }
 
@@ -100,6 +101,42 @@ impl VisitMut for Scrub {
                 syn::visit_mut::visit_expr_mut(self, i);
                 self.is_xforming = prev;
             }
+            syn::Expr::ForLoop(expr) => {
+                syn::visit_mut::visit_expr_for_loop_mut(self, expr);
+                // TODO: Should we allow other attributes?
+                if expr.attrs.len() != 1 || !expr.attrs[0].path.is_ident("await") {
+                    return;
+                }
+                let syn::ExprForLoop {
+                    attrs,
+                    label,
+                    pat,
+                    expr,
+                    body,
+                    ..
+                } = expr;
+
+                let attr = attrs.pop().unwrap();
+                if let Err(e) = syn::parse2::<syn::parse::Nothing>(attr.tokens) {
+                    *i = syn::parse2(e.to_compile_error()).unwrap();
+                    return;
+                }
+
+                *i = syn::parse_quote! {{
+                    let mut __pinned = #expr;
+                    let mut __pinned = unsafe {
+                        ::async_stream::reexport::Pin::new_unchecked(&mut __pinned)
+                    };
+                    #label
+                    loop {
+                        let #pat = match ::async_stream::reexport::next(&mut __pinned).await {
+                            ::async_stream::reexport::Some(e) => e,
+                            ::async_stream::reexport::None => break,
+                        };
+                        #body
+                    }
+                }}
+            }
             _ => syn::visit_mut::visit_expr_mut(self, i),
         }
     }
@@ -117,7 +154,10 @@ pub fn async_stream_impl(input: TokenStream) -> TokenStream {
     let AsyncStreamEnumHack {
         macro_ident,
         mut stmts,
-    } = AsyncStreamEnumHack::parse(input);
+    } = match AsyncStreamEnumHack::parse(input) {
+        Ok(x) => x,
+        Err(e) => return e.to_compile_error().into(),
+    };
 
     let mut scrub = Scrub {
         is_xforming: true,
@@ -156,7 +196,10 @@ pub fn async_try_stream_impl(input: TokenStream) -> TokenStream {
     let AsyncStreamEnumHack {
         macro_ident,
         mut stmts,
-    } = AsyncStreamEnumHack::parse(input);
+    } = match AsyncStreamEnumHack::parse(input) {
+        Ok(x) => x,
+        Err(e) => return e.to_compile_error().into(),
+    };
 
     let mut scrub = Scrub {
         is_xforming: true,
@@ -209,3 +252,30 @@ fn count_bangs(input: TokenStream) -> usize {
 
     count
 }
+
+fn replace_for_await(input: TokenStream2) -> TokenStream2 {
+    let mut input = input.into_iter().peekable();
+    let mut tokens = Vec::new();
+
+    while let Some(token) = input.next() {
+        match token {
+            TokenTree2::Ident(ident) => {
+                match input.peek() {
+                    Some(TokenTree2::Ident(next)) if ident == "for" && next == "await" => {
+                        tokens.extend(quote!(#[#next]));
+                        let _ = input.next();
+                    }
+                    _ => {}
+                }
+                tokens.push(ident.into());
+            }
+            TokenTree2::Group(group) => {
+                let stream = replace_for_await(group.stream());
+                tokens.push(Group::new(group.delimiter(), stream).into());
+            }
+            _ => tokens.push(token),
+        }
+    }
+
+    tokens.into_iter().collect()
+}

async-stream/Cargo.toml

@@ -20,7 +20,7 @@ repository = "https://github.com/tokio-rs/async-stream"
 readme = "README.md"
 
 [dependencies]
-async-stream-impl = { version = "0.1.0" }
+async-stream-impl = { version = "0.1.0", path = "../async-stream-impl" }
 futures-core-preview = "=0.3.0-alpha.18"
 
 [dev-dependencies]

async-stream/src/lib.rs

@@ -169,6 +169,7 @@
 //! [`proc-macro-hack`]: https://github.com/dtolnay/proc-macro-hack/
 
 mod async_stream;
+mod next;
 #[doc(hidden)]
 pub mod yielder;
 
@@ -179,6 +180,16 @@ pub use crate::async_stream::AsyncStream;
 #[doc(hidden)]
 pub use async_stream_impl::{AsyncStreamHack, AsyncTryStreamHack};
 
+#[doc(hidden)]
+pub mod reexport {
+    #[doc(hidden)]
+    pub use crate::next::next;
+    #[doc(hidden)]
+    pub use std::option::Option::{None, Some};
+    #[doc(hidden)]
+    pub use std::pin::Pin;
+}
+
 /// Asynchronous stream
 ///
 /// See [crate](index.html) documentation for more details.

async-stream/src/next.rs

@@ -0,0 +1,32 @@
+use futures_core::Stream;
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+// This is equivalent to the `futures::StreamExt::next` method.
+// But we want to make this crate dependency as small as possible, so we define our `next` function.
+#[doc(hidden)]
+pub fn next<S>(stream: &mut S) -> impl Future<Output = Option<S::Item>> + '_
+where
+    S: Stream + Unpin,
+{
+    Next { stream }
+}
+
+#[derive(Debug)]
+struct Next<'a, S> {
+    stream: &'a mut S,
+}
+
+impl<S> Unpin for Next<'_, S> where S: Unpin {}
+
+impl<S> Future for Next<'_, S>
+where
+    S: Stream + Unpin,
+{
+    type Output = Option<S::Item>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        Pin::new(&mut self.stream).poll_next(cx)
+    }
+}

async-stream/tests/for_await.rs

@@ -0,0 +1,25 @@
+#![feature(async_await)]
+
+use tokio::prelude::*;
+
+use async_stream::stream;
+
+#[tokio::test]
+async fn test() {
+    let s = stream! {
+        yield "hello";
+        yield "world";
+    };
+
+    let s = stream! {
+        for await x in s {
+            yield x.to_owned() + "!";
+        }
+    };
+
+    let values: Vec<_> = s.collect().await;
+
+    assert_eq!(2, values.len());
+    assert_eq!("hello!", values[0]);
+    assert_eq!("world!", values[1]);
+}