for await

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 })
     }
 }
 
@@ -153,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,
@@ -192,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,
@@ -245,3 +252,32 @@ 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.to_string().as_str() == "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/tests/for_await.rs

@@ -12,8 +12,7 @@ async fn test() {
     };
 
     let s = stream! {
-        #[await]
-        for x in s {
+        for await x in s {
             yield x.to_owned() + "!";
         }
     };