zng_var_proc_macros/
expr_var.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
use proc_macro2::{Group, TokenStream, TokenTree};
use quote::ToTokens;
use syn::{
    Expr, Ident, Path, Token,
    parse::{Parse, ParseStream},
    parse_macro_input, parse2,
    spanned::Spanned,
    token,
};

use crate::util::{token_stream_eq, tokens_to_ident_str};

pub fn expand(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let VarExpr { mod_, vars, expr } = parse_macro_input!(input as VarExpr);

    let r = if vars.is_empty() {
        // no interpolation, just eval to var.

        if parse2::<Expr>(expr.clone()).is_ok() {
            quote_spanned! {expr.span()=>
                #mod_::types::expr_var_into(#expr)
            }
        } else {
            // support statement blocks using the macro braces, if we just add the braces for
            // all input it can cause the `unused_braces` lint, and we need the entire expression to have
            // the span so that type mismatch gets highlighted correctly, so we *try* parse as expr and only
            // add the braces if not.
            quote_spanned! {expr.span()=>
                #mod_::types::expr_var_into({#expr})
            }
        }
    } else if vars.len() == 1 {
        let (ident, eval) = &vars[0];

        if token_stream_eq(expr.clone(), quote!(#ident)) || token_stream_eq(expr.clone(), quote!(*#ident)) {
            // full expr is an interpolation, just return the var.
            quote_spanned! {expr.span()=>
                #mod_::types::expr_var_as(#eval)
            }
        } else {
            quote_spanned! {expr.span()=>
                // single var interpolation, use map.
                #mod_::types::expr_var_map(#eval, move |#[allow(non_snake_case)]#ident|{ #expr })
            }
        }
    } else {
        // multiple var interpolation, use merge.
        let idents = vars.iter().map(|(id, _)| id);
        let evals = vars.iter().map(|(_, ev)| ev);
        quote_spanned! {expr.span()=>
            #mod_::types::expr_var_as(
                #mod_::merge_var!{ #({#evals}),* , move |#(#[allow(non_snake_case)]#idents),*| { #expr } }
            )
        }
    };

    r.into()
}

struct VarExpr {
    mod_: Path,
    vars: Vec<(Ident, TokenStream)>,
    expr: TokenStream,
}
impl Parse for VarExpr {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let mod_ = input.parse().unwrap_or_else(|e| non_user_error!(e));
        input.parse::<Token![,]>().unwrap_or_else(|e| non_user_error!(e));
        let mut vars = vec![];
        let expr = parse_replace_expr(input, &mut vars);

        Ok(VarExpr { mod_, vars, expr })
    }
}

fn parse_replace_expr(input: ParseStream, vars: &mut Vec<(Ident, TokenStream)>) -> TokenStream {
    let mut expr = TokenStream::default();

    while !input.is_empty() {
        // look for variable interpolation `#{<block>}` :
        if input.peek(Token![#]) && input.peek2(token::Brace) {
            input.parse::<Token![#]>().unwrap();
            let var = input.parse::<Group>().unwrap().stream();
            if let Some((var_ident, _)) = vars.iter().find(|(_, v)| token_stream_eq(v.clone(), var.clone())) {
                var_ident.to_tokens(&mut expr)
            } else {
                let var_ident = ident_spanned!(var.span()=> "__{}_{}", vars.len(), tokens_to_ident_str(&var));
                var_ident.to_tokens(&mut expr);
                vars.push((var_ident, var))
            }
        }
        // recursive parse groups:
        else if input.peek(token::Brace) {
            assert_group(|| {
                let inner;
                let group = syn::braced!(inner in input);
                let inner = parse_replace_expr(&inner, vars);
                group.surround(&mut expr, |e| e.extend(inner));
                Ok(())
            });
        } else if input.peek(token::Paren) {
            assert_group(|| {
                let inner;
                let group = syn::parenthesized!(inner in input);
                let inner = parse_replace_expr(&inner, vars);
                group.surround(&mut expr, |e| e.extend(inner));
                Ok(())
            });
        } else if input.peek(token::Bracket) {
            assert_group(|| {
                let inner;
                let group = syn::bracketed!(inner in input);
                let inner = parse_replace_expr(&inner, vars);
                group.surround(&mut expr, |e| e.extend(inner));
                Ok(())
            });
        }
        // keep other tokens the same:
        else {
            let tt = input.parse::<TokenTree>().unwrap();
            tt.to_tokens(&mut expr)
        }
    }

    expr
}
/// syn::braced! generates an error return.
fn assert_group(f: impl FnOnce() -> syn::parse::Result<()>) {
    f().unwrap()
}