xref: /linux/rust/macros/fmt.rs (revision 808c999fc9e7c366fd47da564e69d579c1dc8279)
1 // SPDX-License-Identifier: GPL-2.0
2 
3 use proc_macro::{Ident, TokenStream, TokenTree};
4 use std::collections::BTreeSet;
5 
6 /// Please see [`crate::fmt`] for documentation.
7 pub(crate) fn fmt(input: TokenStream) -> TokenStream {
8     let mut input = input.into_iter();
9 
10     let first_opt = input.next();
11     let first_owned_str;
12     let mut names = BTreeSet::new();
13     let first_span = {
14         let Some((mut first_str, first_span)) = (match first_opt.as_ref() {
15             Some(TokenTree::Literal(first_lit)) => {
16                 first_owned_str = first_lit.to_string();
17                 Some(first_owned_str.as_str()).and_then(|first| {
18                     let first = first.strip_prefix('"')?;
19                     let first = first.strip_suffix('"')?;
20                     Some((first, first_lit.span()))
21                 })
22             }
23             _ => None,
24         }) else {
25             return first_opt.into_iter().chain(input).collect();
26         };
27 
28         // Parse `identifier`s from the format string.
29         //
30         // See https://doc.rust-lang.org/std/fmt/index.html#syntax.
31         while let Some((_, rest)) = first_str.split_once('{') {
32             first_str = rest;
33             if let Some(rest) = first_str.strip_prefix('{') {
34                 first_str = rest;
35                 continue;
36             }
37             if let Some((name, rest)) = first_str.split_once('}') {
38                 first_str = rest;
39                 let name = name.split_once(':').map_or(name, |(name, _)| name);
40                 if !name.is_empty() && !name.chars().all(|c| c.is_ascii_digit()) {
41                     names.insert(name);
42                 }
43             }
44         }
45         first_span
46     };
47 
48     let adapter = quote_spanned!(first_span => ::kernel::fmt::Adapter);
49 
50     let mut args = TokenStream::from_iter(first_opt);
51     {
52         let mut flush = |args: &mut TokenStream, current: &mut TokenStream| {
53             let current = std::mem::take(current);
54             if !current.is_empty() {
55                 let (lhs, rhs) = (|| {
56                     let mut current = current.into_iter();
57                     let mut acc = TokenStream::new();
58                     while let Some(tt) = current.next() {
59                         // Split on `=` only once to handle cases like `a = b = c`.
60                         if matches!(&tt, TokenTree::Punct(p) if p.as_char() == '=') {
61                             names.remove(acc.to_string().as_str());
62                             // Include the `=` itself to keep the handling below uniform.
63                             acc.extend([tt]);
64                             return (Some(acc), current.collect::<TokenStream>());
65                         }
66                         acc.extend([tt]);
67                     }
68                     (None, acc)
69                 })();
70                 args.extend(quote_spanned!(first_span => #lhs #adapter(&#rhs)));
71             }
72         };
73 
74         let mut current = TokenStream::new();
75         for tt in input {
76             match &tt {
77                 TokenTree::Punct(p) if p.as_char() == ',' => {
78                     flush(&mut args, &mut current);
79                     &mut args
80                 }
81                 _ => &mut current,
82             }
83             .extend([tt]);
84         }
85         flush(&mut args, &mut current);
86     }
87 
88     for name in names {
89         let name = Ident::new(name, first_span);
90         args.extend(quote_spanned!(first_span => , #name = #adapter(&#name)));
91     }
92 
93     quote_spanned!(first_span => ::core::format_args!(#args))
94 }
95