xref: /linux/rust/macros/paste.rs (revision e467705a9fb37f51595aa6deaca085ccb4005454)
1 // SPDX-License-Identifier: GPL-2.0
2 
3 use proc_macro::{Delimiter, Group, Ident, Spacing, Span, TokenTree};
4 
5 fn concat(tokens: &[TokenTree], group_span: Span) -> TokenTree {
6     let mut tokens = tokens.iter();
7     let mut segments = Vec::new();
8     let mut span = None;
9     loop {
10         match tokens.next() {
11             None => break,
12             Some(TokenTree::Literal(lit)) => {
13                 // Allow us to concat string literals by stripping quotes
14                 let mut value = lit.to_string();
15                 if value.starts_with('"') && value.ends_with('"') {
16                     value.remove(0);
17                     value.pop();
18                 }
19                 segments.push((value, lit.span()));
20             }
21             Some(TokenTree::Ident(ident)) => {
22                 let mut value = ident.to_string();
23                 if value.starts_with("r#") {
24                     value.replace_range(0..2, "");
25                 }
26                 segments.push((value, ident.span()));
27             }
28             Some(TokenTree::Punct(p)) if p.as_char() == ':' => {
29                 let Some(TokenTree::Ident(ident)) = tokens.next() else {
30                     panic!("expected identifier as modifier");
31                 };
32 
33                 let (mut value, sp) = segments.pop().expect("expected identifier before modifier");
34                 match ident.to_string().as_str() {
35                     // Set the overall span of concatenated token as current span
36                     "span" => {
37                         assert!(
38                             span.is_none(),
39                             "span modifier should only appear at most once"
40                         );
41                         span = Some(sp);
42                     }
43                     "lower" => value = value.to_lowercase(),
44                     "upper" => value = value.to_uppercase(),
45                     v => panic!("unknown modifier `{v}`"),
46                 };
47                 segments.push((value, sp));
48             }
49             _ => panic!("unexpected token in paste segments"),
50         };
51     }
52 
53     let pasted: String = segments.into_iter().map(|x| x.0).collect();
54     TokenTree::Ident(Ident::new(&pasted, span.unwrap_or(group_span)))
55 }
56 
57 pub(crate) fn expand(tokens: &mut Vec<TokenTree>) {
58     for token in tokens.iter_mut() {
59         if let TokenTree::Group(group) = token {
60             let delimiter = group.delimiter();
61             let span = group.span();
62             let mut stream: Vec<_> = group.stream().into_iter().collect();
63             // Find groups that looks like `[< A B C D >]`
64             if delimiter == Delimiter::Bracket
65                 && stream.len() >= 3
66                 && matches!(&stream[0], TokenTree::Punct(p) if p.as_char() == '<')
67                 && matches!(&stream[stream.len() - 1], TokenTree::Punct(p) if p.as_char() == '>')
68             {
69                 // Replace the group with concatenated token
70                 *token = concat(&stream[1..stream.len() - 1], span);
71             } else {
72                 // Recursively expand tokens inside the group
73                 expand(&mut stream);
74                 let mut group = Group::new(delimiter, stream.into_iter().collect());
75                 group.set_span(span);
76                 *token = TokenTree::Group(group);
77             }
78         }
79     }
80 
81     // Path segments cannot contain invisible delimiter group, so remove them if any.
82     for i in (0..tokens.len().saturating_sub(3)).rev() {
83         // Looking for a double colon
84         if matches!(
85             (&tokens[i + 1], &tokens[i + 2]),
86             (TokenTree::Punct(a), TokenTree::Punct(b))
87                 if a.as_char() == ':' && a.spacing() == Spacing::Joint && b.as_char() == ':'
88         ) {
89             match &tokens[i + 3] {
90                 TokenTree::Group(group) if group.delimiter() == Delimiter::None => {
91                     tokens.splice(i + 3..i + 4, group.stream());
92                 }
93                 _ => (),
94             }
95 
96             match &tokens[i] {
97                 TokenTree::Group(group) if group.delimiter() == Delimiter::None => {
98                     tokens.splice(i..i + 1, group.stream());
99                 }
100                 _ => (),
101             }
102         }
103     }
104 }
105