1
// SPDX-License-Identifier: GPL-2.0
2

3
use std::collections::BTreeSet;
4

5
use proc_macro2::{Ident, TokenStream, TokenTree};
6
use quote::quote_spanned;
7

8
/// Please see [`crate::fmt`] for documentation.
9
pub(crate) fn fmt(input: TokenStream) -> TokenStream {
10
    let mut input = input.into_iter();
11

12
    let first_opt = input.next();
13
    let first_owned_str;
14
    let mut names = BTreeSet::new();
15
    let first_span = {
16
        let Some((mut first_str, first_span)) = (match first_opt.as_ref() {
17
            Some(TokenTree::Literal(first_lit)) => {
18
                first_owned_str = first_lit.to_string();
19
                Some(first_owned_str.as_str()).and_then(|first| {
20
                    let first = first.strip_prefix('"')?;
21
                    let first = first.strip_suffix('"')?;
22
                    Some((first, first_lit.span()))
23
                })
24
            }
25
            _ => None,
26
        }) else {
27
            return first_opt.into_iter().chain(input).collect();
28
        };
29

30
        // Parse `identifier`s from the format string.
31
        //
32
        // See https://doc.rust-lang.org/std/fmt/index.html#syntax.
33
        while let Some((_, rest)) = first_str.split_once('{') {
34
            first_str = rest;
35
            if let Some(rest) = first_str.strip_prefix('{') {
36
                first_str = rest;
37
                continue;
38
            }
39
            if let Some((name, rest)) = first_str.split_once('}') {
40
                first_str = rest;
41
                let name = name.split_once(':').map_or(name, |(name, _)| name);
42
                if !name.is_empty() && !name.chars().all(|c| c.is_ascii_digit()) {
43
                    names.insert(name);
44
                }
45
            }
46
        }
47
        first_span
48
    };
49

50
    let adapter = quote_spanned!(first_span => ::kernel::fmt::Adapter);
51

52
    let mut args = TokenStream::from_iter(first_opt);
53
    {
54
        let mut flush = |args: &mut TokenStream, current: &mut TokenStream| {
55
            let current = std::mem::take(current);
56
            if !current.is_empty() {
57
                let (lhs, rhs) = (|| {
58
                    let mut current = current.into_iter();
59
                    let mut acc = TokenStream::new();
60
                    while let Some(tt) = current.next() {
61
                        // Split on `=` only once to handle cases like `a = b = c`.
62
                        if matches!(&tt, TokenTree::Punct(p) if p.as_char() == '=') {
63
                            names.remove(acc.to_string().as_str());
64
                            // Include the `=` itself to keep the handling below uniform.
65
                            acc.extend([tt]);
66
                            return (Some(acc), current.collect::<TokenStream>());
67
                        }
68
                        acc.extend([tt]);
69
                    }
70
                    (None, acc)
71
                })();
72
                args.extend(quote_spanned!(first_span => #lhs #adapter(&(#rhs))));
73
            }
74
        };
75

76
        let mut current = TokenStream::new();
77
        for tt in input {
78
            match &tt {
79
                TokenTree::Punct(p) if p.as_char() == ',' => {
80
                    flush(&mut args, &mut current);
81
                    &mut args
82
                }
83
                _ => &mut current,
84
            }
85
            .extend([tt]);
86
        }
87
        flush(&mut args, &mut current);
88
    }
89

90
    for name in names {
91
        let name = Ident::new(name, first_span);
92
        args.extend(quote_spanned!(first_span => , #name = #adapter(&#name)));
93
    }
94

95
    quote_spanned!(first_span => ::core::format_args!(#args))
96
}
97

98