xref: /linux/rust/pin-init/internal/src/pinned_drop.rs (revision 50426bde1577d17e61362bd199d487dbeb159110)
1 // SPDX-License-Identifier: Apache-2.0 OR MIT
2 
3 use proc_macro2::{TokenStream, TokenTree};
4 use quote::quote;
5 
6 pub(crate) fn pinned_drop(_args: TokenStream, input: TokenStream) -> TokenStream {
7     let mut toks = input.into_iter().collect::<Vec<_>>();
8     assert!(!toks.is_empty());
9     // Ensure that we have an `impl` item.
10     assert!(matches!(&toks[0], TokenTree::Ident(i) if i == "impl"));
11     // Ensure that we are implementing `PinnedDrop`.
12     let mut nesting: usize = 0;
13     let mut pinned_drop_idx = None;
14     for (i, tt) in toks.iter().enumerate() {
15         match tt {
16             TokenTree::Punct(p) if p.as_char() == '<' => {
17                 nesting += 1;
18             }
19             TokenTree::Punct(p) if p.as_char() == '>' => {
20                 nesting = nesting.checked_sub(1).unwrap();
21                 continue;
22             }
23             _ => {}
24         }
25         if i >= 1 && nesting == 0 {
26             // Found the end of the generics, this should be `PinnedDrop`.
27             assert!(
28                 matches!(tt, TokenTree::Ident(i) if i == "PinnedDrop"),
29                 "expected 'PinnedDrop', found: '{tt:?}'"
30             );
31             pinned_drop_idx = Some(i);
32             break;
33         }
34     }
35     let idx = pinned_drop_idx
36         .unwrap_or_else(|| panic!("Expected an `impl` block implementing `PinnedDrop`."));
37     // Fully qualify the `PinnedDrop`, as to avoid any tampering.
38     toks.splice(idx..idx, quote!(::pin_init::));
39     // Take the `{}` body and call the declarative macro.
40     if let Some(TokenTree::Group(last)) = toks.pop() {
41         let last = last.stream();
42         quote!(::pin_init::__pinned_drop! {
43             @impl_sig(#(#toks)*),
44             @impl_body(#last),
45         })
46     } else {
47         TokenStream::from_iter(toks)
48     }
49 }
50