async_gen_macros/
lib.rs

1use proc_macro::*;
2
3#[proc_macro]
4pub fn gen_inner(input: TokenStream) -> TokenStream {
5    let mut tokens = input.into_iter();
6
7    let Some(TokenTree::Group(crate_path)) = tokens.next() else { unimplemented!() };
8    let crate_path = crate_path.stream();
9
10    let mut has_yielded = false;
11    let output = out(tokens, &mut has_yielded);
12
13    let mut o = TokenStream::new();
14    o.extend(crate_path.clone());
15    o.push_colon2();
16    o.push_ident("gen");
17
18    o.push_group(Delimiter::Parenthesis, |o| {
19        o.push_punct('|');
20        o.push_ident("mut");
21        o.push_ident("yield_");
22
23        if !has_yielded {
24            o.push_punct(':');
25            o.extend(crate_path);
26            o.push_colon2();
27            o.push_ident("Yield");
28        }
29
30        o.push_punct('|');
31        o.push_ident("async");
32        o.push_ident("move");
33        o.push_group(Delimiter::Brace, |o| {
34            o.push_ident("let");
35            o.push_ident("v");
36            o.push_punct('=');
37            o.push_ident("async");
38
39            o.push(output);
40
41            o.push_punct('.');
42            o.push_ident("await");
43            o.push_punct(';');
44            o.push_ident("yield_");
45            o.push_punct('.');
46            o.push_ident("return_");
47            o.push_group(Delimiter::Parenthesis, |o| {
48                o.push_ident("v");
49            });
50        });
51    });
52    o
53}
54
55fn out(mut tokens: token_stream::IntoIter, has_yielded: &mut bool) -> Group {
56    let mut o = TokenStream::new();
57
58    while let Some(tt) = tokens.next() {
59        match tt {
60            TokenTree::Ident(name) if name.to_string() == "yield" => {
61                *has_yielded = true;
62                let mut expr = TokenStream::new();
63                for tt in &mut tokens {
64                    match tt {
65                        TokenTree::Punct(p) if p.as_char() == ';' => break,
66                        _ => expr.push(tt),
67                    }
68                }
69                o.push_ident("yield_");
70                o.push_punct('.');
71                o.push_ident("yield_");
72                o.push(Group::new(Delimiter::Parenthesis, expr));
73                o.push_punct('.');
74                o.push_ident("await");
75                o.push_punct(';');
76            }
77            TokenTree::Group(g) if g.delimiter() == Delimiter::Brace => {
78                o.push(out(g.stream().into_iter(), has_yielded));
79            }
80            _ => o.push(tt),
81        }
82    }
83    Group::new(Delimiter::Brace, o)
84}
85
86trait TokenStreamExt {
87    fn push<U>(&mut self, token: U)
88    where
89        U: Into<TokenTree>;
90
91    #[inline]
92    fn push_punct(&mut self, ch: char) {
93        self.push(Punct::new(ch, Spacing::Alone))
94    }
95
96    #[inline]
97    fn push_ident(&mut self, name: &str) {
98        self.push(Ident::new(name, Span::call_site()))
99    }
100
101    #[inline]
102    fn push_group(&mut self, delimiter: Delimiter, f: impl FnOnce(&mut TokenStream)) {
103        let mut stream = TokenStream::new();
104        f(&mut stream);
105        self.push(Group::new(delimiter, stream))
106    }
107
108    #[inline]
109    fn push_colon2(&mut self) {
110        self.push(Punct::new(':', Spacing::Joint));
111        self.push(Punct::new(':', Spacing::Alone));
112    }
113}
114
115impl TokenStreamExt for TokenStream {
116    #[inline]
117    fn push<U>(&mut self, token: U)
118    where
119        U: Into<TokenTree>,
120    {
121        self.extend(std::iter::once(token.into()));
122    }
123}
async_gen_macros/lib.rs

async_gen_macros/
lib.rs
