extern crate proc_macro;
use proc_macro::TokenStream;
use quote::*;
use syn::{
parse_macro_input, parse_quote, ArgCaptured, ArgSelf, ArgSelfRef, AttributeArgs, FnArg, ItemFn,
NestedMeta,
};
#[proc_macro_attribute]
pub fn no_alloc(args: TokenStream, item: TokenStream) -> TokenStream {
let args = parse_macro_input!(args as AttributeArgs);
let mut item = parse_macro_input!(item as ItemFn);
let mut mode = quote!(deny_alloc);
for arg in &args {
match arg {
NestedMeta::Meta(meta) if meta.name() == "forbid" => {
mode = quote!(forbid_alloc);
}
NestedMeta::Meta(meta) => {
panic!("Invalid meta argument for #[no_alloc]. {}", quote!(#meta));
}
NestedMeta::Literal(lit) => {
panic!("Invalid literal argument for #[no_alloc]. {}", quote!(#lit));
}
}
}
let mut self_hack = None;
let force_move = item.decl.inputs.iter().filter_map(|a| match a {
FnArg::SelfRef(ArgSelfRef { .. }) => None,
FnArg::SelfValue(ArgSelf { self_token, .. }) => {
self_hack = Some(quote!(
fn _self_hack<T: Copy>(t: T) {}
_self_hack(#self_token);
));
None
}
FnArg::Captured(ArgCaptured { pat, .. }) => Some(quote!( let #pat = #pat; )),
_ => panic!("FIXME: unhandled function argument in #[no_alloc]."),
});
let block = item.block;
let output = &item.decl.output;
item.block =
if item.asyncness.is_none() {
parse_quote!({
alloc_counter::#mode(move || #output {
#( #force_move )*
#self_hack
#block
})
})
} else {
parse_quote!({
use core::{ops::{Generator, GeneratorState}, pin::Pin};
let mut gen = move || #output {
if false { yield }
#( #force_move )*
#self_hack
#block
};
loop {
match alloc_counter::#mode(|| Pin::new(&mut gen).resume()) {
GeneratorState::Yielded(y) => yield y,
GeneratorState::Complete(r) => return r,
}
}
})
};
item.into_token_stream().into()
}