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