Skip to main content

rustmeter_beacon_proc_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{
4    Ident, ItemFn, LitStr, Result, Token,
5    parse::{Parse, ParseStream},
6    parse_macro_input,
7};
8
9extern crate proc_macro;
10
11/// Helper struct to parse arguments for the `monitor_fn` attribute macro
12struct MonitorArgs {
13    name: Option<String>,
14}
15
16impl Parse for MonitorArgs {
17    fn parse(input: ParseStream) -> Result<Self> {
18        let mut name = None;
19        if input.is_empty() {
20            return Ok(MonitorArgs { name });
21        }
22
23        // Case 1: #[monitor_fn("Name")]
24        // `lookahead` checks if the next token is a string literal
25        if input.peek(syn::LitStr) {
26            let lit: LitStr = input.parse()?;
27            name = Some(lit.value());
28        }
29        // Case 2: Key-Value Pair: #[monitor_fn(name = "Name")]
30        else if input.peek(syn::Ident) {
31            let key: Ident = input.parse()?;
32            if key == "name" {
33                input.parse::<Token![=]>()?; // Consume the '='
34                let lit: LitStr = input.parse()?;
35                name = Some(lit.value());
36            } else {
37                return Err(syn::Error::new(
38                    key.span(),
39                    "Unknown argument (expected 'name')",
40                ));
41            }
42        }
43
44        // More arguments could be parsed here in the future
45
46        Ok(MonitorArgs { name })
47    }
48}
49
50/// Instruments a function to log execution for rustmeter
51///
52/// # Examples
53///
54/// Basic usage using the function's name:
55///
56/// ```rust
57/// #[monitor_fn]
58/// fn process_data(data: u8) {
59///     // Function implementation
60/// }
61/// ```
62#[proc_macro_attribute]
63pub fn monitor_fn(attr: TokenStream, item: TokenStream) -> TokenStream {
64    let input = parse_macro_input!(item as ItemFn);
65    let args = parse_macro_input!(attr as MonitorArgs);
66
67    let fn_name = &input.sig.ident;
68    let vis = &input.vis;
69    let sig = &input.sig;
70    let block = &input.block;
71    let attrs = &input.attrs; // Important: Keep other attributes (e.g., #[inline])
72
73    let mut output_name = fn_name.to_string();
74
75    // Handle output name from args (if provided)
76    if let Some(custom_name) = args.name {
77        output_name = custom_name;
78    }
79
80    if input.sig.asyncness.is_some() {
81        // ASYNC FUNCTION not supported yet
82        quote! {
83            compile_error!("`monitor_fn` macro does not support async functions yet.");
84        }
85        .into()
86    } else {
87        // SYNC FUNCTION
88        quote! {
89            #(#attrs)*
90            #vis #sig {
91                {
92                    let core_id = rustmeter_beacon::core_id::get_current_core_id();
93
94                    // Get or register monitor ID
95                    use rustmeter_beacon::monitors::VALUE_MONITOR_REGISTRY;
96                    let (local_id, registered_newly) = rustmeter_beacon::get_static_id_by_registry!(
97                        rustmeter_beacon::monitors::CODE_MONITOR_REGISTRY
98                    );
99
100                    // Send TypeDefinition event if newly registered
101                    if registered_newly {
102                        let fn_addr = #fn_name as usize;
103                        let payload = rustmeter_beacon::protocol::TypeDefinitionPayload::FunctionMonitor {
104                            monitor_id: local_id as u8,
105                            fn_address: fn_addr as u32,
106                        };
107                        rustmeter_beacon::tracing::write_tracing_event(
108                            rustmeter_beacon::protocol::EventPayload::TypeDefinition(payload)
109                        );
110                    
111                        rustmeter_beacon::monitors::defmt_trace_new_function_monitor(#output_name, local_id);
112                    }
113
114                    // Create guard to signal end of scope
115                    let _guard = rustmeter_beacon::monitors::DropGuard::new(|| {
116                        rustmeter_beacon::protocol::raw_writers::write_monitor_end();
117                    });
118
119                    // Send MonitorStart event (after guard-created to lower tracing impact on measured scope)
120                    rustmeter_beacon::protocol::raw_writers::write_monitor_start(local_id as u8);
121                
122
123                    // Execute original function body
124                    { #block }
125                }
126            }           
127        }.into()
128    }
129}