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