rust/macros/paste.rs

   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 }