aformat_macros/
lib.rs

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
34/// A piece of the formatted string.
35enum Piece {
36    /// A literal `&'static str` piece.
37    Literal(ByteString),
38    /// A argument that should be passed to `ToArrayString`, then into `aformat_inner`.
39    Argument {
40        /// The expression that produces the argument value.
41        expr: Expr,
42        /// The name of the function argument inside `aformat_inner`
43        ident: Ident,
44    },
45    /// A duplicate argument that was already passed to `aformat_inner`.
46    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/// A no-alloc version of [`format!`], producing an [`ArrayString`].
235///
236/// ## Usage
237/// Usage is similar to `format!`, although there are multiple limitations:
238/// - No support for formatting flags, as we are not reinventing all of `format!`.
239/// - No support for `format!("{name}", name=username)` syntax, may be lifted in future.
240///
241// Workaround for a rustdoc bug.
242/// [`format!`]: https://doc.rust-lang.org/stable/std/macro.format.html
243/// [`ArrayString`]: https://docs.rs/arrayvec/latest/arrayvec/struct.ArrayString.html
244#[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            // Fixes type inferrence
283            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/// [`aformat!`], but you provide your own [`ArrayString`].
295///
296/// The length of the [`ArrayString`] is checked at compile-time to fit all the arguments, although is not checked to be optimal.
297///
298/// ## Usage
299/// The first argument should be the identifier of the [`ArrayString`], then the normal [`aformat!`] arguments follow.
300///
301/// ## Examples
302/// ```
303/// let mut out_buf = ArrayString::<32>::new();
304///
305/// let age = 18_u8;
306/// aformat_into!(out_buf, "You are {} years old!", age);
307///
308/// assert_eq!(out_buf.as_str(), "You are 18 years old!");
309/// ```
310///
311/// ```compile_fail
312/// // Buffer is too small, so compile failure!
313/// let mut out_buf = ArrayString::<4>::new();
314///
315/// let age = 18_u8;
316/// aformat_into!(out_buf, "You are {} years old!", age);
317/// ```
318///
319// Workaround for a rustdoc bug.
320/// [`ArrayString`]: https://docs.rs/arrayvec/latest/arrayvec/struct.ArrayString.html
321#[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}
aformat_macros/lib.rs

aformat_macros/
lib.rs
