diff --git a/examples/example_custom_format.rs b/examples/example_custom_format.rs new file mode 100644 index 0000000..31705a5 --- /dev/null +++ b/examples/example_custom_format.rs @@ -0,0 +1,19 @@ +use trace::trace; + +trace::init_depth_var!(); + +fn main() { + foo(5, 4); +} + +#[trace(format_enter = "{y} and {z} {{7}}", format_exit = "{r} * {r}")] +fn foo(z: u32, y: u32) -> u32 { + z +} + +#[cfg(test)] +#[macro_use] +mod trace_test; + +#[cfg(test)] +trace_test!(test_custom_format, main()); diff --git a/examples/expected_test_outputs/test_custom_format.expected b/examples/expected_test_outputs/test_custom_format.expected new file mode 100644 index 0000000..1a2a777 --- /dev/null +++ b/examples/expected_test_outputs/test_custom_format.expected @@ -0,0 +1,2 @@ +[+] Entering foo(4 and 5 {7}) +[-] Exiting foo = 5 * 5 diff --git a/src/args.rs b/src/args.rs index 480472c..947992d 100644 --- a/src/args.rs +++ b/src/args.rs @@ -5,6 +5,8 @@ use syn::{self, spanned::Spanned}; pub(crate) struct Args { pub(crate) prefix_enter: String, pub(crate) prefix_exit: String, + pub(crate) format_enter: Option, + pub(crate) format_exit: Option, pub(crate) filter: Filter, pub(crate) pause: bool, pub(crate) pretty: bool, @@ -23,6 +25,17 @@ const DEFAULT_PAUSE: bool = false; const DEFAULT_PRETTY: bool = false; const DEFAULT_LOGGING: bool = false; +macro_rules! try_extract_str { + ($lit:expr, $meta:expr, $arg_ty:ident) => {{ + match *$lit { + syn::Lit::Str(ref lit_str) => Ok(Arg::$arg_ty($meta.span(), lit_str.value())), + _ => Err(vec![syn::Error::new_spanned( + $lit, + format!("`{}` must have a string value", stringify!($arg_ty)), + )]), + } + }}; +} impl Args { pub(crate) fn from_raw_args(raw_args: syn::AttributeArgs) -> Result> { // Different types of arguments accepted by `#[trace]`; @@ -35,6 +48,8 @@ impl Args { Pause(proc_macro2::Span, bool), Pretty(proc_macro2::Span, bool), Logging(proc_macro2::Span, bool), + FormatEnter(proc_macro2::Span, String), + FormatExit(proc_macro2::Span, String), } // Parse arguments @@ -43,6 +58,8 @@ impl Args { enum ArgName { PrefixEnter, PrefixExit, + FormatEnter, + FormatExit, Enable, Disable, Pause, @@ -54,6 +71,8 @@ impl Args { let arg_name = match ident.to_string().as_str() { "prefix_enter" => ArgName::PrefixEnter, "prefix_exit" => ArgName::PrefixExit, + "format_enter" => ArgName::FormatEnter, + "format_exit" => ArgName::FormatExit, "enable" => ArgName::Enable, "disable" => ArgName::Disable, "pause" => ArgName::Pause, @@ -79,6 +98,18 @@ impl Args { "`prefix_exit` requires a string value", )] }; + let format_enter_type_error = || { + vec![syn::Error::new_spanned( + ident.clone(), + "`format_enter` requires a string value", + )] + }; + let format_exit_type_error = || { + vec![syn::Error::new_spanned( + ident.clone(), + "`format_exit` requires a string value", + )] + }; let enable_type_error = || { vec![syn::Error::new_spanned( ident.clone(), @@ -115,11 +146,12 @@ impl Args { ArgName::Pause => Ok(Arg::Pause(meta.span(), true)), ArgName::Pretty => Ok(Arg::Pretty(meta.span(), true)), ArgName::Logging => Ok(Arg::Logging(meta.span(), true)), - ArgName::PrefixEnter => Err(prefix_enter_type_error()), ArgName::PrefixExit => Err(prefix_exit_type_error()), ArgName::Enable => Err(enable_type_error()), ArgName::Disable => Err(disable_type_error()), + ArgName::FormatEnter => Err(format_enter_type_error()), + ArgName::FormatExit => Err(format_exit_type_error()), }, syn::Meta::List(syn::MetaList { ref nested, .. }) => match arg_name { ArgName::Enable => { @@ -172,27 +204,14 @@ impl Args { ArgName::Pause => Err(pause_type_error()), ArgName::Pretty => Err(pretty_type_error()), ArgName::Logging => Err(logging_type_error()), + ArgName::FormatEnter => Err(format_enter_type_error()), + ArgName::FormatExit => Err(format_exit_type_error()), }, syn::Meta::NameValue(syn::MetaNameValue { ref lit, .. }) => match arg_name { - ArgName::PrefixEnter => match *lit { - syn::Lit::Str(ref lit_str) => { - Ok(Arg::PrefixEnter(meta.span(), lit_str.value())) - } - _ => Err(vec![syn::Error::new_spanned( - lit, - "`prefix_enter` must have a string value", - )]), - }, - ArgName::PrefixExit => match *lit { - syn::Lit::Str(ref lit_str) => { - Ok(Arg::PrefixExit(meta.span(), lit_str.value())) - } - _ => Err(vec![syn::Error::new_spanned( - lit, - "`prefix_exit` must have a string value", - )]), - }, - + ArgName::PrefixEnter => try_extract_str!(lit, meta, PrefixEnter), + ArgName::PrefixExit => try_extract_str!(lit, meta, PrefixExit), + ArgName::FormatEnter => try_extract_str!(lit, meta, FormatEnter), + ArgName::FormatExit => try_extract_str!(lit, meta, FormatExit), ArgName::Enable => Err(enable_type_error()), ArgName::Disable => Err(disable_type_error()), ArgName::Pause => Err(pause_type_error()), @@ -209,6 +228,8 @@ impl Args { let mut prefix_enter_args = vec![]; let mut prefix_exit_args = vec![]; + let mut format_enter_args = vec![]; + let mut format_exit_args = vec![]; let mut enable_args = vec![]; let mut disable_args = vec![]; let mut pause_args = vec![]; @@ -227,6 +248,8 @@ impl Args { Arg::Pause(span, b) => pause_args.push((span, b)), Arg::Pretty(span, b) => pretty_args.push((span, b)), Arg::Logging(span, b) => logging_args.push((span, b)), + Arg::FormatEnter(span, s) => format_enter_args.push((span, s)), + Arg::FormatExit(span, s) => format_exit_args.push((span, s)), }, Err(es) => errors.extend(es), } @@ -247,6 +270,20 @@ impl Args { .map(|(span, _)| syn::Error::new(*span, "duplicate `prefix_exit`")), ); } + if format_enter_args.len() >= 2 { + errors.extend( + format_enter_args + .iter() + .map(|(span, _)| syn::Error::new(*span, "duplicate `format_enter`")), + ); + } + if format_exit_args.len() >= 2 { + errors.extend( + format_exit_args + .iter() + .map(|(span, _)| syn::Error::new(*span, "duplicate `format_exit`")), + ); + } if enable_args.len() >= 2 { errors.extend( enable_args @@ -294,6 +331,26 @@ impl Args { "cannot have both `enable` and `disable`", )); } + if pretty_args.len() == 1 && format_enter_args.len() == 1 { + errors.push(syn::Error::new( + pretty_args[0].0, + "cannot have both `pretty` and `format_enter`", + )); + errors.push(syn::Error::new( + format_enter_args[0].0, + "cannot have both `pretty` and `format_enter`", + )); + } + if pretty_args.len() == 1 && format_exit_args.len() == 1 { + errors.push(syn::Error::new( + pretty_args[0].0, + "cannot have both `pretty` and `format_exit`", + )); + errors.push(syn::Error::new( + format_exit_args[0].0, + "cannot have both `pretty` and `format_exit`", + )); + } if errors.is_empty() { macro_rules! first_no_span { @@ -306,6 +363,8 @@ impl Args { .unwrap_or_else(|| DEFAULT_PREFIX_ENTER.to_owned()); let prefix_exit = first_no_span!(prefix_exit_args).unwrap_or_else(|| DEFAULT_PREFIX_EXIT.to_owned()); + let format_enter = first_no_span!(format_enter_args); + let format_exit = first_no_span!(format_exit_args); let filter = match (first_no_span!(enable_args), first_no_span!(disable_args)) { (None, None) => Filter::None, (Some(idents), None) => Filter::Enable(idents), @@ -323,6 +382,8 @@ impl Args { pause, pretty, logging, + format_enter, + format_exit, }) } else { Err(errors) diff --git a/src/lib.rs b/src/lib.rs index 2ff6076..3716e8c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -74,6 +74,9 @@ mod args; +use std::{iter::Peekable, str::Chars}; + +use proc_macro2::{Span, TokenStream}; use quote::{quote, ToTokens}; use syn::{ parse::{Parse, Parser}, @@ -139,7 +142,30 @@ pub fn init_depth_var(input: proc_macro::TokenStream) -> proc_macro::TokenStream /// /// - `logging` - Use `log::trace!` from the `log` crate instead of `println`. Disabled by default. /// -/// Note that `enable` and `disable` can not be used together, and doing so will result in an error. +/// - `format_enter` - The format (anything after the prefix) of `println!` statements when a function +/// is entered. Allows parameter interpolation like: +/// ```rust +/// #[trace(format_enter = "i is {i}")] +/// fn foo(i: i32) { +/// println!("foo") +/// } +/// ``` +/// Interpolation follows the same rules as `format!()` besides for the fact that there is no pretty printing, +/// that is anything interpolated will be debug formatted. Disabled by default. + +/// - `format_exit` - The format (anything after the prefix) of `println!` statements when a function +/// is exited. To interpolate the return value use `{r}`: +/// ```rust +/// #[trace(format_exit = "returning {r}")] +/// fn foo() -> i32 { +/// 1 +/// } +/// ``` +/// Otherwise formatting follows the same rules as `format_enter`. Disabled by default. +/// +/// Note that `enable` and `disable` cannot be used together, and doing so will result in an error. +/// +/// Further note that `format_enter` or `format_exit` cannot be used together with with `pretty`, and doing so will result in an error. #[proc_macro_attribute] pub fn trace( args: proc_macro::TokenStream, @@ -319,21 +345,58 @@ fn construct_traced_block( 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 + let arg_idents = extract_arg_idents(args, attr_applied, sig) .iter() - .map(|arg_ident| format!("{} = {{:?}}", arg_ident)) - .collect::>() - .join(", "); - - let pretty = if args.pretty { "#" } else { "" }; + .map(|ident| ident.to_token_stream()) + .collect(); + let (enter_format, arg_idents) = if let Some(fmt_str) = &args.format_enter { + parse_fmt_str(fmt_str, arg_idents) + } else { + ( + Ok(arg_idents + .iter() + .map(|arg_ident| format!("{} = {{:?}}", arg_ident)) + .collect::>() + .join(", ")), + arg_idents, + ) + }; + // we set set exit val to be a vector with one element which is Ident called r + // this means that the format parser can indentify when then return value should be interprolated + // so if we want to use a different symbol to denote return value interpolation we just need to change the symbol in the following quote + // ie: `let exit_val = vec![quote!(return_value)];` if we wanted to use return_value to denote return value interpolation + let exit_val = vec![quote!(r)]; + let (exit_format, exit_val) = if let Some(fmt_str) = &args.format_exit { + parse_fmt_str(fmt_str, exit_val) + } else if args.pretty { + (Ok("{:#?}".to_string()), exit_val) + } else { + (Ok("{:?}".to_string()), exit_val) + }; + let should_interpolate = !exit_val.is_empty(); let entering_format = format!( "{{:depth$}}{} Entering {}({})", - args.prefix_enter, sig.ident, arg_idents_format + args.prefix_enter, + sig.ident, + match enter_format { + Ok(ok) => ok, + Err(e) => { + let error = e.into_compile_error(); + return parse_quote! {{#error}}; + } + } ); let exiting_format = format!( - "{{:depth$}}{} Exiting {} = {{:{}?}}", - args.prefix_exit, sig.ident, pretty + "{{:depth$}}{} Exiting {} = {}", + args.prefix_exit, + sig.ident, + match exit_format { + Ok(ok) => ok, + Err(e) => { + let error = e.into_compile_error(); + return parse_quote! {{#error}}; + } + } ); let pause_stmt = if args.pause { @@ -351,19 +414,137 @@ fn construct_traced_block( } else { quote! { println! } }; - + let print_exit = if should_interpolate { + quote! {{#printer(#exiting_format, "",fn_return_value, depth = DEPTH.with(|d| d.get()));}} + } else { + quote!(#printer(#exiting_format, "", depth = DEPTH.with(|d| d.get()));) + }; parse_quote! {{ #printer(#entering_format, "", #(#arg_idents,)* depth = DEPTH.with(|d| d.get())); #pause_stmt DEPTH.with(|d| d.set(d.get() + 1)); let fn_return_value = #original_block; DEPTH.with(|d| d.set(d.get() - 1)); - #printer(#exiting_format, "", fn_return_value, depth = DEPTH.with(|d| d.get())); + #print_exit #pause_stmt fn_return_value }} } +fn parse_fmt_str( + fmt_str: &str, + mut arg_idents: Vec, +) -> (Result, Vec) { + let mut fixed_format_str = String::new(); + let mut kept_arg_idents = Vec::new(); + let mut fmt_iter = fmt_str.chars().peekable(); + while let Some(fmt_char) = fmt_iter.next() { + match fmt_char { + '{' => { + if let Some('{') = fmt_iter.peek() { + fixed_format_str.push_str("{{"); + fmt_iter.next(); + } else { + match parse_interpolated(&mut fmt_iter, &mut arg_idents, &mut kept_arg_idents) { + Ok(interpolated) => fixed_format_str.push_str(&interpolated), + Err(e) => return (Err(e), kept_arg_idents), + } + } + } + '}' => { + if fmt_iter.next() != Some('}') { + return (Err(syn::Error::new( + Span::call_site(), + "invalid format string: unmatched `}` found\nif you intended to print `}`, you can escape it using `}}`" + )), kept_arg_idents); + } + + fixed_format_str.push_str("}}") + } + _ => fixed_format_str.push(fmt_char), + } + } + (Ok(fixed_format_str), kept_arg_idents) +} + +fn fix_interpolated( + last_char: char, + ident: String, + arg_idents: &mut Vec, + kept_arg_idents: &mut Vec, +) -> Result { + if last_char != '}' { + return Err(syn::Error::new( + Span::call_site(), + "invalid format string: expected `'}}'` but string was terminated\nif you intended to print `{{`, you can escape it using `{{`.", + )); + } + let predicate = |arg_ident: &TokenStream| arg_ident.to_string() == ident; + if let Some(index) = kept_arg_idents.iter().position(predicate) { + Ok(format!("{{{}}}", index + 1)) + } else if let Some(index) = arg_idents.iter().position(predicate) { + kept_arg_idents.push(arg_idents.remove(index)); + Ok(format!("{{{}}}", kept_arg_idents.len())) + } else { + Err(syn::Error::new( + Span::call_site(), + // TODO: better error message + format!("cannot find `{ident}` in this scope."), + )) + } +} + +fn parse_interpolated( + fmt_iter: &mut Peekable, + arg_idents: &mut Vec, + kept_arg_idents: &mut Vec, +) -> Result { + let mut last_char = ' '; + let mut ident = String::new(); + while let Some(ident_char) = fmt_iter.next() { + match ident_char { + '}' => { + last_char = '}'; + break; + } + _ => { + last_char = ident_char; + if !ident_char.is_whitespace() { + ident.push(ident_char); + } else { + skip_whitespace_and_check(fmt_iter, &mut last_char, ident_char)?; + } + } + } + } + fix_interpolated(last_char, ident, arg_idents, kept_arg_idents) +} + +fn skip_whitespace_and_check( + fmt_iter: &mut Peekable, + last_char: &mut char, + ident_char: char, +) -> Result<(), syn::Error> { + for blank_char in fmt_iter.by_ref() { + match blank_char { + '}' => { + *last_char = '}'; + break; + } + c if c.is_whitespace() => { + *last_char = ident_char; + } + _ => { + return Err(syn::Error::new( + Span::call_site(), + format!("invalid format string: expected `'}}'`, found `'{blank_char}'`\nif you intended to print `{{`, you can escape it using `{{`."), + )) + } + } + } + Ok(()) +} + fn extract_arg_idents( args: &args::Args, attr_applied: AttrApplied,