Prepare for landing await syntax

futures-async-macro/README.md

@@ -17,7 +17,7 @@ Processes streams using a for loop.
 This is a reimplement of [futures-await]'s `#[async]` for loops for futures 0.3 and is an experimental implementation of [the idea listed as the next step of async/await](https://github.com/rust-lang/rfcs/blob/master/text/2394-async_await.md#for-await-and-processing-streams).
 
 ```rust
-#![feature(async_await, generators, stmt_expr_attributes, proc_macro_hygiene)]
+#![feature(async_await, await_macro, stmt_expr_attributes, proc_macro_hygiene)]
 use futures::for_await;
 use futures::prelude::*;
 

futures-async-macro/src/lib.rs

@@ -241,42 +241,61 @@ impl Expand {
             ));
         }
 
+        let ExprForLoop { label, pat, expr, body, .. } = &expr;
         // It needs to adjust the type yielded by the macro because generators used internally by
         // async fn yield `()` type, but generators used internally by `async_stream` yield
         // `Poll<U>` type.
-        let yield_ = match self.0 {
-            Future => TokenStream2::new(),
-            Stream => quote! { ::futures::core_reexport::task::Poll::Pending },
-            Closure => return outside_of_async_error!(expr, "#[for_await]"),
-        };
-        let ExprForLoop { label, pat, expr, body, .. } = expr;
-
-        // Basically just expand to a `poll` loop
-        syn::parse_quote! {{
-            let mut __pinned = #expr;
-            #label
-            loop {
-                let #pat = {
-                    match ::futures::async_stream::poll_next_with_tls_context(unsafe {
-                            ::futures::core_reexport::pin::Pin::new_unchecked(&mut __pinned)
-                        })
-                    {
-                        ::futures::core_reexport::task::Poll::Ready(e) => {
-                            match e {
+        match self.0 {
+            Future => {
+                // Basically just expand to a `poll` loop
+                syn::parse_quote! {{
+                    let mut __pinned = #expr;
+                    let mut __pinned = unsafe {
+                        ::futures::core_reexport::pin::Pin::new_unchecked(&mut __pinned)
+                    };
+                    #label
+                    loop {
+                        let #pat = {
+                            match r#await!(::futures::stream::StreamExt::next(&mut __pinned)) {
                                 ::futures::core_reexport::option::Option::Some(e) => e,
                                 ::futures::core_reexport::option::Option::None => break,
                             }
-                        }
-                        ::futures::core_reexport::task::Poll::Pending => {
-                            yield #yield_;
-                            continue
-                        }
+                        };
+
+                        #body
                     }
-                };
+                }}
+            }
+            Stream => {
+                // Basically just expand to a `poll` loop
+                syn::parse_quote! {{
+                    let mut __pinned = #expr;
+                    #label
+                    loop {
+                        let #pat = {
+                            match ::futures::async_stream::poll_next_with_tls_context(unsafe {
+                                    ::futures::core_reexport::pin::Pin::new_unchecked(&mut __pinned)
+                                })
+                            {
+                                ::futures::core_reexport::task::Poll::Ready(e) => {
+                                    match e {
+                                        ::futures::core_reexport::option::Option::Some(e) => e,
+                                        ::futures::core_reexport::option::Option::None => break,
+                                    }
+                                }
+                                ::futures::core_reexport::task::Poll::Pending => {
+                                    yield ::futures::core_reexport::task::Poll::Pending;
+                                    continue
+                                }
+                            }
+                        };
 
-                #body
+                        #body
+                    }
+                }}
             }
-        }}
+            Closure => return outside_of_async_error!(expr, "#[for_await]"),
+        }
     }
 
     /// Expands `yield expr` in `async_stream` scope.

futures/testcrate/ui/nested.rs

@@ -1,4 +1,4 @@
-#![feature(async_await, generators)]
+#![feature(async_await, generators, await_macro)]
 
 use futures::*;