1#![warn(clippy::pedantic, rust_2018_idioms)]
2
3use std::array;
4
5use bytestring::ByteString;
6use proc_macro2::TokenStream;
7use quote::{format_ident, quote};
8use syn::{punctuated::Punctuated, spanned::Spanned as _, Error, Expr, Ident, LitStr, Token};
9
10fn ident_to_expr(ident: Ident) -> syn::Expr {
11 syn::Expr::Path(syn::ExprPath {
12 attrs: Vec::new(),
13 qself: None,
14 path: syn::Path {
15 leading_colon: None,
16 segments: Punctuated::from_iter([syn::PathSegment {
17 arguments: syn::PathArguments::None,
18 ident,
19 }]),
20 },
21 })
22}
23
24fn expr_to_ident(expr: &syn::Expr) -> Option<&Ident> {
25 if let syn::Expr::Path(path) = expr {
26 if path.attrs.is_empty() && path.qself.is_none() {
27 return path.path.get_ident();
28 }
29 }
30
31 None
32}
33
34enum Piece {
36 Literal(ByteString),
38 Argument {
40 expr: Expr,
42 ident: Ident,
44 },
45 ArgumentRef { ident: Ident },
47}
48
49impl Piece {
50 fn new_arg_dedupe(new_ident: Ident, existing_pieces: &[Piece]) -> Self {
51 let expr = ident_to_expr(new_ident.clone());
52 let new_ident = format_ident!("arg_{new_ident}");
53
54 let existing = existing_pieces.iter().find_map(|p| {
55 if let Piece::Argument { ident, .. } = p {
56 if &new_ident == ident {
57 return Some(ident.clone());
58 }
59 }
60
61 None
62 });
63
64 if let Some(ident) = existing {
65 Piece::ArgumentRef { ident }
66 } else {
67 Piece::Argument {
68 expr,
69 ident: new_ident,
70 }
71 }
72 }
73
74 fn as_ident(&self) -> Option<&Ident> {
75 if let Self::Argument { ident, .. } = self {
76 Some(ident)
77 } else {
78 None
79 }
80 }
81
82 fn as_expr(&self) -> Option<&Expr> {
83 if let Self::Argument { expr, .. } = self {
84 Some(expr)
85 } else {
86 None
87 }
88 }
89}
90
91impl quote::ToTokens for Piece {
92 fn to_tokens(&self, tokens: &mut TokenStream) {
93 match self {
94 Self::Literal(str) if str.is_empty() => {}
95 Self::Literal(str) => {
96 let str: &str = str;
97 quote!(out.push_str(#str);).to_tokens(tokens)
98 }
99 Self::Argument { ident, .. } | Self::ArgumentRef { ident } => {
100 quote!(out.push_str(#ident.as_str());).to_tokens(tokens)
101 }
102 }
103 }
104}
105
106struct ConcatInvoke(Punctuated<LitStr, Token![,]>);
107
108impl ConcatInvoke {
109 fn evaluate(self) -> LitStr {
110 let span = self.0.span();
111 let mut strings = self.0.into_iter();
112 if let Some(first) = strings.next() {
113 let mut out_string = first.value();
114 for string in strings {
115 out_string.push_str(&string.value());
116 }
117
118 LitStr::new(&out_string, span)
119 } else {
120 LitStr::new("", span)
121 }
122 }
123}
124
125impl syn::parse::Parse for ConcatInvoke {
126 fn parse(input: syn::parse::ParseStream<'_>) -> syn::Result<Self> {
127 let call: syn::Macro = input.parse()?;
128 let Some(ident) = call.path.get_ident() else {
129 return Err(input.error("expected `concat!`"));
130 };
131
132 if ident != "concat" {
133 return Err(input.error("expected `concat!`"));
134 }
135
136 call.parse_body_with(Punctuated::parse_terminated).map(Self)
137 }
138}
139
140struct Arguments {
141 str_base_len: usize,
142 pieces: Vec<Piece>,
143}
144
145impl syn::parse::Parse for Arguments {
146 fn parse(input: syn::parse::ParseStream<'_>) -> syn::Result<Self> {
147 let format_str = {
148 let input_copy = input.fork();
149 if input_copy.peek(LitStr) {
150 input.parse()?
151 } else if let Ok(invoke) = input.parse::<ConcatInvoke>() {
152 invoke.evaluate()
153 } else {
154 return Err(input_copy.error("Expected literal string or concat! invoke"));
155 }
156 };
157
158 input.parse::<Option<Token![,]>>()?;
159
160 let create_err = |msg| Error::new(format_str.span(), msg);
161 let unterminated_fmt = move || create_err("Unterminated format argument");
162
163 let format_string = ByteString::from(format_str.value());
164 let mut current: &str = &*format_string;
165
166 let mut str_base_len = 0;
167 let mut pieces = Vec::new();
168
169 for arg_num in 0_u8.. {
170 let Some((text, rest)) = current.split_once('{') else {
171 str_base_len += current.len();
172 pieces.push(Piece::Literal(format_string.slice_ref(current)));
173 break;
174 };
175
176 let (arg_name, rest) = rest.split_once('}').ok_or_else(unterminated_fmt)?;
177
178 let arg_piece = if arg_name.is_empty() {
179 if input.is_empty() {
180 return Err(create_err("Not enough arguments for format string"));
181 }
182
183 let argument = input.parse::<syn::Expr>()?;
184 if input.parse::<Option<Token![,]>>()?.is_none() && !input.is_empty() {
185 return Err(Error::new(input.span(), "Missing argument seperator (`,`)"));
186 };
187
188 if let Some(ident) = expr_to_ident(&argument) {
189 Piece::new_arg_dedupe(ident.clone(), &pieces)
190 } else {
191 Piece::Argument {
192 expr: argument,
193 ident: format_ident!("arg_{arg_num}"),
194 }
195 }
196 } else {
197 let new_ident = Ident::new(arg_name, format_str.span());
198 Piece::new_arg_dedupe(new_ident, &pieces)
199 };
200
201 current = rest;
202
203 str_base_len += text.len();
204 pieces.push(Piece::Literal(format_string.slice_ref(text)));
205 pieces.push(arg_piece);
206 }
207
208 Ok(Self {
209 str_base_len,
210 pieces,
211 })
212 }
213}
214
215struct FormatIntoArguments {
216 write_into: Ident,
217 arguments: Arguments,
218}
219
220impl syn::parse::Parse for FormatIntoArguments {
221 fn parse(input: syn::parse::ParseStream<'_>) -> syn::Result<Self> {
222 let write_into = input.parse()?;
223 input.parse::<Token![,]>()?;
224
225 let arguments = Arguments::parse(input)?;
226
227 Ok(Self {
228 write_into,
229 arguments,
230 })
231 }
232}
233
234#[proc_macro]
245pub fn aformat(tokens: proc_macro::TokenStream) -> proc_macro::TokenStream {
246 let Arguments {
247 str_base_len,
248 pieces,
249 } = match syn::parse(tokens) {
250 Ok(args) => args,
251 Err(err) => return err.into_compile_error().into(),
252 };
253
254 let caller_arguments = pieces.iter().filter_map(Piece::as_expr);
255 let [arguments_iter_1, arguments_iter_2] =
256 array::from_fn(|_| pieces.iter().filter_map(Piece::as_ident));
257
258 let argument_count = arguments_iter_1.count();
259 let [const_args_1, const_args_2, const_args_3, const_args_4, const_args_5] =
260 array::from_fn(|_| (0..argument_count).map(|i| format_ident!("N{i}")));
261
262 let return_adder = const_args_1.fold(
263 quote!(StrBaseLen),
264 |current, ident| quote!(RunAdd<#current, U<#ident>>),
265 );
266
267 let return_close_angle_braces = (0..argument_count).map(|_| <Token![>]>::default());
268
269 quote!({
270 use ::aformat::{ArrayString, ToArrayString, __internal::*};
271
272 #[allow(non_snake_case, clippy::too_many_arguments)]
273 fn aformat_inner<StrBaseLen, #(const #const_args_2: usize),*>(
274 #(#arguments_iter_2: ArrayString<#const_args_3>),*
275 ) -> RunTypeToArrayString<#return_adder>
276 where
277 Const<#str_base_len>: ToUInt<Output = StrBaseLen>,
278 #(Const<#const_args_4>: ToUInt,)*
279 StrBaseLen: #(Add<U<#const_args_5>, Output: )* TypeNumToArrayString #(#return_close_angle_braces)*
280 {
281 let mut out = ArrayStringLike::new();
282 if false { return out; }
284
285 #(#pieces)*
286 out
287 }
288
289 aformat_inner(#(ToArrayString::to_arraystring(#caller_arguments)),*)
290 })
291 .into()
292}
293
294#[proc_macro]
322pub fn aformat_into(tokens: proc_macro::TokenStream) -> proc_macro::TokenStream {
323 let FormatIntoArguments {
324 arguments: Arguments {
325 str_base_len,
326 pieces,
327 },
328 write_into,
329 } = match syn::parse(tokens) {
330 Ok(args) => args,
331 Err(err) => return err.into_compile_error().into(),
332 };
333
334 let caller_arguments = pieces.iter().filter_map(Piece::as_expr);
335 let [arguments_iter_1, arguments_iter_2] =
336 array::from_fn(|_| pieces.iter().filter_map(Piece::as_ident));
337
338 let argument_count = arguments_iter_1.clone().count();
339 let [const_args_1, const_args_2, const_args_3, const_args_4] =
340 array::from_fn(|_| (0..argument_count).map(|i| format_ident!("N{i}")));
341
342 let return_close_angle_braces = (0..argument_count).map(|_| <Token![>]>::default());
343
344 quote!({
345 use ::aformat::{ArrayString, ToArrayString, __internal::*};
346
347 #[allow(non_snake_case, clippy::too_many_arguments)]
348 fn aformat_into_inner<StrBaseLen, const OUT: usize, #(const #const_args_2: usize),*>(
349 out: &mut ArrayString<OUT>,
350 #(#arguments_iter_2: ArrayString<#const_args_3>),*
351 )
352 where
353 Const<#str_base_len>: ToUInt<Output = StrBaseLen>,
354 Const<OUT>: ToUInt,
355 #(Const<#const_args_4>: ToUInt,)*
356
357 StrBaseLen: #(Add<U<#const_args_1>, Output: )* IsLessOrEqual<U<OUT>, Output: BufferFits> #(#return_close_angle_braces)*
358 {
359 #(#pieces)*
360 }
361
362 aformat_into_inner(&mut #write_into, #(ToArrayString::to_arraystring(#caller_arguments)),*)
363 })
364 .into()
365}