1extern crate proc_macro;
2use proc_macro::TokenStream;
3use quote::{quote, ToTokens};
4use syn::{parse_macro_input, parse_quote, ImplItem, ItemFn, ItemImpl};
5
6#[proc_macro_attribute]
7pub fn function(
8 _attr: TokenStream,
9 item: TokenStream,
10) -> TokenStream {
11 let mut function: ItemFn = parse_macro_input!(item as ItemFn);
12 let instrumented_function_name: String = function.sig.ident.to_string();
13
14 let body: &Box = &function.block;
15 let new_body: syn::Block = impl_block(body, &instrumented_function_name);
16
17 function.block = Box::new(new_body);
18
19 (quote! {
20 #function
21 })
22 .into()
23}
24
25#[proc_macro_attribute]
26pub fn skip(
27 _attr: TokenStream,
28 item: TokenStream,
29) -> TokenStream {
30 item
31}
32
33#[proc_macro_attribute]
34pub fn all_functions(
35 _attr: TokenStream,
36 item: TokenStream,
37) -> TokenStream {
38 let mut content = parse_macro_input!(item as ItemImpl);
39 let struct_name = content.self_ty.to_token_stream().to_string();
40
41 'func_loop: for block in &mut content.items {
42 // Currently, we only care about the function impl part.
43 // In the future, expand the code to following if we are interested in other parts
44 //
45 // match block {
46 // ImplItem::Fn(ref mut func) => {
47 // for func_attr in &func.attrs {
48 // if let syn::Meta::Path(ref func_attr_info) = func_attr.meta {
49 // let attr_seg = func_attr_info.segments.last().unwrap();
50 // if attr_seg.ident.to_string() == "skip".to_string() {
51 // continue 'func_loop;
52 // }
53 // }
54 // }
55 // let prev_block = &func.block;
56 // let func_name = func.sig.ident.to_string();
57 // func.block = impl_block(prev_block, &func_name);
58 // }
59 // ImplItem::Macro(_) => { // some code... },
60 // ImplItem::Type(_) => { // some code... },
61 // _ => {}
62 // }
63 let ImplItem::Fn(ref mut func) = block else {
64 continue;
65 };
66
67 for func_attr in &func.attrs {
68 let func_attr_info = func_attr.path();
69 if func_attr_info.segments.is_empty() {
70 continue;
71 }
72 if func_attr_info.segments.first().unwrap().ident != "profiling" {
73 continue;
74 }
75 if func_attr_info.segments.last().unwrap().ident == "skip" {
76 continue 'func_loop;
77 }
78 }
79 let prev_block = &func.block;
80 let calling_info = format!("{}: {}", struct_name, func.sig.ident);
81 func.block = impl_block(prev_block, &calling_info);
82 }
83
84 (quote!(
85 #content
86 ))
87 .into()
88}
89
90#[cfg(not(any(
91 feature = "profile-with-puffin",
92 feature = "profile-with-optick",
93 feature = "profile-with-superluminal",
94 feature = "profile-with-tracing",
95 feature = "profile-with-tracy"
96)))]
97fn impl_block(
98 body: &syn::Block,
99 _instrumented_function_name: &str,
100) -> syn::Block {
101 parse_quote! {
102 {
103 #body
104 }
105 }
106}
107
108#[cfg(any(
109 feature = "profile-with-puffin",
110 feature = "profile-with-optick",
111 feature = "profile-with-superluminal",
112 feature = "profile-with-tracy"
113))]
114fn impl_block(
115 body: &syn::Block,
116 _instrumented_function_name: &str,
117) -> syn::Block {
118 parse_quote! {
119 {
120 profiling::function_scope!();
121
122 #body
123 }
124 }
125}
126
127#[cfg(feature = "profile-with-tracing")]
128fn impl_block(
129 body: &syn::Block,
130 instrumented_function_name: &str,
131) -> syn::Block {
132 parse_quote! {
133 {
134 let _fn_span = profiling::tracing::span!(profiling::tracing::Level::INFO, #instrumented_function_name);
135 let _fn_span_entered = _fn_span.enter();
136
137 #body
138 }
139 }
140}
141