mod args;
use quote::{quote, ToTokens};
use syn::{
parse::{Parse, Parser},
parse_quote,
};
#[proc_macro]
pub fn init_depth_var(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
let output = if input.is_empty() {
quote! {
::std::thread_local! {
static DEPTH: ::std::cell::Cell<usize> = ::std::cell::Cell::new(0);
}
}
} else {
let input2 = proc_macro2::TokenStream::from(input);
syn::Error::new_spanned(input2, "`init_depth_var` takes no arguments").to_compile_error()
};
output.into()
}
#[proc_macro_attribute]
pub fn trace(
args: proc_macro::TokenStream,
input: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
let raw_args = syn::parse_macro_input!(args as syn::AttributeArgs);
let args = match args::Args::from_raw_args(raw_args) {
Ok(args) => args,
Err(errors) => {
return errors
.iter()
.map(syn::Error::to_compile_error)
.collect::<proc_macro2::TokenStream>()
.into()
}
};
let output = if let Ok(item) = syn::Item::parse.parse(input.clone()) {
expand_item(&args, item)
} else if let Ok(impl_item) = syn::ImplItem::parse.parse(input.clone()) {
expand_impl_item(&args, impl_item)
} else {
let input2 = proc_macro2::TokenStream::from(input);
syn::Error::new_spanned(input2, "expected one of: `fn`, `impl`, `mod`").to_compile_error()
};
output.into()
}
#[derive(Clone, Copy)]
enum AttrApplied {
Directly,
Indirectly,
}
fn expand_item(args: &args::Args, mut item: syn::Item) -> proc_macro2::TokenStream {
transform_item(args, AttrApplied::Directly, &mut item);
match item {
syn::Item::Fn(_) | syn::Item::Mod(_) | syn::Item::Impl(_) => item.into_token_stream(),
_ => syn::Error::new_spanned(item, "#[trace] is not supported for this item")
.to_compile_error(),
}
}
fn expand_impl_item(args: &args::Args, mut impl_item: syn::ImplItem) -> proc_macro2::TokenStream {
transform_impl_item(args, AttrApplied::Directly, &mut impl_item);
match impl_item {
syn::ImplItem::Method(_) => impl_item.into_token_stream(),
_ => syn::Error::new_spanned(impl_item, "#[trace] is not supported for this impl item")
.to_compile_error(),
}
}
fn transform_item(args: &args::Args, attr_applied: AttrApplied, item: &mut syn::Item) {
match *item {
syn::Item::Fn(ref mut item_fn) => transform_fn(args, attr_applied, item_fn),
syn::Item::Mod(ref mut item_mod) => transform_mod(args, attr_applied, item_mod),
syn::Item::Impl(ref mut item_impl) => transform_impl(args, attr_applied, item_impl),
_ => (),
}
}
fn transform_fn(args: &args::Args, attr_applied: AttrApplied, item_fn: &mut syn::ItemFn) {
item_fn.block = Box::new(construct_traced_block(
&args,
attr_applied,
&item_fn.sig,
&item_fn.block,
));
}
fn transform_mod(args: &args::Args, attr_applied: AttrApplied, item_mod: &mut syn::ItemMod) {
assert!(
(item_mod.content.is_some() && item_mod.semi.is_none())
|| (item_mod.content.is_none() && item_mod.semi.is_some())
);
if item_mod.semi.is_some() {
unimplemented!();
}
if let Some((_, items)) = item_mod.content.as_mut() {
items.iter_mut().for_each(|item| {
if let AttrApplied::Directly = attr_applied {
match *item {
syn::Item::Fn(syn::ItemFn {
sig: syn::Signature { ref ident, .. },
..
})
| syn::Item::Mod(syn::ItemMod { ref ident, .. }) => match args.filter {
args::Filter::Enable(ref idents) if !idents.contains(ident) => {
return;
}
args::Filter::Disable(ref idents) if idents.contains(ident) => {
return;
}
_ => (),
},
_ => (),
}
}
transform_item(args, AttrApplied::Indirectly, item);
});
items.insert(
0,
parse_quote! {
::std::thread_local! {
static DEPTH: ::std::cell::Cell<usize> = ::std::cell::Cell::new(0);
}
},
);
}
}
fn transform_impl(args: &args::Args, attr_applied: AttrApplied, item_impl: &mut syn::ItemImpl) {
item_impl.items.iter_mut().for_each(|impl_item| {
if let syn::ImplItem::Method(ref mut impl_item_method) = *impl_item {
if let AttrApplied::Directly = attr_applied {
let ident = &impl_item_method.sig.ident;
match args.filter {
args::Filter::Enable(ref idents) if !idents.contains(ident) => {
return;
}
args::Filter::Disable(ref idents) if idents.contains(ident) => {
return;
}
_ => (),
}
}
impl_item_method.block = construct_traced_block(
&args,
AttrApplied::Indirectly,
&impl_item_method.sig,
&impl_item_method.block,
);
}
});
}
fn transform_impl_item(
args: &args::Args,
attr_applied: AttrApplied,
impl_item: &mut syn::ImplItem,
) {
#[allow(clippy::single_match)]
match *impl_item {
syn::ImplItem::Method(ref mut impl_item_method) => {
transform_method(args, attr_applied, impl_item_method)
}
_ => (),
}
}
fn transform_method(
args: &args::Args,
attr_applied: AttrApplied,
impl_item_method: &mut syn::ImplItemMethod,
) {
impl_item_method.block = construct_traced_block(
&args,
attr_applied,
&impl_item_method.sig,
&impl_item_method.block,
);
}
fn construct_traced_block(
args: &args::Args,
attr_applied: AttrApplied,
sig: &syn::Signature,
original_block: &syn::Block,
) -> syn::Block {
let arg_idents = extract_arg_idents(args, attr_applied, &sig);
let arg_idents_format = arg_idents
.iter()
.map(|arg_ident| format!("{} = {{:?}}", arg_ident))
.collect::<Vec<_>>()
.join(", ");
let pretty = if args.pretty { "#" } else { "" };
let entering_format = format!(
"{{:depth$}}{} Entering {}({})",
args.prefix_enter, sig.ident, arg_idents_format
);
let exiting_format = format!(
"{{:depth$}}{} Exiting {} = {{:{}?}}",
args.prefix_exit, sig.ident, pretty
);
let pause_stmt = if args.pause {
quote! {{
use std::io::{self, BufRead};
let stdin = io::stdin();
stdin.lock().lines().next();
}}
} else {
quote!()
};
let printer = if args.logging {
quote! { log::trace! }
} else {
quote! { println! }
};
parse_quote! {{
#printer(#entering_format, "", #(#arg_idents,)* depth = DEPTH.with(|d| d.get()));
#pause_stmt
DEPTH.with(|d| d.set(d.get() + 1));
let mut fn_closure = move || #original_block;
let fn_return_value = fn_closure();
DEPTH.with(|d| d.set(d.get() - 1));
#printer(#exiting_format, "", fn_return_value, depth = DEPTH.with(|d| d.get()));
#pause_stmt
fn_return_value
}}
}
fn extract_arg_idents(
args: &args::Args,
attr_applied: AttrApplied,
sig: &syn::Signature,
) -> Vec<proc_macro2::Ident> {
fn process_pat(
args: &args::Args,
attr_applied: AttrApplied,
pat: &syn::Pat,
arg_idents: &mut Vec<proc_macro2::Ident>,
) {
match *pat {
syn::Pat::Ident(ref pat_ident) => {
let ident = &pat_ident.ident;
if let AttrApplied::Directly = attr_applied {
match args.filter {
args::Filter::Enable(ref idents) if !idents.contains(ident) => {
return;
}
args::Filter::Disable(ref idents) if idents.contains(ident) => {
return;
}
_ => (),
}
}
arg_idents.push(ident.clone());
}
syn::Pat::Tuple(ref pat_tuple) => {
pat_tuple.elems.iter().for_each(|pat| {
process_pat(args, attr_applied, pat, arg_idents);
});
}
_ => unimplemented!(),
}
}
let mut arg_idents = vec![];
for input in &sig.inputs {
match input {
syn::FnArg::Receiver(_) => (),
syn::FnArg::Typed(arg_typed) => {
process_pat(args, attr_applied, &arg_typed.pat, &mut arg_idents);
}
}
}
arg_idents
}