fixup! rust: macros: rewrite #[pinned_drop] using syn

Rust-for-Linux · BennoLossin · Apr 7, 2024 · Apr 7, 2024 · Apr 10, 2024 · Oct 9, 2022
commit 66b9b59688cfe2c21090d592400d6c6f5f829def
diff --git a/rust/macros/lib.rs b/rust/macros/lib.rs
@@ -292,21 +292,9 @@ pub fn pin_data(args: TokenStream, input: TokenStream) -> TokenStream {
 #[proc_macro_attribute]
 pub fn pinned_drop(args: TokenStream, input: TokenStream) -> TokenStream {
     parse_macro_input!(args as syn::parse::Nothing);
-    match syn::parse(input.clone()) {
-        Ok(input) => match pinned_drop::pinned_drop(input) {
-            Ok(output) => output,
-            Err((err, impl_)) => {
-                let err = err.into_compile_error();
-                quote! {
-                    #impl_
-                    #err
-                }
-            }
-        }
-        .into(),
-        // Let the compiler handle the error.
-        Err(_) => input,
-    }
+    pinned_drop::pinned_drop(parse_macro_input!(input))
+        .unwrap_or_else(|e| e.into_compile_error())
+        .into()
 }
 
 /// Paste identifiers together.

diff --git a/rust/macros/pinned_drop.rs b/rust/macros/pinned_drop.rs
@@ -1,112 +1,69 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
 use proc_macro2::TokenStream;
-use quote::{quote, ToTokens};
-use syn::{parse_quote, Error, ImplItem, ImplItemFn, ItemImpl, Token};
+use quote::quote;
+use syn::{parse_quote, spanned::Spanned, Error, ImplItem, ImplItemFn, ItemImpl, Result, Token};
 
-struct PinnedDropImpl(ItemImpl, ImplItemFn);
-
-impl TryFrom<ItemImpl> for PinnedDropImpl {
-    type Error = (Error, TokenStream);
-    fn try_from(mut impl_: ItemImpl) -> Result<Self, Self::Error> {
-        let aux_impl = match &impl_ {
-            ItemImpl {
-                attrs,
-                impl_token,
-                generics,
-                trait_: Some((polarity, path, for_)),
-                self_ty,
-                ..
-            } if path.is_ident("PinnedDrop") => {
-                let (impl_generics, _, whr) = generics.split_for_impl();
-                quote! {
-                    #(#attrs)*
-                    unsafe #impl_token #impl_generics #polarity ::kernel::init::PinnedDrop #for_ #self_ty
-                        #whr
-                    {
-                        fn drop(
-                            self: ::core::pin::Pin<&mut Self>,
-                            _: ::kernel::init::__internal::OnlyCallFromDrop,
-                        ) {}
-                    }
-                }
-            }
-            _ => quote!(#impl_),
-        };
-        if impl_.unsafety.is_some() {
-            return Err((
-                Error::new_spanned(
-                    impl_.unsafety,
-                    "The `PinnedDrop` impl must not be `unsafe`.",
-                ),
-                aux_impl,
-            ));
-        }
-        let trait_tokens = match &impl_.trait_ {
-            None => quote!(),
-            Some((None, path, _)) => quote!(#path),
-            Some((Some(not), path, _)) => quote!(#not #path),
-        };
-        match &impl_.trait_ {
-            Some((None, path, _)) if path.is_ident("PinnedDrop") => {}
-            None | Some((None, ..)) => {
-                return Err((
-                    Error::new_spanned(
-                        trait_tokens,
-                        "`#[pinned_drop]` can only be used on `PinnedDrop` impls.",
-                    ),
-                    aux_impl,
-                ));
-            }
-            Some((Some(_), ..)) => {
-                return Err((
-                    Error::new_spanned(
-                        trait_tokens,
-                        "`#[pinned_drop]` can only be used on positive `PinnedDrop` impls.",
-                    ),
-                    aux_impl,
-                ));
+pub(crate) fn pinned_drop(mut input: ItemImpl) -> Result<TokenStream> {
+    let Some((_, path, _)) = &mut input.trait_ else {
+        return Err(Error::new_spanned(
+            input,
+            "expected an `impl` block implementing `PinnedDrop`",
+        ));
+    };
+    if !is_pinned_drop(&path) {
+        return Err(Error::new_spanned(
+            input,
+            "expected an `impl` block implementing `PinnedDrop`",
+        ));
+    }
+    let mut error = None;
+    if let Some(unsafety) = input.unsafety.take() {
+        error = Some(
+            Error::new_spanned(
+                unsafety,
+                "implementing the trait `PinnedDrop` via `#[pinned_drop]` is not unsafe",
+            )
+            .into_compile_error(),
+        );
+    }
+    input.unsafety = Some(Token![unsafe](input.impl_token.span()));
+    if path.segments.len() != 2 {
+        path.segments.insert(0, parse_quote!(pinned_init));
+    }
+    path.leading_colon.get_or_insert(Token![::](path.span()));
+    for item in &mut input.items {
+        match item {
+            ImplItem::Fn(ImplItemFn { sig, .. }) if sig.ident == "drop" => {
+                sig.inputs
+                    .push(parse_quote!(_: ::kernel::init::__internal::OnlyCallFromDrop));
             }
+            _ => {}
         }
-        if impl_.items.len() != 1 {
-            return Err((
-                Error::new(
-                    impl_.brace_token.span.join(),
-                    "Expected exactly one function in the `PinnedDrop` impl.",
-                ),
-                aux_impl,
-            ));
-        }
-        let ImplItem::Fn(drop) = impl_.items.pop().unwrap() else {
-            return Err((
-                Error::new(impl_.brace_token.span.join(), "Expected a function."),
-                aux_impl,
-            ));
-        };
-
-        Ok(Self(impl_, drop))
     }
+    Ok(quote! {
+        #error
+        #input
+    })
 }
 
-impl ToTokens for PinnedDropImpl {
-    fn to_tokens(&self, tokens: &mut TokenStream) {
-        self.0.to_tokens(tokens)
+fn is_pinned_drop(path: &syn::Path) -> bool {
+    if path.segments.len() > 2 {
+        return false;
     }
-}
-
-pub(crate) fn pinned_drop(drop_impl: ItemImpl) -> Result<TokenStream, (Error, TokenStream)> {
-    let PinnedDropImpl(mut drop_impl, mut drop) = drop_impl.try_into()?;
-    drop.sig
-        .inputs
-        .push(parse_quote!(_: ::kernel::init::__internal::OnlyCallFromDrop));
-    let span = drop_impl.impl_token.span;
-
-    drop_impl.items.push(ImplItem::Fn(drop));
-    drop_impl.unsafety = Some(Token![unsafe](span));
-    drop_impl.trait_ = Some((
-        None,
-        parse_quote!(::kernel::init::PinnedDrop),
-        Token![for](span),
-    ));
-    Ok(quote! { #drop_impl })
+    // If there is a `::`, then the path needs to be `::kernel::init::PinnedDrop`.
+    if path.leading_colon.is_some() && path.segments.len() != 2 {
+        return false;
+    }
+    for (actual, expected) in path
+        .segments
+        .iter()
+        .rev()
+        .zip(["PinnedDrop", "pinned_init"])
+    {
+        if actual.ident != expected {
+            return false;
+        }
+    }
+    true
 }