extern crate proc_macro;
extern crate quote;
extern crate syn;
use self::proc_macro::TokenStream;
use quote::quote;
use syn::fold::{self, Fold};
use syn::parse::{Parse, ParseStream, Result};
use syn::punctuated::Punctuated;
use syn::{
parse_macro_input, parse_quote, Attribute, ImplItemFn, Item, ItemFn, ItemImpl, ItemMod,
ItemTrait, Token, TraitItemFn,
};
#[derive(Default)]
struct Flamer {
id_stack: Vec<String>,
}
impl Flamer {
fn push(&mut self, ident: String) {
self.id_stack.push(ident);
}
fn pop(&mut self) {
let _ = self.id_stack.pop();
}
fn name(&self) -> String {
self.id_stack.join("::")
}
fn is_noflame(&self, i: &[Attribute]) -> bool {
i.iter().any(|ref a| {
if a.path().segments.len() == 1 {
let ident = &a.path().segments.iter().next().unwrap().ident;
ident == "flame" || ident == "noflame"
} else {
false
}
})
}
}
impl Parse for Flamer {
fn parse(input: ParseStream) -> Result<Self> {
let vars = Punctuated::<syn::LitStr, Token![,]>::parse_terminated(input)?;
Ok(Flamer {
id_stack: vars.into_iter().map(|s| s.value()).collect::<Vec<_>>(),
})
}
}
impl Fold for Flamer {
fn fold_item_mod(&mut self, i: ItemMod) -> ItemMod {
if self.is_noflame(&i.attrs) {
return i;
}
self.push(i.ident.to_string());
let item_mod = fold::fold_item_mod(self, i);
self.pop();
item_mod
}
fn fold_item_trait(&mut self, i: ItemTrait) -> ItemTrait {
if self.is_noflame(&i.attrs) {
return i;
}
self.push(i.ident.to_string());
let t = fold::fold_item_trait(self, i);
self.pop();
t
}
fn fold_trait_item_fn(&mut self, i: TraitItemFn) -> TraitItemFn {
if i.sig.constness.is_some() || self.is_noflame(&i.attrs) {
return i;
}
self.push(i.sig.ident.to_string());
let m = fold::fold_trait_item_fn(self, i);
self.pop();
m
}
fn fold_item_impl(&mut self, i: ItemImpl) -> ItemImpl {
if self.is_noflame(&i.attrs) {
return i;
}
if let Some((_, ref path, _)) = i.trait_ {
self.push(format!("{:?} as {:?}", &i.self_ty, path));
} else {
self.push(format!("{:?}", &i.self_ty));
}
let ii = fold::fold_item_impl(self, i);
self.pop();
ii
}
fn fold_impl_item_fn(&mut self, i: ImplItemFn) -> ImplItemFn {
if i.sig.constness.is_some() || self.is_noflame(&i.attrs) {
return i;
}
self.push(i.sig.ident.to_string());
let method = fold::fold_impl_item_fn(self, i);
self.pop();
method
}
fn fold_item_fn(&mut self, i: ItemFn) -> ItemFn {
if self.is_noflame(&i.attrs) {
return i;
}
if i.sig.constness.is_some() {
return fold::fold_item_fn(self, i);
}
let mut i = i;
self.push(i.sig.ident.to_string());
let name = self.name();
let stmts = ::std::mem::replace(
&mut i.block.stmts,
vec![parse_quote! {
let _flame_guard = ::flame::start_guard(#name);
}],
);
for stmt in stmts {
i.block.stmts.push(fold::fold_stmt(self, stmt));
}
self.pop();
i
}
}
#[proc_macro_attribute]
pub fn flame(attrs: TokenStream, code: TokenStream) -> TokenStream {
let input = parse_macro_input!(code as Item);
let mut flamer = parse_macro_input!(attrs as Flamer);
let item = fold::fold_item(&mut flamer, input);
TokenStream::from(quote!(#item))
}
#[proc_macro_attribute]
pub fn noflame(_attrs: TokenStream, code: TokenStream) -> TokenStream {
code
}