@@ -10,8 +10,8 @@ use proc_macro2::{Span, TokenStream as TokenStream2, TokenTree as TokenTree2};
10
10
use quote:: { quote, ToTokens } ;
11
11
use syn:: {
12
12
fold:: { self , Fold } ,
13
- token, ArgCaptured , Error , Expr , ExprCall , ExprForLoop , ExprMacro , ExprYield , FnArg , FnDecl ,
14
- Ident , Item , ItemFn , Pat , PatIdent , ReturnType , TypeTuple ,
13
+ token, ArgCaptured , Error , Expr , ExprCall , ExprField , ExprForLoop , ExprMacro , ExprYield , FnArg ,
14
+ FnDecl , Ident , Item , ItemFn , Member , Pat , PatIdent , ReturnType , TypeTuple ,
15
15
} ;
16
16
17
17
#[ macro_use]
@@ -210,7 +210,7 @@ pub fn async_stream_block(input: TokenStream) -> TokenStream {
210
210
tokens. into ( )
211
211
}
212
212
213
- /// The scope in which `#[for_await]`, `await! ` was called.
213
+ /// The scope in which `#[for_await]`, `. await` was called.
214
214
///
215
215
/// The type of generator depends on which scope is called.
216
216
#[ derive( Debug , Clone , Copy , PartialEq , Eq ) ]
@@ -221,7 +221,7 @@ enum Scope {
221
221
Stream ,
222
222
/// `static move ||`, `||`
223
223
///
224
- /// It cannot call `#[for_await]`, `await! ` in this scope.
224
+ /// It cannot call `#[for_await]`, `. await` in this scope.
225
225
Closure ,
226
226
}
227
227
@@ -241,49 +241,66 @@ impl Expand {
241
241
) ) ;
242
242
}
243
243
244
+ let ExprForLoop { label, pat, expr, body, .. } = & expr;
244
245
// It needs to adjust the type yielded by the macro because generators used internally by
245
246
// async fn yield `()` type, but generators used internally by `async_stream` yield
246
247
// `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 {
248
+ match self . 0 {
249
+ Future => {
250
+ // Basically just expand to a `poll` loop
251
+ syn:: parse_quote! { {
252
+ let mut __pinned = #expr;
253
+ let mut __pinned = unsafe {
254
+ :: futures:: core_reexport:: pin:: Pin :: new_unchecked( & mut __pinned)
255
+ } ;
256
+ #label
257
+ loop {
258
+ let #pat = {
259
+ match :: futures:: stream:: StreamExt :: next( & mut __pinned) . await {
266
260
:: futures:: core_reexport:: option:: Option :: Some ( e) => e,
267
261
:: futures:: core_reexport:: option:: Option :: None => break ,
268
262
}
269
- }
270
- :: futures:: core_reexport:: task:: Poll :: Pending => {
271
- yield #yield_;
272
- continue
273
- }
263
+ } ;
264
+
265
+ #body
274
266
}
275
- } ;
267
+ } }
268
+ }
269
+ Stream => {
270
+ // Basically just expand to a `poll` loop
271
+ syn:: parse_quote! { {
272
+ let mut __pinned = #expr;
273
+ #label
274
+ loop {
275
+ let #pat = {
276
+ match :: futures:: async_stream:: poll_next_with_tls_context( unsafe {
277
+ :: futures:: core_reexport:: pin:: Pin :: new_unchecked( & mut __pinned)
278
+ } )
279
+ {
280
+ :: futures:: core_reexport:: task:: Poll :: Ready ( e) => {
281
+ match e {
282
+ :: futures:: core_reexport:: option:: Option :: Some ( e) => e,
283
+ :: futures:: core_reexport:: option:: Option :: None => break ,
284
+ }
285
+ }
286
+ :: futures:: core_reexport:: task:: Poll :: Pending => {
287
+ yield :: futures:: core_reexport:: task:: Poll :: Pending ;
288
+ continue
289
+ }
290
+ }
291
+ } ;
276
292
277
- #body
293
+ #body
294
+ }
295
+ } }
278
296
}
279
- } }
297
+ Closure => return outside_of_async_error ! ( expr, "#[for_await]" ) ,
298
+ }
280
299
}
281
300
282
301
/// Expands `yield expr` in `async_stream` scope.
283
302
fn expand_yield ( & self , expr : ExprYield ) -> ExprYield {
284
- if self . 0 != Stream {
285
- return expr;
286
- }
303
+ if self . 0 != Stream { return expr }
287
304
288
305
let ExprYield { attrs, yield_token, expr } = expr;
289
306
let expr = expr. map_or_else ( || quote ! ( ( ) ) , ToTokens :: into_token_stream) ;
@@ -293,28 +310,30 @@ impl Expand {
293
310
ExprYield { attrs, yield_token, expr : Some ( Box :: new ( expr) ) }
294
311
}
295
312
296
- /// Expands a macro.
313
+ /// Expands `async_stream_block!` macro.
297
314
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" ) {
315
+ if expr. mac . path . is_ident ( "async_stream_block" ) {
301
316
let mut e: ExprCall = syn:: parse ( async_stream_block ( expr. mac . tts . into ( ) ) ) . unwrap ( ) ;
302
317
e. attrs . append ( & mut expr. attrs ) ;
303
- return Expr :: Call ( e) ;
318
+ Expr :: Call ( e)
319
+ } else {
320
+ Expr :: Macro ( expr)
304
321
}
305
-
306
- Expr :: Macro ( expr)
307
322
}
308
323
309
- /// Expands `await!( expr) ` in `async_stream` scope.
324
+ /// Expands `expr.await ` in `async_stream` scope.
310
325
///
311
326
/// It needs to adjust the type yielded by the macro because generators used internally by
312
327
/// async fn yield `()` type, but generators used internally by `async_stream` yield
313
328
/// `Poll<U>` type.
314
- fn expand_await_macros ( & mut self , expr : ExprMacro ) -> Expr {
315
- assert_eq ! ( self . 0 , Stream ) ;
329
+ fn expand_await ( & mut self , expr : ExprField ) -> Expr {
330
+ if self . 0 != Stream { return Expr :: Field ( expr ) }
316
331
317
- let expr = expr. mac . tts ;
332
+ match & expr. member {
333
+ Member :: Named ( x) if x == "await" => { }
334
+ _ => return Expr :: Field ( expr) ,
335
+ }
336
+ let expr = expr. base ;
318
337
319
338
// Because macro input (`#expr`) is untrusted, use `syn::parse2` + `expr_compile_error`
320
339
// instead of `syn::parse_quote!` to generate better error messages (`syn::parse_quote!`
@@ -349,8 +368,9 @@ impl Fold for Expand {
349
368
}
350
369
351
370
let expr = match fold:: fold_expr ( self , expr) {
352
- Expr :: ForLoop ( expr) => self . expand_for_await ( expr) ,
353
371
Expr :: Yield ( expr) => Expr :: Yield ( self . expand_yield ( expr) ) ,
372
+ Expr :: Field ( expr) => self . expand_await ( expr) ,
373
+ Expr :: ForLoop ( expr) => self . expand_for_await ( expr) ,
354
374
Expr :: Macro ( expr) => self . expand_macro ( expr) ,
355
375
expr => expr,
356
376
} ;
0 commit comments