1 extern crate proc_macro;
2 use proc_macro::TokenStream;
3 use proc_macro2::{Group, TokenStream as TokenStream2, TokenTree};
4 use quote::quote;
5 use syn::parse::{Parse, ParseStream, Parser, Result};
6 use syn::visit_mut::VisitMut;
7 
8 struct Scrub<'a> {
9     /// Whether the stream is a try stream.
10     is_try: bool,
11     /// The unit expression, `()`.
12     unit: Box<syn::Expr>,
13     has_yielded: bool,
14     crate_path: &'a TokenStream2,
15 }
16 
parse_input(input: TokenStream) -> syn::Result<(TokenStream2, Vec<syn::Stmt>)>17 fn parse_input(input: TokenStream) -> syn::Result<(TokenStream2, Vec<syn::Stmt>)> {
18     let mut input = TokenStream2::from(input).into_iter();
19     let crate_path = match input.next().unwrap() {
20         TokenTree::Group(group) => group.stream(),
21         _ => panic!(),
22     };
23     let stmts = syn::Block::parse_within.parse2(replace_for_await(input))?;
24     Ok((crate_path, stmts))
25 }
26 
27 impl<'a> Scrub<'a> {
new(is_try: bool, crate_path: &'a TokenStream2) -> Self28     fn new(is_try: bool, crate_path: &'a TokenStream2) -> Self {
29         Self {
30             is_try,
31             unit: syn::parse_quote!(()),
32             has_yielded: false,
33             crate_path,
34         }
35     }
36 }
37 
38 struct Partial<T>(T, TokenStream2);
39 
40 impl<T: Parse> Parse for Partial<T> {
parse(input: ParseStream) -> Result<Self>41     fn parse(input: ParseStream) -> Result<Self> {
42         Ok(Partial(input.parse()?, input.parse()?))
43     }
44 }
45 
visit_token_stream_impl( visitor: &mut Scrub<'_>, tokens: TokenStream2, modified: &mut bool, out: &mut TokenStream2, )46 fn visit_token_stream_impl(
47     visitor: &mut Scrub<'_>,
48     tokens: TokenStream2,
49     modified: &mut bool,
50     out: &mut TokenStream2,
51 ) {
52     use quote::ToTokens;
53     use quote::TokenStreamExt;
54 
55     let mut tokens = tokens.into_iter().peekable();
56     while let Some(tt) = tokens.next() {
57         match tt {
58             TokenTree::Ident(i) if i == "yield" => {
59                 let stream = std::iter::once(TokenTree::Ident(i)).chain(tokens).collect();
60                 match syn::parse2(stream) {
61                     Ok(Partial(yield_expr, rest)) => {
62                         let mut expr = syn::Expr::Yield(yield_expr);
63                         visitor.visit_expr_mut(&mut expr);
64                         expr.to_tokens(out);
65                         *modified = true;
66                         tokens = rest.into_iter().peekable();
67                     }
68                     Err(e) => {
69                         out.append_all(&mut e.to_compile_error().into_iter());
70                         *modified = true;
71                         return;
72                     }
73                 }
74             }
75             TokenTree::Ident(i) if i == "stream" || i == "try_stream" => {
76                 out.append(TokenTree::Ident(i));
77                 match tokens.peek() {
78                     Some(TokenTree::Punct(p)) if p.as_char() == '!' => {
79                         out.extend(tokens.next()); // !
80                         if let Some(TokenTree::Group(_)) = tokens.peek() {
81                             out.extend(tokens.next()); // { .. } or [ .. ] or ( .. )
82                         }
83                     }
84                     _ => {}
85                 }
86             }
87             TokenTree::Group(group) => {
88                 let mut content = group.stream();
89                 *modified |= visitor.visit_token_stream(&mut content);
90                 let mut new = Group::new(group.delimiter(), content);
91                 new.set_span(group.span());
92                 out.append(new);
93             }
94             other => out.append(other),
95         }
96     }
97 }
98 
99 impl Scrub<'_> {
visit_token_stream(&mut self, tokens: &mut TokenStream2) -> bool100     fn visit_token_stream(&mut self, tokens: &mut TokenStream2) -> bool {
101         let (mut out, mut modified) = (TokenStream2::new(), false);
102         visit_token_stream_impl(self, tokens.clone(), &mut modified, &mut out);
103 
104         if modified {
105             *tokens = out;
106         }
107 
108         modified
109     }
110 }
111 
112 impl VisitMut for Scrub<'_> {
visit_expr_mut(&mut self, i: &mut syn::Expr)113     fn visit_expr_mut(&mut self, i: &mut syn::Expr) {
114         match i {
115             syn::Expr::Yield(yield_expr) => {
116                 self.has_yielded = true;
117 
118                 syn::visit_mut::visit_expr_yield_mut(self, yield_expr);
119 
120                 let value_expr = yield_expr.expr.as_ref().unwrap_or(&self.unit);
121 
122                 // let ident = &self.yielder;
123 
124                 *i = if self.is_try {
125                     syn::parse_quote! { __yield_tx.send(::core::result::Result::Ok(#value_expr)).await }
126                 } else {
127                     syn::parse_quote! { __yield_tx.send(#value_expr).await }
128                 };
129             }
130             syn::Expr::Try(try_expr) => {
131                 syn::visit_mut::visit_expr_try_mut(self, try_expr);
132                 // let ident = &self.yielder;
133                 let e = &try_expr.expr;
134 
135                 *i = syn::parse_quote! {
136                     match #e {
137                         ::core::result::Result::Ok(v) => v,
138                         ::core::result::Result::Err(e) => {
139                             __yield_tx.send(::core::result::Result::Err(e.into())).await;
140                             return;
141                         }
142                     }
143                 };
144             }
145             syn::Expr::Closure(_) | syn::Expr::Async(_) => {
146                 // Don't transform inner closures or async blocks.
147             }
148             syn::Expr::ForLoop(expr) => {
149                 syn::visit_mut::visit_expr_for_loop_mut(self, expr);
150                 // TODO: Should we allow other attributes?
151                 if expr.attrs.len() != 1 || !expr.attrs[0].meta.path().is_ident(AWAIT_ATTR_NAME) {
152                     return;
153                 }
154                 let syn::ExprForLoop {
155                     attrs,
156                     label,
157                     pat,
158                     expr,
159                     body,
160                     ..
161                 } = expr;
162 
163                 attrs.pop().unwrap();
164 
165                 let crate_path = self.crate_path;
166                 *i = syn::parse_quote! {{
167                     let mut __pinned = #expr;
168                     let mut __pinned = unsafe {
169                         ::core::pin::Pin::new_unchecked(&mut __pinned)
170                     };
171                     #label
172                     loop {
173                         let #pat = match #crate_path::__private::next(&mut __pinned).await {
174                             ::core::option::Option::Some(e) => e,
175                             ::core::option::Option::None => break,
176                         };
177                         #body
178                     }
179                 }}
180             }
181             _ => syn::visit_mut::visit_expr_mut(self, i),
182         }
183     }
184 
visit_macro_mut(&mut self, mac: &mut syn::Macro)185     fn visit_macro_mut(&mut self, mac: &mut syn::Macro) {
186         let mac_ident = mac.path.segments.last().map(|p| &p.ident);
187         if mac_ident.map_or(false, |i| i == "stream" || i == "try_stream") {
188             return;
189         }
190 
191         self.visit_token_stream(&mut mac.tokens);
192     }
193 
visit_item_mut(&mut self, i: &mut syn::Item)194     fn visit_item_mut(&mut self, i: &mut syn::Item) {
195         // Recurse into macros but otherwise don't transform inner items.
196         if let syn::Item::Macro(i) = i {
197             self.visit_macro_mut(&mut i.mac);
198         }
199     }
200 }
201 
202 /// The first token tree in the stream must be a group containing the path to the `async-stream`
203 /// crate.
204 #[proc_macro]
205 #[doc(hidden)]
stream_inner(input: TokenStream) -> TokenStream206 pub fn stream_inner(input: TokenStream) -> TokenStream {
207     let (crate_path, mut stmts) = match parse_input(input) {
208         Ok(x) => x,
209         Err(e) => return e.to_compile_error().into(),
210     };
211 
212     let mut scrub = Scrub::new(false, &crate_path);
213 
214     for stmt in &mut stmts {
215         scrub.visit_stmt_mut(stmt);
216     }
217 
218     let dummy_yield = if scrub.has_yielded {
219         None
220     } else {
221         Some(quote!(if false {
222             __yield_tx.send(()).await;
223         }))
224     };
225 
226     quote!({
227         let (mut __yield_tx, __yield_rx) = unsafe { #crate_path::__private::yielder::pair() };
228         #crate_path::__private::AsyncStream::new(__yield_rx, async move {
229             #dummy_yield
230             #(#stmts)*
231         })
232     })
233     .into()
234 }
235 
236 /// The first token tree in the stream must be a group containing the path to the `async-stream`
237 /// crate.
238 #[proc_macro]
239 #[doc(hidden)]
try_stream_inner(input: TokenStream) -> TokenStream240 pub fn try_stream_inner(input: TokenStream) -> TokenStream {
241     let (crate_path, mut stmts) = match parse_input(input) {
242         Ok(x) => x,
243         Err(e) => return e.to_compile_error().into(),
244     };
245 
246     let mut scrub = Scrub::new(true, &crate_path);
247 
248     for stmt in &mut stmts {
249         scrub.visit_stmt_mut(stmt);
250     }
251 
252     let dummy_yield = if scrub.has_yielded {
253         None
254     } else {
255         Some(quote!(if false {
256             __yield_tx.send(()).await;
257         }))
258     };
259 
260     quote!({
261         let (mut __yield_tx, __yield_rx) = unsafe { #crate_path::__private::yielder::pair() };
262         #crate_path::__private::AsyncStream::new(__yield_rx, async move {
263             #dummy_yield
264             #(#stmts)*
265         })
266     })
267     .into()
268 }
269 
270 // syn 2.0 wont parse `#[await] for x in xs {}`
271 // because `await` is a keyword, use `await_` instead
272 const AWAIT_ATTR_NAME: &str = "await_";
273 
274 /// Replace `for await` with `#[await] for`, which will be later transformed into a `next` loop.
replace_for_await(input: impl IntoIterator<Item = TokenTree>) -> TokenStream2275 fn replace_for_await(input: impl IntoIterator<Item = TokenTree>) -> TokenStream2 {
276     let mut input = input.into_iter().peekable();
277     let mut tokens = Vec::new();
278 
279     while let Some(token) = input.next() {
280         match token {
281             TokenTree::Ident(ident) => {
282                 match input.peek() {
283                     Some(TokenTree::Ident(next)) if ident == "for" && next == "await" => {
284                         let next_span = next.span();
285                         let next = syn::Ident::new(AWAIT_ATTR_NAME, next_span);
286                         tokens.extend(quote!(#[#next]));
287                         let _ = input.next();
288                     }
289                     _ => {}
290                 }
291                 tokens.push(ident.into());
292             }
293             TokenTree::Group(group) => {
294                 let stream = replace_for_await(group.stream());
295                 let mut new_group = Group::new(group.delimiter(), stream);
296                 new_group.set_span(group.span());
297                 tokens.push(new_group.into());
298             }
299             _ => tokens.push(token),
300         }
301     }
302 
303     tokens.into_iter().collect()
304 }
305