macros/
fmt.rs

1// SPDX-License-Identifier: GPL-2.0
2
3use proc_macro::{Ident, TokenStream, TokenTree};
4use std::collections::BTreeSet;
5
6/// Please see [`crate::fmt`] for documentation.
7pub(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}