xref: /linux/rust/pin-init/internal/src/pin_data.rs (revision 50426bde1577d17e61362bd199d487dbeb159110)
1 // SPDX-License-Identifier: Apache-2.0 OR MIT
2 
3 use crate::helpers::{parse_generics, Generics};
4 use proc_macro2::{Group, Punct, Spacing, TokenStream, TokenTree};
5 use quote::quote;
6 
7 pub(crate) fn pin_data(args: TokenStream, input: TokenStream) -> TokenStream {
8     // This proc-macro only does some pre-parsing and then delegates the actual parsing to
9     // `pin_init::__pin_data!`.
10 
11     let (
12         Generics {
13             impl_generics,
14             decl_generics,
15             ty_generics,
16         },
17         rest,
18     ) = parse_generics(input);
19     // The struct definition might contain the `Self` type. Since `__pin_data!` will define a new
20     // type with the same generics and bounds, this poses a problem, since `Self` will refer to the
21     // new type as opposed to this struct definition. Therefore we have to replace `Self` with the
22     // concrete name.
23 
24     // Errors that occur when replacing `Self` with `struct_name`.
25     let mut errs = TokenStream::new();
26     // The name of the struct with ty_generics.
27     let struct_name = rest
28         .iter()
29         .skip_while(|tt| !matches!(tt, TokenTree::Ident(i) if i == "struct"))
30         .nth(1)
31         .and_then(|tt| match tt {
32             TokenTree::Ident(_) => {
33                 let tt = tt.clone();
34                 let mut res = vec![tt];
35                 if !ty_generics.is_empty() {
36                     // We add this, so it is maximally compatible with e.g. `Self::CONST` which
37                     // will be replaced by `StructName::<$generics>::CONST`.
38                     res.push(TokenTree::Punct(Punct::new(':', Spacing::Joint)));
39                     res.push(TokenTree::Punct(Punct::new(':', Spacing::Alone)));
40                     res.push(TokenTree::Punct(Punct::new('<', Spacing::Alone)));
41                     res.extend(ty_generics.iter().cloned());
42                     res.push(TokenTree::Punct(Punct::new('>', Spacing::Alone)));
43                 }
44                 Some(res)
45             }
46             _ => None,
47         })
48         .unwrap_or_else(|| {
49             // If we did not find the name of the struct then we will use `Self` as the replacement
50             // and add a compile error to ensure it does not compile.
51             errs.extend(
52                 "::core::compile_error!(\"Could not locate type name.\");"
53                     .parse::<TokenStream>()
54                     .unwrap(),
55             );
56             "Self".parse::<TokenStream>().unwrap().into_iter().collect()
57         });
58     let impl_generics = impl_generics
59         .into_iter()
60         .flat_map(|tt| replace_self_and_deny_type_defs(&struct_name, tt, &mut errs))
61         .collect::<Vec<_>>();
62     let mut rest = rest
63         .into_iter()
64         .flat_map(|tt| {
65             // We ignore top level `struct` tokens, since they would emit a compile error.
66             if matches!(&tt, TokenTree::Ident(i) if i == "struct") {
67                 vec![tt]
68             } else {
69                 replace_self_and_deny_type_defs(&struct_name, tt, &mut errs)
70             }
71         })
72         .collect::<Vec<_>>();
73     // This should be the body of the struct `{...}`.
74     let last = rest.pop();
75     let mut quoted = quote!(::pin_init::__pin_data! {
76         parse_input:
77         @args(#args),
78         @sig(#(#rest)*),
79         @impl_generics(#(#impl_generics)*),
80         @ty_generics(#(#ty_generics)*),
81         @decl_generics(#(#decl_generics)*),
82         @body(#last),
83     });
84     quoted.extend(errs);
85     quoted
86 }
87 
88 /// Replaces `Self` with `struct_name` and errors on `enum`, `trait`, `struct` `union` and `impl`
89 /// keywords.
90 ///
91 /// The error is appended to `errs` to allow normal parsing to continue.
92 fn replace_self_and_deny_type_defs(
93     struct_name: &Vec<TokenTree>,
94     tt: TokenTree,
95     errs: &mut TokenStream,
96 ) -> Vec<TokenTree> {
97     match tt {
98         TokenTree::Ident(ref i)
99             if i == "enum" || i == "trait" || i == "struct" || i == "union" || i == "impl" =>
100         {
101             errs.extend(
102                 format!(
103                     "::core::compile_error!(\"Cannot use `{i}` inside of struct definition with \
104                         `#[pin_data]`.\");"
105                 )
106                 .parse::<TokenStream>()
107                 .unwrap()
108                 .into_iter()
109                 .map(|mut tok| {
110                     tok.set_span(tt.span());
111                     tok
112                 }),
113             );
114             vec![tt]
115         }
116         TokenTree::Ident(i) if i == "Self" => struct_name.clone(),
117         TokenTree::Literal(_) | TokenTree::Punct(_) | TokenTree::Ident(_) => vec![tt],
118         TokenTree::Group(g) => vec![TokenTree::Group(Group::new(
119             g.delimiter(),
120             g.stream()
121                 .into_iter()
122                 .flat_map(|tt| replace_self_and_deny_type_defs(struct_name, tt, errs))
123                 .collect(),
124         ))],
125     }
126 }
127