concat_strs_impl/
lib.rs

1//! Implementation crate for `concat_strs`, needed for [`proc_macro_hack`][proc_macro_hack]
2//!
3//! [proc_macro_hack]: https://docs.rs/proc-macro-hack/
4
5use proc_macro as pm;
6use proc_macro2::{Span, TokenStream};
7use proc_macro_hack::proc_macro_hack;
8use quote::{quote, ToTokens};
9use syn::parse::{Parse, ParseStream};
10use syn::punctuated::Punctuated;
11use syn::token;
12
13/// `concat_strs` macro; see documentation for the `concat_strs` crate for more information.
14#[proc_macro_hack]
15pub fn concat_strs(input: pm::TokenStream) -> pm::TokenStream {
16    let macro_input = syn::parse_macro_input!(input as ConcatStrs);
17    impl_concat_strs(macro_input).into()
18}
19
20struct ConcatStrs {
21    exprs: Punctuated<syn::Expr, token::Comma>,
22}
23
24impl Parse for ConcatStrs {
25    fn parse(input: ParseStream) -> syn::Result<Self> {
26        Ok(Self {
27            exprs: Punctuated::parse_terminated(input)?,
28        })
29    }
30}
31
32impl ConcatStrs {
33    fn bytes_needed(&self) -> usize {
34        self.exprs.iter().map(|e| byte_count(e)).sum()
35    }
36}
37
38fn byte_count(expr: &syn::Expr) -> usize {
39    match expr {
40        syn::Expr::Lit(lit) => match &lit.lit {
41            syn::Lit::Str(lit) => lit.value().len(),
42            syn::Lit::Char(lit) => lit.value().len_utf8(),
43            syn::Lit::Int(_) | syn::Lit::Float(_) => lit.lit.to_token_stream().to_string().len(),
44            _ => 0,
45        },
46        _ => 0,
47    }
48}
49
50fn tmp_assign_tokens(
51    expr: &syn::Expr,
52    tmp_ident: &syn::Ident,
53) -> Option<(syn::Ident, TokenStream)> {
54    match expr {
55        syn::Expr::Lit(_) => None,
56        _ => {
57            let len_ident = syn::Ident::new(&format!("{}_len", tmp_ident), Span::call_site());
58            Some((
59                len_ident.clone(),
60                quote! {
61                    let #tmp_ident = #expr;
62                    let #len_ident = #tmp_ident.len();
63                },
64            ))
65        }
66    }
67}
68
69fn push_tokens(expr: &syn::Expr, tmp_ident: &syn::Ident, string_ident: &syn::Ident) -> TokenStream {
70    match expr {
71        syn::Expr::Lit(expr) => match expr.lit {
72            syn::Lit::Str(_) => {
73                quote! {
74                    #string_ident.push_str(#expr);
75                }
76            }
77            syn::Lit::Char(_) => {
78                quote! {
79                    #string_ident.push(#expr);
80                }
81            }
82            syn::Lit::Float(_) | syn::Lit::Int(_) => {
83                let expr_string =
84                    syn::LitStr::new(&expr.lit.to_token_stream().to_string(), Span::call_site());
85                quote! {
86                    #string_ident.push_str(#expr_string);
87                }
88            }
89
90            _ => panic!("Unsupported expression {}", expr.to_token_stream()),
91        },
92        _ => {
93            quote! {
94                #string_ident.push_str(#tmp_ident);
95            }
96        }
97    }
98}
99
100fn tmp_idents() -> impl Iterator<Item = syn::Ident> {
101    (0..)
102        .into_iter()
103        .map(|n| syn::Ident::new(&format!("__tmp_{}", n), Span::call_site()))
104}
105
106fn impl_concat_strs(items: ConcatStrs) -> TokenStream {
107    let bytes = items.bytes_needed();
108    let string_ident: syn::Ident = syn::Ident::new("__ret", Span::call_site());
109
110    let tmp_assignments_idents = items
111        .exprs
112        .iter()
113        .zip(tmp_idents())
114        .map(|(expr, ident)| (tmp_assign_tokens(expr, &ident), ident))
115        .collect::<Vec<_>>();
116
117    let tmp_len_idents = tmp_assignments_idents
118        .iter()
119        .filter_map(|(assign, _ident)| {
120            if let Some((len_ident, _assign)) = assign {
121                Some(len_ident)
122            } else {
123                None
124            }
125        })
126        .collect::<Vec<_>>();
127
128    let tmp_assignments = tmp_assignments_idents
129        .iter()
130        .filter_map(|(assign, _ident)| assign.clone().map(|(_len_ident, assign)| assign))
131        .collect::<Vec<_>>();
132
133    let push_exprs = items
134        .exprs
135        .iter()
136        .zip(&tmp_assignments_idents)
137        .map(|(expr, (_expr, tmp_ident))| push_tokens(expr, &tmp_ident, &string_ident))
138        .collect::<Vec<_>>();
139
140    quote! {{
141        #(#tmp_assignments)*
142        let mut #string_ident = String::with_capacity(#(#tmp_len_idents +)* #bytes);
143        #(#push_exprs)*
144        #string_ident
145    }}
146}