metrics_fn_codegen/
lib.rs1use proc_macro::TokenStream;
12use proc_macro2::{Span, TokenStream as TokenStream2};
13use quote::{ToTokens, quote_spanned};
14use syn::parse_macro_input;
15
16mod call_type;
17mod function;
18mod return_type_classification;
19
20use function::*;
21
22use crate::return_type_classification::ReturnTypeClassification;
23
24#[proc_macro_attribute]
25pub fn dummy(_attr: TokenStream, item: TokenStream) -> TokenStream {
26 item
27}
28
29#[proc_macro_attribute]
30pub fn measure(attrs: TokenStream, item: TokenStream) -> TokenStream {
31 let span = proc_macro2::Span::call_site();
32
33 if !attrs.is_empty() {
34 return syn::Error::new(span, "#[measure] does not take arguments.")
35 .to_compile_error()
36 .into();
37 }
38
39 let original_fn = parse_macro_input!(item as Function);
40 let wrapped_fn =
41 original_fn.rename(format!("{}__{}", original_fn.function.sig.ident.clone().to_string(), "wrapped").as_str());
42
43 let wrapped_attrs_tokens = original_fn.attributes_tokens();
44 let wrapped_call_tokens = wrapped_fn.call(span);
45 let wrapped_call_fn_name = original_fn.function.sig.ident.clone().to_string();
46 let wrapped_sig_tokens = wrapped_fn.function.sig.into_token_stream();
47 let wrapped_body_tokens = original_fn.function.block.clone().into_token_stream();
48 let wrapper_sig_tokens = original_fn.signature_full();
49 let record_call_tokens = build_record_call(span, original_fn, wrapped_call_fn_name);
50
51 let output = quote_spanned! { span =>
52 #wrapped_attrs_tokens
53 #wrapper_sig_tokens {
54
55 let start__ = std::time::Instant::now();
56 let output__ = #wrapped_call_tokens;
57 let end__ = std::time::Instant::now();
58
59 let module_name = module_path!();
60 #record_call_tokens
61
62 return output__;
63 }
64
65 #[allow(non_snake_case)]
66 #wrapped_sig_tokens
67 #wrapped_body_tokens
68 };
69
70 output.into()
71}
72
73fn build_record_call(span: Span, original_fn: Function, wrapped_call_fn_name: String) -> TokenStream2 {
74 match original_fn.return_type() {
75 ReturnTypeClassification::Result => {
76 quote_spanned! { span =>
77 let result__: core::result::Result<(), ()> = if output__.is_ok() {
78 Ok(())
79 } else {
80 Err(())
81 };
82 metrics_fn::record(module_name, #wrapped_call_fn_name, result__, end__.duration_since(start__).as_secs_f64());
83 }
84 },
85 _ => {
86 quote_spanned! { span =>
87 metrics_fn::record(module_name, #wrapped_call_fn_name, Ok(()), end__.duration_since(start__).as_secs_f64());
88 }
89 },
90 }
91}