1
1
extern crate proc_macro;
2
2
3
3
use proc_macro:: { TokenStream , TokenTree } ;
4
- use proc_macro2:: Span ;
4
+ use proc_macro2:: { Group , Span , TokenStream as TokenStream2 , TokenTree as TokenTree2 } ;
5
5
use quote:: quote;
6
6
use syn:: visit_mut:: VisitMut ;
7
7
@@ -19,7 +19,7 @@ struct AsyncStreamEnumHack {
19
19
}
20
20
21
21
impl AsyncStreamEnumHack {
22
- fn parse ( input : TokenStream ) -> Self {
22
+ fn parse ( input : TokenStream ) -> syn :: Result < Self > {
23
23
macro_rules! n {
24
24
( $i: ident) => {
25
25
$i. next( ) . unwrap( )
@@ -44,14 +44,15 @@ impl AsyncStreamEnumHack {
44
44
n ! ( braces) ; // !
45
45
46
46
let inner = n ! ( braces) ;
47
- let syn:: Block { stmts, .. } = syn:: parse ( inner. clone ( ) . into ( ) ) . unwrap ( ) ;
47
+ let inner = replace_for_await ( TokenStream2 :: from ( TokenStream :: from ( inner) ) ) ;
48
+ let syn:: Block { stmts, .. } = syn:: parse2 ( inner. clone ( ) ) ?;
48
49
49
50
let macro_ident = syn:: Ident :: new (
50
51
& format ! ( "stream_{}" , count_bangs( inner. into( ) ) ) ,
51
52
Span :: call_site ( ) ,
52
53
) ;
53
54
54
- AsyncStreamEnumHack { stmts, macro_ident }
55
+ Ok ( AsyncStreamEnumHack { stmts, macro_ident } )
55
56
}
56
57
}
57
58
@@ -153,7 +154,10 @@ pub fn async_stream_impl(input: TokenStream) -> TokenStream {
153
154
let AsyncStreamEnumHack {
154
155
macro_ident,
155
156
mut stmts,
156
- } = AsyncStreamEnumHack :: parse ( input) ;
157
+ } = match AsyncStreamEnumHack :: parse ( input) {
158
+ Ok ( x) => x,
159
+ Err ( e) => return e. to_compile_error ( ) . into ( ) ,
160
+ } ;
157
161
158
162
let mut scrub = Scrub {
159
163
is_xforming : true ,
@@ -192,7 +196,10 @@ pub fn async_try_stream_impl(input: TokenStream) -> TokenStream {
192
196
let AsyncStreamEnumHack {
193
197
macro_ident,
194
198
mut stmts,
195
- } = AsyncStreamEnumHack :: parse ( input) ;
199
+ } = match AsyncStreamEnumHack :: parse ( input) {
200
+ Ok ( x) => x,
201
+ Err ( e) => return e. to_compile_error ( ) . into ( ) ,
202
+ } ;
196
203
197
204
let mut scrub = Scrub {
198
205
is_xforming : true ,
@@ -245,3 +252,32 @@ fn count_bangs(input: TokenStream) -> usize {
245
252
246
253
count
247
254
}
255
+
256
+ fn replace_for_await ( input : TokenStream2 ) -> TokenStream2 {
257
+ let mut input = input. into_iter ( ) . peekable ( ) ;
258
+ let mut tokens = Vec :: new ( ) ;
259
+
260
+ while let Some ( token) = input. next ( ) {
261
+ match token {
262
+ TokenTree2 :: Ident ( ident) => {
263
+ match input. peek ( ) {
264
+ Some ( TokenTree2 :: Ident ( next) )
265
+ if ident. to_string ( ) . as_str ( ) == "for" && next == "await" =>
266
+ {
267
+ tokens. extend ( quote ! ( #[ #next] ) ) ;
268
+ let _ = input. next ( ) ;
269
+ }
270
+ _ => { }
271
+ }
272
+ tokens. push ( ident. into ( ) ) ;
273
+ }
274
+ TokenTree2 :: Group ( group) => {
275
+ let stream = replace_for_await ( group. stream ( ) ) ;
276
+ tokens. push ( Group :: new ( group. delimiter ( ) , stream) . into ( ) ) ;
277
+ }
278
+ _ => tokens. push ( token) ,
279
+ }
280
+ }
281
+
282
+ tokens. into_iter ( ) . collect ( )
283
+ }
0 commit comments