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

// SPDX-License-Identifier: Apache-2.0 OR MIT

#[cfg(not(kernel))]
use proc_macro2 as proc_macro;

use proc_macro::{TokenStream, TokenTree};

/// 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,
    )
}