rust/macros/helpers.rs

   1 // SPDX-License-Identifier: GPL-2.0
   2
   3 use proc_macro::{token_stream, Group, TokenStream, TokenTree};
   4
   5 pub(crate) fn try_ident(it: &mut token_stream::IntoIter) -> Option<String> {
   6     if let Some(TokenTree::Ident(ident)) = it.next() {
   7         Some(ident.to_string())
   8     } else {
   9         None
  10     }
  11 }
  12
  13 pub(crate) fn try_literal(it: &mut token_stream::IntoIter) -> Option<String> {
  14     if let Some(TokenTree::Literal(literal)) = it.next() {
  15         Some(literal.to_string())
  16     } else {
  17         None
  18     }
  19 }
  20
  21 pub(crate) fn try_string(it: &mut token_stream::IntoIter) -> Option<String> {
  22     try_literal(it).and_then(|string| {
  23         if string.starts_with('\"') && string.ends_with('\"') {
  24             let content = &string[1..string.len() - 1];
  25             if content.contains('\\') {
  26                 panic!("Escape sequences in string literals not yet handled");
  27             }
  28             Some(content.to_string())
  29         } else if string.starts_with("r\"") {
  30             panic!("Raw string literals are not yet handled");
  31         } else {
  32             None
  33         }
  34     })
  35 }
  36
  37 pub(crate) fn expect_ident(it: &mut token_stream::IntoIter) -> String {
  38     try_ident(it).expect("Expected Ident")
  39 }
  40
  41 pub(crate) fn expect_punct(it: &mut token_stream::IntoIter) -> char {
  42     if let TokenTree::Punct(punct) = it.next().expect("Reached end of token stream for Punct") {
  43         punct.as_char()
  44     } else {
  45         panic!("Expected Punct");
  46     }
  47 }
  48
  49 pub(crate) fn expect_string(it: &mut token_stream::IntoIter) -> String {
  50     try_string(it).expect("Expected string")
  51 }
  52
  53 pub(crate) fn expect_string_ascii(it: &mut token_stream::IntoIter) -> String {
  54     let string = try_string(it).expect("Expected string");
  55     assert!(string.is_ascii(), "Expected ASCII string");
  56     string
  57 }
  58
  59 pub(crate) fn expect_group(it: &mut token_stream::IntoIter) -> Group {
  60     if let TokenTree::Group(group) = it.next().expect("Reached end of token stream for Group") {
  61         group
  62     } else {
  63         panic!("Expected Group");
  64     }
  65 }
  66
  67 pub(crate) fn expect_end(it: &mut token_stream::IntoIter) {
  68     if it.next().is_some() {
  69         panic!("Expected end");
  70     }
  71 }
  72
  73 /// Parsed generics.
  74 ///
  75 /// See the field documentation for an explanation what each of the fields represents.
  76 ///
  77 /// # Examples
  78 ///
  79 /// ```rust,ignore
  80 /// # let input = todo!();
  81 /// let (Generics { decl_generics, impl_generics, ty_generics }, rest) = parse_generics(input);
  82 /// quote! {
  83 ///     struct Foo<$($decl_generics)*> {
  84 ///         // ...
  85 ///     }
  86 ///
  87 ///     impl<$impl_generics> Foo<$ty_generics> {
  88 ///         fn foo() {
  89 ///             // ...
  90 ///         }
  91 ///     }
  92 /// }
  93 /// ```
  94 pub(crate) struct Generics {
  95     /// The generics with bounds and default values (e.g. `T: Clone, const N: usize = 0`).
  96     ///
  97     /// Use this on type definitions e.g. `struct Foo<$decl_generics> ...` (or `union`/`enum`).
  98     pub(crate) decl_generics: Vec<TokenTree>,
  99     /// The generics with bounds (e.g. `T: Clone, const N: usize`).
 100     ///
 101     /// Use this on `impl` blocks e.g. `impl<$impl_generics> Trait for ...`.
 102     pub(crate) impl_generics: Vec<TokenTree>,
 103     /// The generics without bounds and without default values (e.g. `T, N`).
 104     ///
 105     /// Use this when you use the type that is declared with these generics e.g.
 106     /// `Foo<$ty_generics>`.
 107     pub(crate) ty_generics: Vec<TokenTree>,
 108 }
 109
 110 /// Parses the given `TokenStream` into `Generics` and the rest.
 111 ///
 112 /// The generics are not present in the rest, but a where clause might remain.
 113 pub(crate) fn parse_generics(input: TokenStream) -> (Generics, Vec<TokenTree>) {
 114     // The generics with bounds and default values.
 115     let mut decl_generics = vec![];
 116     // `impl_generics`, the declared generics with their bounds.
 117     let mut impl_generics = vec![];
 118     // Only the names of the generics, without any bounds.
 119     let mut ty_generics = vec![];
 120     // Tokens not related to the generics e.g. the `where` token and definition.
 121     let mut rest = vec![];
 122     // The current level of `<`.
 123     let mut nesting = 0;
 124     let mut toks = input.into_iter();
 125     // If we are at the beginning of a generic parameter.
 126     let mut at_start = true;
 127     let mut skip_until_comma = false;
 128     while let Some(tt) = toks.next() {
 129         if nesting == 1 && matches!(&tt, TokenTree::Punct(p) if p.as_char() == '>') {
 130             // Found the end of the generics.
 131             break;
 132         } else if nesting >= 1 {
 133             decl_generics.push(tt.clone());
 134         }
 135         match tt.clone() {
 136             TokenTree::Punct(p) if p.as_char() == '<' => {
 137                 if nesting >= 1 && !skip_until_comma {
 138                     // This is inside of the generics and part of some bound.
 139                     impl_generics.push(tt);
 140                 }
 141                 nesting += 1;
 142             }
 143             TokenTree::Punct(p) if p.as_char() == '>' => {
 144                 // This is a parsing error, so we just end it here.
 145                 if nesting == 0 {
 146                     break;
 147                 } else {
 148                     nesting -= 1;
 149                     if nesting >= 1 && !skip_until_comma {
 150                         // We are still inside of the generics and part of some bound.
 151                         impl_generics.push(tt);
 152                     }
 153                 }
 154             }
 155             TokenTree::Punct(p) if skip_until_comma && p.as_char() == ',' => {
 156                 if nesting == 1 {
 157                     impl_generics.push(tt.clone());
 158                     impl_generics.push(tt);
 159                     skip_until_comma = false;
 160                 }
 161             }
 162             _ if !skip_until_comma => {
 163                 match nesting {
 164                     // If we haven't entered the generics yet, we still want to keep these tokens.
 165                     0 => rest.push(tt),
 166                     1 => {
 167                         // Here depending on the token, it might be a generic variable name.
 168                         match tt.clone() {
 169                             TokenTree::Ident(i) if at_start && i.to_string() == "const" => {
 170                                 let Some(name) = toks.next() else {
 171                                     // Parsing error.
 172                                     break;
 173                                 };
 174                                 impl_generics.push(tt);
 175                                 impl_generics.push(name.clone());
 176                                 ty_generics.push(name.clone());
 177                                 decl_generics.push(name);
 178                                 at_start = false;
 179                             }
 180                             TokenTree::Ident(_) if at_start => {
 181                                 impl_generics.push(tt.clone());
 182                                 ty_generics.push(tt);
 183                                 at_start = false;
 184                             }
 185                             TokenTree::Punct(p) if p.as_char() == ',' => {
 186                                 impl_generics.push(tt.clone());
 187                                 ty_generics.push(tt);
 188                                 at_start = true;
 189                             }
 190                             // Lifetimes begin with `'`.
 191                             TokenTree::Punct(p) if p.as_char() == '\'' && at_start => {
 192                                 impl_generics.push(tt.clone());
 193                                 ty_generics.push(tt);
 194                             }
 195                             // Generics can have default values, we skip these.
 196                             TokenTree::Punct(p) if p.as_char() == '=' => {
 197                                 skip_until_comma = true;
 198                             }
 199                             _ => impl_generics.push(tt),
 200                         }
 201                     }
 202                     _ => impl_generics.push(tt),
 203                 }
 204             }
 205             _ => {}
 206         }
 207     }
 208     rest.extend(toks);
 209     (
 210         Generics {
 211             impl_generics,
 212             decl_generics,
 213             ty_generics,
 214         },
 215         rest,
 216     )
 217 }