light_profiler_macro/
lib.rs

1use proc_macro::TokenStream;
2
3#[proc_macro_attribute]
4pub fn profile(_attr: TokenStream, item: TokenStream) -> TokenStream {
5    // If profile-program feature is not enabled, return the original function completely unchanged
6    #[cfg(not(feature = "profile-program"))]
7    {
8        return item;
9    }
10
11    #[cfg(feature = "profile-program")]
12    {
13        use quote::quote;
14        use syn::{parse_macro_input, ItemFn, ReturnType};
15        let mut f = parse_macro_input!(item as ItemFn);
16        let original_body = f.block.clone();
17        let sig = f.sig.clone();
18        let ident = sig.ident.clone();
19        let fn_name_str = ident.to_string();
20        let returns_value = !matches!(sig.output, ReturnType::Default);
21
22        // Add padding before function name to align with program ID position
23        // and padding after to align "consumed" with program's "consumed"
24        let program_id_width = 43;
25        let fn_name_len = fn_name_str.len();
26        let front_padding = " ".repeat(8); // Same as "Program " length to align start position
27        let back_padding = if fn_name_len < program_id_width {
28            " ".repeat(program_id_width - fn_name_len) // Add 1 extra space to fix alignment
29        } else {
30            " ".to_string() // minimum one space
31        };
32
33        // Create profiling start and end calls with feature flag and compile-time caller info
34        let profile_start = quote! {
35            #[cfg(all(target_os = "solana", feature = "profile-program"))]
36            {
37                extern "C" {
38                    #[inline(always)]
39                    fn sol_log_compute_units_start(id_addr: u64, id_len: u64, heap_value: u64, with_heap: u64, _arg5: u64);
40                }
41                // Dynamic padding calculated at compile time
42                const PROFILE_ID: &str = concat!(#fn_name_str, "\n", #front_padding, file!(), ":", line!(), #back_padding);
43
44                #[cfg(feature = "profile-heap")]
45                unsafe {
46                    sol_log_compute_units_start(
47                        PROFILE_ID.as_ptr() as u64,
48                        PROFILE_ID.len() as u64,
49                     ::light_heap::GLOBAL_ALLOCATOR.get_used_heap(),
50                        1u64,
51                        0
52                    );
53                }
54
55                #[cfg(not(feature = "profile-heap"))]
56                unsafe {
57                    sol_log_compute_units_start(
58                        PROFILE_ID.as_ptr() as u64,
59                        PROFILE_ID.len() as u64,
60                        0u64,
61                        0u64,
62                        0
63                    );
64                }
65            }
66        };
67
68        let profile_end = quote! {
69                   #[cfg(all(target_os = "solana", feature = "profile-program"))]
70                   {
71                       extern "C" {
72                           #[inline(always)]
73                           fn sol_log_compute_units_end(id_addr: u64, id_len: u64, heap_value: u64, with_heap: u64, _arg5: u64);
74                       }
75                       // Dynamic padding calculated at compile time
76                       const PROFILE_ID: &str = concat!(#fn_name_str, "\n", #front_padding, file!(), ":", line!(), #back_padding);
77                       #[cfg(feature = "profile-heap")]
78                       unsafe {
79                           sol_log_compute_units_end(
80                               PROFILE_ID.as_ptr() as u64,
81                               PROFILE_ID.len() as u64,
82                               ::light_heap::GLOBAL_ALLOCATOR.get_used_heap(),
83                               1u64,
84                               0
85                           );
86                       }
87
88                       #[cfg(target_os = "solana")]
89        #[cfg(not(feature = "profile-heap"))]
90                       unsafe {
91                           sol_log_compute_units_end(
92                               PROFILE_ID.as_ptr() as u64,
93                               PROFILE_ID.len() as u64,
94                               0u64,
95                               0u64,
96                               0
97                           );
98                       }
99                   }
100               };
101
102        // Build the new function body by wrapping the original body with profiling calls
103        let original_stmts = &original_body.stmts;
104        let new_body = if returns_value {
105            quote! {
106                {
107                    #profile_start
108                    let __result = {
109                        #(#original_stmts)*
110                    };
111                    #profile_end
112                    __result
113                }
114            }
115        } else {
116            quote! {
117                {
118                    #profile_start
119                    #(#original_stmts)*
120                    #profile_end
121                }
122            }
123        };
124
125        // Filter out the profile attribute, add inline(always), and replace the function body
126        f.attrs.retain(|a| !a.path().is_ident("profile"));
127        #[cfg(feature = "inline")]
128        {
129            f.attrs.push(syn::parse_quote!(#[inline(always)]));
130        }
131        f.block = syn::parse2(new_body).unwrap();
132
133        TokenStream::from(quote! {
134            #f
135        })
136    }
137}
light_profiler_macro/lib.rs

light_profiler_macro/
lib.rs
