zng_var_proc_macros/
expr_var.rs

Help
1use proc_macro2::{Group, TokenStream, TokenTree};
2use quote::ToTokens;
3use syn::{
4    Expr, Ident, Path, Token,
5    parse::{Parse, ParseStream},
6    parse_macro_input, parse2,
7    spanned::Spanned,
8    token,
9};
10
11use crate::util::{token_stream_eq, tokens_to_ident_str};
12
13pub fn expand(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
14    let VarExpr { mod_, vars, expr } = parse_macro_input!(input as VarExpr);
15
16    let r = if vars.is_empty() {
17        // no interpolation, just eval to var.
18
19        if parse2::<Expr>(expr.clone()).is_ok() {
20            quote_spanned! {expr.span()=>
21                #mod_::types::expr_var_into(#expr)
22            }
23        } else {
24            // support statement blocks using the macro braces, if we just add the braces for
25            // all input it can cause the `unused_braces` lint, and we need the entire expression to have
26            // the span so that type mismatch gets highlighted correctly, so we *try* parse as expr and only
27            // add the braces if not.
28            quote_spanned! {expr.span()=>
29                #mod_::types::expr_var_into({#expr})
30            }
31        }
32    } else if vars.len() == 1 {
33        let (ident, eval) = &vars[0];
34
35        if token_stream_eq(expr.clone(), quote!(#ident)) || token_stream_eq(expr.clone(), quote!(*#ident)) {
36            // full expr is an interpolation, just return the var.
37            quote_spanned! {expr.span()=>
38                #mod_::types::expr_var_as(#eval)
39            }
40        } else {
41            quote_spanned! {expr.span()=>
42                // single var interpolation, use map.
43                #mod_::types::expr_var_map(#eval, move |#[allow(non_snake_case)]#ident|{ #expr })
44            }
45        }
46    } else {
47        // multiple var interpolation, use merge.
48        let idents = vars.iter().map(|(id, _)| id);
49        let evals = vars.iter().map(|(_, ev)| ev);
50        quote_spanned! {expr.span()=>
51            #mod_::types::expr_var_as(
52                #mod_::merge_var!{ #({#evals}),* , move |#(#[allow(non_snake_case)]#idents),*| { #expr } }
53            )
54        }
55    };
56
57    r.into()
58}
59
60struct VarExpr {
61    mod_: Path,
62    vars: Vec<(Ident, TokenStream)>,
63    expr: TokenStream,
64}
65impl Parse for VarExpr {
66    fn parse(input: ParseStream) -> syn::Result<Self> {
67        let mod_ = input.parse().unwrap_or_else(|e| non_user_error!(e));
68        input.parse::<Token![,]>().unwrap_or_else(|e| non_user_error!(e));
69        let mut vars = vec![];
70        let expr = parse_replace_expr(input, &mut vars);
71
72        Ok(VarExpr { mod_, vars, expr })
73    }
74}
75
76fn parse_replace_expr(input: ParseStream, vars: &mut Vec<(Ident, TokenStream)>) -> TokenStream {
77    let mut expr = TokenStream::default();
78
79    while !input.is_empty() {
80        // look for variable interpolation `#{<block>}` :
81        if input.peek(Token![#]) && input.peek2(token::Brace) {
82            input.parse::<Token![#]>().unwrap();
83            let var = input.parse::<Group>().unwrap().stream();
84            if let Some((var_ident, _)) = vars.iter().find(|(_, v)| token_stream_eq(v.clone(), var.clone())) {
85                var_ident.to_tokens(&mut expr)
86            } else {
87                let var_ident = ident_spanned!(var.span()=> "__{}_{}", vars.len(), tokens_to_ident_str(&var));
88                var_ident.to_tokens(&mut expr);
89                vars.push((var_ident, var))
90            }
91        }
92        // recursive parse groups:
93        else if input.peek(token::Brace) {
94            assert_group(|| {
95                let inner;
96                let group = syn::braced!(inner in input);
97                let inner = parse_replace_expr(&inner, vars);
98                group.surround(&mut expr, |e| e.extend(inner));
99                Ok(())
100            });
101        } else if input.peek(token::Paren) {
102            assert_group(|| {
103                let inner;
104                let group = syn::parenthesized!(inner in input);
105                let inner = parse_replace_expr(&inner, vars);
106                group.surround(&mut expr, |e| e.extend(inner));
107                Ok(())
108            });
109        } else if input.peek(token::Bracket) {
110            assert_group(|| {
111                let inner;
112                let group = syn::bracketed!(inner in input);
113                let inner = parse_replace_expr(&inner, vars);
114                group.surround(&mut expr, |e| e.extend(inner));
115                Ok(())
116            });
117        }
118        // keep other tokens the same:
119        else {
120            let tt = input.parse::<TokenTree>().unwrap();
121            tt.to_tokens(&mut expr)
122        }
123    }
124
125    expr
126}
127/// syn::braced! generates an error return.
128fn assert_group(f: impl FnOnce() -> syn::parse::Result<()>) {
129    f().unwrap()
130}
zng_var_proc_macros/expr_var.rs

zng_var_proc_macros/
expr_var.rs
