1use 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#[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}