xref: /linux/rust/macros/helpers.rs (revision 460ea8980511c01c1551012b9a6ec6a06d02da59)
1 // SPDX-License-Identifier: GPL-2.0
2 
3 use proc_macro::{token_stream, Group, Punct, Spacing, 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 pub(crate) struct Generics {
74     pub(crate) impl_generics: Vec<TokenTree>,
75     pub(crate) ty_generics: Vec<TokenTree>,
76 }
77 
78 /// Parses the given `TokenStream` into `Generics` and the rest.
79 ///
80 /// The generics are not present in the rest, but a where clause might remain.
81 pub(crate) fn parse_generics(input: TokenStream) -> (Generics, Vec<TokenTree>) {
82     // `impl_generics`, the declared generics with their bounds.
83     let mut impl_generics = vec![];
84     // Only the names of the generics, without any bounds.
85     let mut ty_generics = vec![];
86     // Tokens not related to the generics e.g. the `where` token and definition.
87     let mut rest = vec![];
88     // The current level of `<`.
89     let mut nesting = 0;
90     let mut toks = input.into_iter();
91     // If we are at the beginning of a generic parameter.
92     let mut at_start = true;
93     for tt in &mut toks {
94         match tt.clone() {
95             TokenTree::Punct(p) if p.as_char() == '<' => {
96                 if nesting >= 1 {
97                     // This is inside of the generics and part of some bound.
98                     impl_generics.push(tt);
99                 }
100                 nesting += 1;
101             }
102             TokenTree::Punct(p) if p.as_char() == '>' => {
103                 // This is a parsing error, so we just end it here.
104                 if nesting == 0 {
105                     break;
106                 } else {
107                     nesting -= 1;
108                     if nesting >= 1 {
109                         // We are still inside of the generics and part of some bound.
110                         impl_generics.push(tt);
111                     }
112                     if nesting == 0 {
113                         break;
114                     }
115                 }
116             }
117             tt => {
118                 if nesting == 1 {
119                     // Here depending on the token, it might be a generic variable name.
120                     match &tt {
121                         // Ignore const.
122                         TokenTree::Ident(i) if i.to_string() == "const" => {}
123                         TokenTree::Ident(_) if at_start => {
124                             ty_generics.push(tt.clone());
125                             // We also already push the `,` token, this makes it easier to append
126                             // generics.
127                             ty_generics.push(TokenTree::Punct(Punct::new(',', Spacing::Alone)));
128                             at_start = false;
129                         }
130                         TokenTree::Punct(p) if p.as_char() == ',' => at_start = true,
131                         // Lifetimes begin with `'`.
132                         TokenTree::Punct(p) if p.as_char() == '\'' && at_start => {
133                             ty_generics.push(tt.clone());
134                         }
135                         _ => {}
136                     }
137                 }
138                 if nesting >= 1 {
139                     impl_generics.push(tt);
140                 } else if nesting == 0 {
141                     // If we haven't entered the generics yet, we still want to keep these tokens.
142                     rest.push(tt);
143                 }
144             }
145         }
146     }
147     rest.extend(toks);
148     (
149         Generics {
150             impl_generics,
151             ty_generics,
152         },
153         rest,
154     )
155 }
156