2 files changed, 52 insertions, 38 deletions

diff --git a/rust/pin-init/internal/src/lib.rs b/rust/pin-init/internal/src/lib.rs
index 4cc9b7b0cda1..a75b99b58189 100644
--- a/rust/pin-init/internal/src/lib.rs
+++ b/rust/pin-init/internal/src/lib.rs
@@ -28,7 +28,9 @@ pub fn pin_data(inner: TokenStream, item: TokenStream) -> TokenStream {
 
 #[proc_macro_attribute]
 pub fn pinned_drop(args: TokenStream, input: TokenStream) -> TokenStream {
-    pinned_drop::pinned_drop(args.into(), input.into()).into()
+    let args = parse_macro_input!(args);
+    let input = parse_macro_input!(input);
+    DiagCtxt::with(|dcx| pinned_drop::pinned_drop(args, input, dcx)).into()
 }
 
 #[proc_macro_derive(Zeroable)]
diff --git a/rust/pin-init/internal/src/pinned_drop.rs b/rust/pin-init/internal/src/pinned_drop.rs
index cf8cd1c42984..a20ac314ca82 100644
--- a/rust/pin-init/internal/src/pinned_drop.rs
+++ b/rust/pin-init/internal/src/pinned_drop.rs
@@ -1,49 +1,61 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
-use proc_macro2::{TokenStream, TokenTree};
+use proc_macro2::TokenStream;
 use quote::quote;
+use syn::{parse::Nothing, parse_quote, spanned::Spanned, ImplItem, ItemImpl, Token};
 
-pub(crate) fn pinned_drop(_args: TokenStream, input: TokenStream) -> TokenStream {
-    let mut toks = input.into_iter().collect::<Vec<_>>();
-    assert!(!toks.is_empty());
-    // Ensure that we have an `impl` item.
-    assert!(matches!(&toks[0], TokenTree::Ident(i) if i == "impl"));
-    // Ensure that we are implementing `PinnedDrop`.
-    let mut nesting: usize = 0;
-    let mut pinned_drop_idx = None;
-    for (i, tt) in toks.iter().enumerate() {
-        match tt {
-            TokenTree::Punct(p) if p.as_char() == '<' => {
-                nesting += 1;
+use crate::diagnostics::{DiagCtxt, ErrorGuaranteed};
+
+pub(crate) fn pinned_drop(
+    _args: Nothing,
+    mut input: ItemImpl,
+    dcx: &mut DiagCtxt,
+) -> Result<TokenStream, ErrorGuaranteed> {
+    if let Some(unsafety) = input.unsafety {
+        dcx.error(unsafety, "implementing `PinnedDrop` is safe");
+    }
+    input.unsafety = Some(Token![unsafe](input.impl_token.span));
+    match &mut input.trait_ {
+        Some((not, path, _for)) => {
+            if let Some(not) = not {
+                dcx.error(not, "cannot implement `!PinnedDrop`");
             }
-            TokenTree::Punct(p) if p.as_char() == '>' => {
-                nesting = nesting.checked_sub(1).unwrap();
-                continue;
+            for (seg, expected) in path
+                .segments
+                .iter()
+                .rev()
+                .zip(["PinnedDrop", "pin_init", ""])
+            {
+                if expected.is_empty() || seg.ident != expected {
+                    dcx.error(seg, "bad import path for `PinnedDrop`");
+                }
+                if !seg.arguments.is_none() {
+                    dcx.error(&seg.arguments, "unexpected arguments for `PinnedDrop` path");
+                }
             }
-            _ => {}
+            *path = parse_quote!(::pin_init::PinnedDrop);
         }
-        if i >= 1 && nesting == 0 {
-            // Found the end of the generics, this should be `PinnedDrop`.
-            assert!(
-                matches!(tt, TokenTree::Ident(i) if i == "PinnedDrop"),
-                "expected 'PinnedDrop', found: '{tt:?}'"
+        None => {
+            let span = input
+                .impl_token
+                .span
+                .join(input.self_ty.span())
+                .unwrap_or(input.impl_token.span);
+            dcx.error(
+                span,
+                "expected `impl ... PinnedDrop for ...`, got inherent impl",
             );
-            pinned_drop_idx = Some(i);
-            break;
         }
     }
-    let idx = pinned_drop_idx
-        .unwrap_or_else(|| panic!("Expected an `impl` block implementing `PinnedDrop`."));
-    // Fully qualify the `PinnedDrop`, as to avoid any tampering.
-    toks.splice(idx..idx, quote!(::pin_init::));
-    // Take the `{}` body and call the declarative macro.
-    if let Some(TokenTree::Group(last)) = toks.pop() {
-        let last = last.stream();
-        quote!(::pin_init::__pinned_drop! {
-            @impl_sig(#(#toks)*),
-            @impl_body(#last),
-        })
-    } else {
-        TokenStream::from_iter(toks)
+    for item in &mut input.items {
+        if let ImplItem::Fn(fn_item) = item {
+            if fn_item.sig.ident == "drop" {
+                fn_item
+                    .sig
+                    .inputs
+                    .push(parse_quote!(_: ::pin_init::__internal::OnlyCallFromDrop));
+            }
+        }
     }
+    Ok(quote!(#input))
 }