ct_for/
lib.rs

1use proc_macro2::Group;
2use quote::{quote, quote_spanned};
3
4/// Expands a code snippet `n` times and substitutes an identifier with given values
5///
6/// # Example
7/// ```
8/// use ct_for::ct_for;
9///
10/// let c = 8;
11/// ct_for!(x in ["5", 6, c] do
12///     println!("{}", x);
13/// );
14/// // expands to
15/// println!("{}", "5");
16/// println!("{}", 6);
17/// println!("{}", c);
18/// ```
19///
20#[proc_macro]
21pub fn ct_for(args: proc_macro::TokenStream) -> proc_macro::TokenStream {
22    let args = proc_macro2::TokenStream::from(args);
23    let mut iter = args.into_iter();
24
25    let subst_name = match iter.next() {
26        Some(proc_macro2::TokenTree::Ident(ident)) => ident,
27        Some(tok) => return error_span("Expected identifier", tok.span()),
28        None => return error("Unexpected end of macro"),
29    };
30
31    match iter.next() {
32        Some(proc_macro2::TokenTree::Ident(ident)) if ident == "in" => {}
33        Some(tok) => return error_span("Expected 'in'", tok.span()),
34        None => return error("Unexpected end of macro"),
35    }
36
37    let subst_group = match iter.next() {
38        Some(proc_macro2::TokenTree::Group(group)) => group,
39        Some(tok) => return error_span("Expected (...)", tok.span()),
40        None => return error("Unexpected end of macro"),
41    };
42
43    let subst_values = collect_subst_values(subst_group.stream().into_iter());
44
45    match iter.next() {
46        Some(proc_macro2::TokenTree::Ident(ident)) if ident == "do" => {}
47        Some(tok) => return error_span("Expected 'do'", tok.span()),
48        None => return error("Unexpected end of macro"),
49    }
50
51    let mut res = proc_macro2::TokenStream::new();
52
53    for subst in subst_values {
54        let iter = iter.clone();
55        res.extend(subst_recursive(iter, subst_name.clone(), subst));
56    }
57
58    res.into()
59}
60
61fn collect_subst_values(
62    iter: impl Iterator<Item = proc_macro2::TokenTree>,
63) -> Vec<proc_macro2::TokenStream> {
64    let mut res = Vec::new();
65
66    let mut val = proc_macro2::TokenStream::new();
67    for tok in iter {
68        match tok {
69            proc_macro2::TokenTree::Punct(punct) if punct.as_char() == ',' => {
70                res.push(val);
71                val = proc_macro2::TokenStream::new();
72            }
73            _ => {
74                val.extend(std::iter::once(tok));
75            }
76        }
77    }
78    if !val.is_empty() {
79        res.push(val);
80    }
81
82    res
83}
84
85fn subst_recursive(
86    iter: impl Iterator<Item = proc_macro2::TokenTree>,
87    subst_name: proc_macro2::Ident,
88    subst_value: proc_macro2::TokenStream,
89) -> proc_macro2::TokenStream {
90    let mut res = proc_macro2::TokenStream::new();
91
92    for tok in iter {
93        match tok {
94            proc_macro2::TokenTree::Group(group) => {
95                let group_tokens = subst_recursive(
96                    group.stream().into_iter(),
97                    subst_name.clone(),
98                    subst_value.clone(),
99                );
100                res.extend(std::iter::once(proc_macro2::TokenTree::Group(Group::new(
101                    group.delimiter(),
102                    group_tokens,
103                ))));
104            }
105            proc_macro2::TokenTree::Ident(ident) if ident == subst_name => {
106                res.extend(subst_value.clone());
107            }
108            _ => {
109                res.extend(std::iter::once(tok));
110            }
111        }
112    }
113
114    res
115}
116
117fn error(msg: &str) -> proc_macro::TokenStream {
118    quote! {
119        compile_error!(#msg);
120    }
121    .into()
122}
123
124fn error_span(msg: &str, span: proc_macro2::Span) -> proc_macro::TokenStream {
125    quote_spanned! {span=>
126        compile_error!(#msg);
127    }
128    .into()
129}
ct_for/lib.rs

ct_for/
lib.rs
