1// SPDX-License-Identifier: Apache-2.0 OR MIT
 2
 3use proc_macro::{TokenStream, TokenTree};
 4
 5pub(crate) fn pinned_drop(_args: TokenStream, input: TokenStream) -> TokenStream {
 6    let mut toks = input.into_iter().collect::<Vec<_>>();
 7    assert!(!toks.is_empty());
 8    // Ensure that we have an `impl` item.
 9    assert!(matches!(&toks[0], TokenTree::Ident(i) if i.to_string() == "impl"));
10    // Ensure that we are implementing `PinnedDrop`.
11    let mut nesting: usize = 0;
12    let mut pinned_drop_idx = None;
13    for (i, tt) in toks.iter().enumerate() {
14        match tt {
15            TokenTree::Punct(p) if p.as_char() == '<' => {
16                nesting += 1;
17            }
18            TokenTree::Punct(p) if p.as_char() == '>' => {
19                nesting = nesting.checked_sub(1).unwrap();
20                continue;
21            }
22            _ => {}
23        }
24        if i >= 1 && nesting == 0 {
25            // Found the end of the generics, this should be `PinnedDrop`.
26            assert!(
27                matches!(tt, TokenTree::Ident(i) if i.to_string() == "PinnedDrop"),
28                "expected 'PinnedDrop', found: '{:?}'",
29                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!(::kernel::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!(::kernel::__pinned_drop! {
43            @impl_sig(#(#toks)*),
44            @impl_body(#last),
45        })
46    } else {
47        TokenStream::from_iter(toks)
48    }
49}