vecmerge_impl/
lib.rs

1extern crate proc_macro;
2
3use proc_macro2::{Ident, Span, TokenStream};
4use proc_macro_hack::proc_macro_hack;
5use proc_quote::quote;
6
7#[proc_macro_hack]
8pub fn vecmerge(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
9    let expr = syn::parse_macro_input!(input as syn::Expr);
10    let mut visitor = Visitor::new();
11    visitor.visit(&expr);
12    let (capacity, parts) = visitor.get();
13    proc_macro::TokenStream::from(quote!({
14        let mut vec = Vec::with_capacity(#(#capacity)+*);
15        #(#parts)*
16        vec
17    }))
18}
19
20struct Visitor {
21    capacity: Vec<TokenStream>,
22    parts: Vec<TokenStream>,
23}
24
25impl Visitor {
26    fn new() -> Self {
27        Self {
28            capacity: Vec::new(),
29            parts: Vec::new(),
30        }
31    }
32
33    fn get(self) -> (Vec<TokenStream>, Vec<TokenStream>) {
34        (self.capacity, self.parts)
35    }
36
37    fn visit(&mut self, expr: &syn::Expr) {
38        match expr {
39            syn::Expr::Binary(ref binary_expr) => match binary_expr.op {
40                syn::BinOp::Add(..) => {
41                    self.visit(binary_expr.left.as_ref());
42                    self.visit(binary_expr.right.as_ref());
43                }
44                op => self.parts.push(quote! {
45                    compile_error!(format!("Unexpected operator {}", #op))
46                }),
47            },
48            syn::Expr::Try(ref try_expr) => match try_expr.expr.as_ref() {
49                syn::Expr::Array(array_expr) => self.visit_try_array_expr(array_expr),
50                _ => self.visit_expr(expr),
51            },
52            syn::Expr::Array(array_expr) => self.visit_array_expr(array_expr),
53            _ => self.visit_expr(expr),
54        }
55    }
56
57    fn visit_array_expr(&mut self, array_expr: &syn::ExprArray) {
58        let len = array_expr.elems.len();
59        self.capacity.push(quote!(#len));
60        if len == 1 {
61            let elem = &array_expr.elems[0];
62            self.parts.push(quote!(vec.push(#elem);));
63        } else {
64            self.parts
65                .push(quote!(vec.extend(<[_]>::into_vec(std::boxed::Box::new(#array_expr)));));
66        }
67    }
68
69    fn visit_try_array_expr(&mut self, try_array_expr: &syn::ExprArray) {
70        let len = try_array_expr.elems.len();
71        let elems = &try_array_expr.elems;
72        if len == 1 {
73            let elem = &elems[0];
74            self.parts.push(quote! {
75                if let Some(elem) = #elem {
76                    vec.push(elem);
77                }
78            });
79        } else {
80            let names = (0..elems.len())
81                .into_iter()
82                .map(|index| Ident::new(&format!("elem_{}", index), Span::call_site()))
83                .collect::<Vec<Ident>>();
84            self.parts.push(quote! {
85                if let (#(Some(#names)),*) = (#elems) {
86                    vec.extend(<[_]>::into_vec(std::boxed::Box::new([#(#names,)*])));
87                }
88            });
89        }
90    }
91
92    fn visit_expr(&mut self, expr: &syn::Expr) {
93        self.capacity.push(quote!(#expr.len()));
94        self.parts.push(quote!(vec.extend(#expr);));
95    }
96}
vecmerge_impl/lib.rs

vecmerge_impl/
lib.rs
