rust/macros/pinned_drop.rs

   1 // SPDX-License-Identifier: Apache-2.0 OR MIT
   2
   3 use proc_macro::{TokenStream, TokenTree};
   4
   5 pub(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 }