xref: /linux/rust/macros/helpers.rs (revision c532de5a67a70f8533d495f8f2aaa9a0491c3ad0)
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 }
218