50 lines · rust
1// SPDX-License-Identifier: Apache-2.0 OR MIT2 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 tt30 );31 pinned_drop_idx = Some(i);32 break;33 }34 }35 let idx = pinned_drop_idx36 .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}50