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