quote2_macros/
lib.rs

1use proc_macro::*;
2use std::mem;
3
4/// # Example
5/// ```rust
6/// use quote2::{proc_macro2::TokenStream, quote, Quote};
7/// let body = quote(|tokens| {
8///     for i in 0..3 {
9///         quote!(tokens, {
10///             println!("{}", #i);
11///         });
12///     }
13/// });
14/// let mut tokens = TokenStream::new();
15/// quote!(tokens, {
16///     fn name() {
17///         #body
18///     }
19/// });
20/// ```
21///
22/// ## Generated Code
23///
24/// ```rust
25/// fn name() {
26///     println!("{}", 0i32);
27///     println!("{}", 1i32);
28///     println!("{}", 2i32);
29/// }
30/// ```
31#[proc_macro]
32pub fn quote(input: TokenStream) -> TokenStream {
33    let mut input = input.into_iter();
34
35    let var = parse_arg(&mut input, "expected `ident`");
36
37    let input = match input.next() {
38        Some(TokenTree::Group(g)) => g.stream(),
39        _ => panic!("expected `{{`"),
40    };
41
42    let mut output = TokenStream::new();
43    expend(input, &mut output, None, var);
44    output
45}
46
47fn parse_arg(input: &mut token_stream::IntoIter, msg: &str) -> Ident {
48    let Some(TokenTree::Ident(var)) = input.next() else {
49        panic!("{msg}")
50    };
51    input.next().expect("expected `,`");
52    var
53}
54
55/// ## Example
56///
57/// ```rust
58/// use quote2::{
59///     proc_macro2::{Span, TokenStream},
60///     quote_spanned, Quote,
61/// };
62/// let span = Span::call_site();
63/// let mut tokens = TokenStream::new();
64/// quote_spanned!(span, tokens, {
65///     fn add(a: i32, b: i32) -> i32 {
66///         a + b
67///     }
68/// });
69/// ```
70#[proc_macro]
71pub fn quote_spanned(input: TokenStream) -> TokenStream {
72    let mut input = input.into_iter();
73
74    let span = parse_arg(&mut input, "expected `span`");
75    let var = parse_arg(&mut input, "expected `ident`");
76
77    let input = match input.next() {
78        Some(TokenTree::Group(g)) => g.stream(),
79        _ => panic!("expected `{{`"),
80    };
81
82    let mut output = TokenStream::new();
83    expend(input, &mut output, Some(&span), var);
84    output
85}
86
87fn expend(input: TokenStream, o: &mut TokenStream, span: Option<&Ident>, var: Ident) {
88    let mut input = input.into_iter().peekable();
89    let mut items = TokenStream::new();
90
91    while let Some(tree) = input.next() {
92        match tree {
93            TokenTree::Punct(punct) => {
94                let ch = punct.as_char();
95                if ch == '#' && matches!(input.peek(), Some(TokenTree::Ident(_))) {
96                    write_extender(&mut items, o, &var);
97                    let v = input.next();
98
99                    o.extend([
100                        tt(var.clone()),
101                        tt::punct('.'),
102                        tt::ident("add_tokens"),
103                        tt::group('(', |o| {
104                            add(o, tt::punct('&'));
105                            o.extend(v);
106                        }),
107                        tt::punct(';'),
108                    ]);
109                } else {
110                    let varient_ty = match (punct.spacing(), span.is_some()) {
111                        (Spacing::Joint, true) => "punct_join_span",
112                        (Spacing::Alone, true) => "punct_span",
113                        (Spacing::Joint, false) => "punct_join",
114                        (Spacing::Alone, false) => "punct",
115                    };
116                    varient(&mut items, varient_ty, |o| {
117                        add_span(o, span);
118                        add(o, tt::char(ch));
119                    });
120                }
121            }
122            TokenTree::Group(group) => {
123                let varient_ty = if span.is_some() {
124                    "group_span"
125                } else {
126                    "group"
127                };
128                varient(&mut items, varient_ty, |o| {
129                    add_span(o, span);
130
131                    let var = Ident::new("__o", Span::call_site());
132                    o.extend([
133                        tt::char(match group.delimiter() {
134                            Delimiter::None => '_',
135                            Delimiter::Brace => '{',
136                            Delimiter::Bracket => '[',
137                            Delimiter::Parenthesis => '(',
138                        }),
139                        tt::punct(','),
140                        tt::punct('|'),
141                        tt(var.clone()),
142                        tt::punct('|'),
143                        tt::group('{', |o| {
144                            expend(group.stream(), o, span, var);
145                        }),
146                    ]);
147                });
148            }
149            TokenTree::Ident(ident) => {
150                let varient_ty = if span.is_some() {
151                    "ident_span"
152                } else {
153                    "ident"
154                };
155                varient(&mut items, varient_ty, |o| {
156                    add_span(o, span);
157                    add(o, Literal::string(&ident.to_string()));
158                });
159            }
160            TokenTree::Literal(lit) => {
161                let varient_ty = if span.is_some() {
162                    "parsed_lit_span"
163                } else {
164                    "parsed_lit"
165                };
166                varient(&mut items, varient_ty, |o| {
167                    add_span(o, span);
168                    add(o, Literal::string(&lit.to_string()));
169                });
170            }
171        }
172    }
173    write_extender(&mut items, o, &var);
174}
175
176fn write_extender(items: &mut TokenStream, o: &mut TokenStream, var: &Ident) {
177    if !items.is_empty() {
178        let items = mem::take(items);
179        o.extend([
180            tt(var.clone()),
181            tt::punct('.'),
182            tt::ident("extend"),
183            tt::group('(', |o| {
184                add(o, Group::new(Delimiter::Bracket, items));
185            }),
186            tt::punct(';'),
187        ]);
188    }
189}
190
191fn add(o: &mut TokenStream, t: impl Into<TokenTree>) {
192    o.extend(Some(t.into()));
193}
194
195fn add_span(o: &mut TokenStream, span: Option<&Ident>) {
196    if let Some(spanned) = span {
197        o.extend([tt(spanned.clone()), tt::punct(',')]);
198    }
199}
200
201fn varient(t: &mut TokenStream, varient_ty: &str, f: impl FnOnce(&mut TokenStream)) {
202    t.extend([
203        tt::ident("quote2"),
204        tt::punct_joined(':'),
205        tt::punct(':'),
206        tt::ident("tt"),
207        tt::punct_joined(':'),
208        tt::punct(':'),
209        tt::ident(varient_ty),
210        tt::group('(', f),
211        tt::punct(','),
212    ]);
213}
214
215fn tt<T: Into<TokenTree>>(tt: T) -> TokenTree {
216    tt.into()
217}
218
219mod tt {
220    use super::*;
221    pub fn punct_joined(ch: char) -> TokenTree {
222        TokenTree::from(Punct::new(ch, Spacing::Joint))
223    }
224    pub fn punct(ch: char) -> TokenTree {
225        TokenTree::from(Punct::new(ch, Spacing::Alone))
226    }
227    pub fn ident(string: &str) -> TokenTree {
228        TokenTree::from(Ident::new(string, Span::call_site()))
229    }
230    pub fn char(ch: char) -> TokenTree {
231        TokenTree::from(Literal::character(ch))
232    }
233
234    pub fn group(delimiter: char, f: impl FnOnce(&mut TokenStream)) -> TokenTree {
235        let mut stream = TokenStream::new();
236        f(&mut stream);
237        let delimiter = match delimiter {
238            '{' => Delimiter::Brace,
239            '[' => Delimiter::Bracket,
240            '(' => Delimiter::Parenthesis,
241            _ => Delimiter::None,
242        };
243        Group::new(delimiter, stream).into()
244    }
245}
quote2_macros/lib.rs

quote2_macros/
lib.rs
