22
33#![ recursion_limit = "128" ]
44#![ warn( rust_2018_idioms) ]
5+ #![ warn( clippy:: all) ]
56
67extern crate proc_macro;
78
@@ -10,8 +11,8 @@ use proc_macro2::{Span, TokenStream as TokenStream2, TokenTree as TokenTree2};
1011use quote:: { quote, ToTokens } ;
1112use syn:: {
1213 fold:: { self , Fold } ,
13- token, ArgCaptured , Error , Expr , ExprCall , ExprForLoop , ExprMacro , ExprYield , FnArg , FnDecl ,
14- Ident , Item , ItemFn , Pat , PatIdent , ReturnType , TypeTuple ,
14+ token, ArgCaptured , Error , Expr , ExprCall , ExprField , ExprForLoop , ExprMacro , ExprYield , FnArg ,
15+ FnDecl , Ident , Item , ItemFn , Member , Pat , PatIdent , ReturnType , TypeTuple ,
1516} ;
1617
1718#[ macro_use]
@@ -210,7 +211,7 @@ pub fn async_stream_block(input: TokenStream) -> TokenStream {
210211 tokens. into ( )
211212}
212213
213- /// The scope in which `#[for_await]`, `await! ` was called.
214+ /// The scope in which `#[for_await]`, `. await` was called.
214215///
215216/// The type of generator depends on which scope is called.
216217#[ derive( Debug , Clone , Copy , PartialEq , Eq ) ]
@@ -221,7 +222,7 @@ enum Scope {
221222 Stream ,
222223 /// `static move ||`, `||`
223224 ///
224- /// It cannot call `#[for_await]`, `await! ` in this scope.
225+ /// It cannot call `#[for_await]`, `. await` in this scope.
225226 Closure ,
226227}
227228
@@ -241,49 +242,66 @@ impl Expand {
241242 ) ) ;
242243 }
243244
245+ let ExprForLoop { label, pat, expr, body, .. } = & expr;
244246 // It needs to adjust the type yielded by the macro because generators used internally by
245247 // async fn yield `()` type, but generators used internally by `async_stream` yield
246248 // `Poll<U>` type.
247- let yield_ = match self . 0 {
248- Future => TokenStream2 :: new ( ) ,
249- Stream => quote ! { :: futures:: core_reexport:: task:: Poll :: Pending } ,
250- Closure => return outside_of_async_error ! ( expr, "#[for_await]" ) ,
251- } ;
252- let ExprForLoop { label, pat, expr, body, .. } = expr;
253-
254- // Basically just expand to a `poll` loop
255- syn:: parse_quote! { {
256- let mut __pinned = #expr;
257- #label
258- loop {
259- let #pat = {
260- match :: futures:: async_stream:: poll_next_with_tls_context( unsafe {
261- :: futures:: core_reexport:: pin:: Pin :: new_unchecked( & mut __pinned)
262- } )
263- {
264- :: futures:: core_reexport:: task:: Poll :: Ready ( e) => {
265- match e {
249+ match self . 0 {
250+ Future => {
251+ // Basically just expand to a `poll` loop
252+ syn:: parse_quote! { {
253+ let mut __pinned = #expr;
254+ let mut __pinned = unsafe {
255+ :: futures:: core_reexport:: pin:: Pin :: new_unchecked( & mut __pinned)
256+ } ;
257+ #label
258+ loop {
259+ let #pat = {
260+ match :: futures:: stream:: StreamExt :: next( & mut __pinned) . await {
266261 :: futures:: core_reexport:: option:: Option :: Some ( e) => e,
267262 :: futures:: core_reexport:: option:: Option :: None => break ,
268263 }
269- }
270- :: futures:: core_reexport:: task:: Poll :: Pending => {
271- yield #yield_;
272- continue
273- }
264+ } ;
265+
266+ #body
274267 }
275- } ;
268+ } }
269+ }
270+ Stream => {
271+ // Basically just expand to a `poll` loop
272+ syn:: parse_quote! { {
273+ let mut __pinned = #expr;
274+ #label
275+ loop {
276+ let #pat = {
277+ match :: futures:: async_stream:: poll_next_with_tls_context( unsafe {
278+ :: futures:: core_reexport:: pin:: Pin :: new_unchecked( & mut __pinned)
279+ } )
280+ {
281+ :: futures:: core_reexport:: task:: Poll :: Ready ( e) => {
282+ match e {
283+ :: futures:: core_reexport:: option:: Option :: Some ( e) => e,
284+ :: futures:: core_reexport:: option:: Option :: None => break ,
285+ }
286+ }
287+ :: futures:: core_reexport:: task:: Poll :: Pending => {
288+ yield :: futures:: core_reexport:: task:: Poll :: Pending ;
289+ continue
290+ }
291+ }
292+ } ;
276293
277- #body
294+ #body
295+ }
296+ } }
278297 }
279- } }
298+ Closure => return outside_of_async_error ! ( expr, "#[for_await]" ) ,
299+ }
280300 }
281301
282302 /// Expands `yield expr` in `async_stream` scope.
283303 fn expand_yield ( & self , expr : ExprYield ) -> ExprYield {
284- if self . 0 != Stream {
285- return expr;
286- }
304+ if self . 0 != Stream { return expr }
287305
288306 let ExprYield { attrs, yield_token, expr } = expr;
289307 let expr = expr. map_or_else ( || quote ! ( ( ) ) , ToTokens :: into_token_stream) ;
@@ -293,28 +311,30 @@ impl Expand {
293311 ExprYield { attrs, yield_token, expr : Some ( Box :: new ( expr) ) }
294312 }
295313
296- /// Expands a macro.
314+ /// Expands `async_stream_block!` macro.
297315 fn expand_macro ( & mut self , mut expr : ExprMacro ) -> Expr {
298- if self . 0 == Stream && expr. mac . path . is_ident ( "await" ) {
299- return self . expand_await_macros ( expr) ;
300- } else if expr. mac . path . is_ident ( "async_stream_block" ) {
316+ if expr. mac . path . is_ident ( "async_stream_block" ) {
301317 let mut e: ExprCall = syn:: parse ( async_stream_block ( expr. mac . tts . into ( ) ) ) . unwrap ( ) ;
302318 e. attrs . append ( & mut expr. attrs ) ;
303- return Expr :: Call ( e) ;
319+ Expr :: Call ( e)
320+ } else {
321+ Expr :: Macro ( expr)
304322 }
305-
306- Expr :: Macro ( expr)
307323 }
308324
309- /// Expands `await!( expr) ` in `async_stream` scope.
325+ /// Expands `expr.await ` in `async_stream` scope.
310326 ///
311327 /// It needs to adjust the type yielded by the macro because generators used internally by
312328 /// async fn yield `()` type, but generators used internally by `async_stream` yield
313329 /// `Poll<U>` type.
314- fn expand_await_macros ( & mut self , expr : ExprMacro ) -> Expr {
315- assert_eq ! ( self . 0 , Stream ) ;
330+ fn expand_await ( & mut self , expr : ExprField ) -> Expr {
331+ if self . 0 != Stream { return Expr :: Field ( expr ) }
316332
317- let expr = expr. mac . tts ;
333+ match & expr. member {
334+ Member :: Named ( x) if x == "await" => { }
335+ _ => return Expr :: Field ( expr) ,
336+ }
337+ let expr = expr. base ;
318338
319339 // Because macro input (`#expr`) is untrusted, use `syn::parse2` + `expr_compile_error`
320340 // instead of `syn::parse_quote!` to generate better error messages (`syn::parse_quote!`
@@ -349,8 +369,9 @@ impl Fold for Expand {
349369 }
350370
351371 let expr = match fold:: fold_expr ( self , expr) {
352- Expr :: ForLoop ( expr) => self . expand_for_await ( expr) ,
353372 Expr :: Yield ( expr) => Expr :: Yield ( self . expand_yield ( expr) ) ,
373+ Expr :: Field ( expr) => self . expand_await ( expr) ,
374+ Expr :: ForLoop ( expr) => self . expand_for_await ( expr) ,
354375 Expr :: Macro ( expr) => self . expand_macro ( expr) ,
355376 expr => expr,
356377 } ;
0 commit comments