pin_data.rs

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