summaryrefslogtreecommitdiff
path: root/rust/macros/helpers.rs
blob: 563dcd2b7ace5e8322d0fddb409571cca2dd31ea (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
// SPDX-License-Identifier: GPL-2.0

use proc_macro::{token_stream, Group, TokenStream, TokenTree};

pub(crate) fn try_ident(it: &mut token_stream::IntoIter) -> Option<String> {
    if let Some(TokenTree::Ident(ident)) = it.next() {
        Some(ident.to_string())
    } else {
        None
    }
}

pub(crate) fn try_literal(it: &mut token_stream::IntoIter) -> Option<String> {
    if let Some(TokenTree::Literal(literal)) = it.next() {
        Some(literal.to_string())
    } else {
        None
    }
}

pub(crate) fn try_string(it: &mut token_stream::IntoIter) -> Option<String> {
    try_literal(it).and_then(|string| {
        if string.starts_with('\"') && string.ends_with('\"') {
            let content = &string[1..string.len() - 1];
            if content.contains('\\') {
                panic!("Escape sequences in string literals not yet handled");
            }
            Some(content.to_string())
        } else if string.starts_with("r\"") {
            panic!("Raw string literals are not yet handled");
        } else {
            None
        }
    })
}

pub(crate) fn expect_ident(it: &mut token_stream::IntoIter) -> String {
    try_ident(it).expect("Expected Ident")
}

pub(crate) fn expect_punct(it: &mut token_stream::IntoIter) -> char {
    if let TokenTree::Punct(punct) = it.next().expect("Reached end of token stream for Punct") {
        punct.as_char()
    } else {
        panic!("Expected Punct");
    }
}

pub(crate) fn expect_string(it: &mut token_stream::IntoIter) -> String {
    try_string(it).expect("Expected string")
}

pub(crate) fn expect_string_ascii(it: &mut token_stream::IntoIter) -> String {
    let string = try_string(it).expect("Expected string");
    assert!(string.is_ascii(), "Expected ASCII string");
    string
}

pub(crate) fn expect_group(it: &mut token_stream::IntoIter) -> Group {
    if let TokenTree::Group(group) = it.next().expect("Reached end of token stream for Group") {
        group
    } else {
        panic!("Expected Group");
    }
}

pub(crate) fn expect_end(it: &mut token_stream::IntoIter) {
    if it.next().is_some() {
        panic!("Expected end");
    }
}

/// Parsed generics.
///
/// See the field documentation for an explanation what each of the fields represents.
///
/// # Examples
///
/// ```rust,ignore
/// # let input = todo!();
/// let (Generics { decl_generics, impl_generics, ty_generics }, rest) = parse_generics(input);
/// quote! {
///     struct Foo<$($decl_generics)*> {
///         // ...
///     }
///
///     impl<$impl_generics> Foo<$ty_generics> {
///         fn foo() {
///             // ...
///         }
///     }
/// }
/// ```
pub(crate) struct Generics {
    /// The generics with bounds and default values (e.g. `T: Clone, const N: usize = 0`).
    ///
    /// Use this on type definitions e.g. `struct Foo<$decl_generics> ...` (or `union`/`enum`).
    pub(crate) decl_generics: Vec<TokenTree>,
    /// The generics with bounds (e.g. `T: Clone, const N: usize`).
    ///
    /// Use this on `impl` blocks e.g. `impl<$impl_generics> Trait for ...`.
    pub(crate) impl_generics: Vec<TokenTree>,
    /// The generics without bounds and without default values (e.g. `T, N`).
    ///
    /// Use this when you use the type that is declared with these generics e.g.
    /// `Foo<$ty_generics>`.
    pub(crate) ty_generics: Vec<TokenTree>,
}

/// Parses the given `TokenStream` into `Generics` and the rest.
///
/// The generics are not present in the rest, but a where clause might remain.
pub(crate) fn parse_generics(input: TokenStream) -> (Generics, Vec<TokenTree>) {
    // The generics with bounds and default values.
    let mut decl_generics = vec![];
    // `impl_generics`, the declared generics with their bounds.
    let mut impl_generics = vec![];
    // Only the names of the generics, without any bounds.
    let mut ty_generics = vec![];
    // Tokens not related to the generics e.g. the `where` token and definition.
    let mut rest = vec![];
    // The current level of `<`.
    let mut nesting = 0;
    let mut toks = input.into_iter();
    // If we are at the beginning of a generic parameter.
    let mut at_start = true;
    let mut skip_until_comma = false;
    while let Some(tt) = toks.next() {
        if nesting == 1 && matches!(&tt, TokenTree::Punct(p) if p.as_char() == '>') {
            // Found the end of the generics.
            break;
        } else if nesting >= 1 {
            decl_generics.push(tt.clone());
        }
        match tt.clone() {
            TokenTree::Punct(p) if p.as_char() == '<' => {
                if nesting >= 1 && !skip_until_comma {
                    // This is inside of the generics and part of some bound.
                    impl_generics.push(tt);
                }
                nesting += 1;
            }
            TokenTree::Punct(p) if p.as_char() == '>' => {
                // This is a parsing error, so we just end it here.
                if nesting == 0 {
                    break;
                } else {
                    nesting -= 1;
                    if nesting >= 1 && !skip_until_comma {
                        // We are still inside of the generics and part of some bound.
                        impl_generics.push(tt);
                    }
                }
            }
            TokenTree::Punct(p) if skip_until_comma && p.as_char() == ',' => {
                if nesting == 1 {
                    impl_generics.push(tt.clone());
                    impl_generics.push(tt);
                    skip_until_comma = false;
                }
            }
            _ if !skip_until_comma => {
                match nesting {
                    // If we haven't entered the generics yet, we still want to keep these tokens.
                    0 => rest.push(tt),
                    1 => {
                        // Here depending on the token, it might be a generic variable name.
                        match tt.clone() {
                            TokenTree::Ident(i) if at_start && i.to_string() == "const" => {
                                let Some(name) = toks.next() else {
                                    // Parsing error.
                                    break;
                                };
                                impl_generics.push(tt);
                                impl_generics.push(name.clone());
                                ty_generics.push(name.clone());
                                decl_generics.push(name);
                                at_start = false;
                            }
                            TokenTree::Ident(_) if at_start => {
                                impl_generics.push(tt.clone());
                                ty_generics.push(tt);
                                at_start = false;
                            }
                            TokenTree::Punct(p) if p.as_char() == ',' => {
                                impl_generics.push(tt.clone());
                                ty_generics.push(tt);
                                at_start = true;
                            }
                            // Lifetimes begin with `'`.
                            TokenTree::Punct(p) if p.as_char() == '\'' && at_start => {
                                impl_generics.push(tt.clone());
                                ty_generics.push(tt);
                            }
                            // Generics can have default values, we skip these.
                            TokenTree::Punct(p) if p.as_char() == '=' => {
                                skip_until_comma = true;
                            }
                            _ => impl_generics.push(tt),
                        }
                    }
                    _ => impl_generics.push(tt),
                }
            }
            _ => {}
        }
    }
    rest.extend(toks);
    (
        Generics {
            impl_generics,
            decl_generics,
            ty_generics,
        },
        rest,
    )
}