concat_strs_derive/
lib.rs

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

concat_strs_derive/
lib.rs
