profiling_procmacros/
lib.rs

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 = parse_macro_input!(item as ItemFn);
12    let instrumented_function_name = function.sig.ident.to_string();
13
14    let body = &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}
profiling_procmacros/lib.rs

profiling_procmacros/
lib.rs
