metrics_fn_codegen/
lib.rs

1//! # Function Metrics Codegen
2//!
3//! **ATTENTION: Don't use this directly!**
4//!
5//! This project is divided in a _build time_ module (this module), and a _runtime_ module (`metrics-fn`).
6//!
7//! The runtime module includes this automatically, so you should pretty much ignore this.
8//!
9//! See [metrics-fn](https://crates.io/crates/metrics-fn/) for more details.
10
11use 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}