radix_engine_profiling_derive/
lib.rs

1use proc_macro::TokenStream;
2#[cfg(feature = "resource_tracker")]
3use proc_macro2::Span;
4#[cfg(feature = "resource_tracker")]
5use quote::{quote, ToTokens};
6#[cfg(feature = "resource_tracker")]
7use syn::{
8    parse::{Parse, Parser},
9    parse_quote, FnArg,
10    Pat::Ident,
11    Type::{Path, Reference},
12};
13
14/// Empty implementation for compilation without 'resource_tracker' feature.
15#[cfg(not(feature = "resource_tracker"))]
16#[proc_macro_attribute]
17pub fn trace_resources(_attr: TokenStream, input: TokenStream) -> TokenStream {
18    input
19}
20
21/// 'trace_resources' attribute macro is used to log function execution cpu instructions count
22/// during QEMU emulation.
23/// Reauires 'resource_tracker' feature.
24///
25/// Macro parameters:
26/// - log=VALUE
27///   Outputs VALUE, multiple log parameters can be used. Output instruction is added before
28///   original function block execution.
29/// - log_after=VALUE
30///   Outputs VALUE, multiple log_after parameters can be used. Output instruction is added after
31///   original function block execution.
32/// - log_after=ret
33///   Outputs original function block return value.
34/// - info="SOME_STRING"
35///   Outputs SOME_STRING.
36///
37/// VALUE can be:
38/// - identifier of signed/unsigned integer, boolean, &str or must implements conversion to string trait
39/// - method call
40/// - block instructions
41///
42///  Complex example:
43///  #[trace_resources(info="function from module X", log=param, log={data.len()}, log_after={param + 1}, log_after=ret)]
44///  fn test(param: u64, data: &mut Vec<u8>) -> u64 {
45///   ...
46///  }
47///
48/// Simple example which will output only instructions count:
49///  #[trace_resources]
50///  fn test(param: u64, data: &mut Vec<u8>) -> u64 {
51///   ...
52///  }
53///
54#[cfg(all(target_family = "unix", feature = "resource_tracker"))]
55#[proc_macro_attribute]
56pub fn trace_resources(attr: TokenStream, input: TokenStream) -> TokenStream {
57    let args_parsed =
58        syn::punctuated::Punctuated::<syn::ExprAssign, syn::Token![,]>::parse_terminated
59            .parse(attr.clone())
60            .expect("Wrong arguments passed");
61
62    let mut log_ident: Vec<syn::Ident> = Vec::new();
63    let mut log_expr_quote = Vec::new();
64    let mut log_ident_after: Vec<syn::Ident> = Vec::new();
65    let mut log_expr_after_quote = Vec::new();
66
67    let mut additional_items: Vec<proc_macro2::TokenStream> = Vec::new();
68    let mut additional_items_after: Vec<proc_macro2::TokenStream> = Vec::new();
69
70    // parse attributes
71    for (idx, i) in args_parsed.into_iter().enumerate() {
72        if let Ok(left_ident) = syn::parse::<syn::Ident>(i.left.as_ref().to_token_stream().into()) {
73            let left = left_ident.to_string().clone();
74            let right_arg = i.right.as_ref().to_token_stream();
75            match left.as_str() {
76                "info" => {
77                    if let Ok(right) = syn::parse::<syn::LitStr>(right_arg.into()) {
78                        let info_value = right.value();
79                        additional_items.push( quote!{ OutputParam { name: "info".into(), value: OutputParamValue::Literal(#info_value.into())} } );
80                    }
81                }
82                "log" => {
83                    if let Ok(right) = syn::parse::<syn::Ident>(right_arg.clone().into()) {
84                        // log variable
85                        log_ident.push(right);
86                    } else if let Ok(right) =
87                        syn::parse::<syn::ExprMethodCall>(right_arg.clone().into())
88                    {
89                        // log method call result
90                        let var = syn::Ident::new(&format!("arg{}", idx), Span::call_site());
91                        log_expr_quote.push(quote! { let #var = #right; });
92                        let var_s = var.to_string();
93                        additional_items.push( quote!{ OutputParam { name: #var_s.into(), value: OutputParamValue::Literal(format!("{:?}", #var).into())} } );
94                    } else if let Ok(right) = syn::parse::<syn::ExprBlock>(right_arg.clone().into())
95                    {
96                        // log block result
97                        let var = syn::Ident::new(&format!("arg{}", idx), Span::call_site());
98                        log_expr_quote.push(quote! { let #var = #right; });
99                        let var_s = var.to_string();
100                        additional_items.push( quote!{ OutputParam { name: #var_s.into(), value: OutputParamValue::Literal(format!("{:?}", #var).into())} } );
101                    } else if let Ok(right) = syn::parse::<syn::ExprUnary>(right_arg.into()) {
102                        // log deref result
103                        let var_name = match right.expr.as_ref() {
104                            syn::Expr::Path(p) => p.path.segments.last().unwrap().ident.clone(),
105                            _ => panic!("Not supported path expression: {:?}", right.expr),
106                        };
107                        let var =
108                            syn::Ident::new(&format!("{}_deref", var_name), Span::call_site());
109                        log_expr_quote.push(quote! { let #var = #right; });
110                        let var_s = var_name.to_string();
111                        additional_items.push( quote!{ OutputParam { name: #var_s.into(), value: OutputParamValue::Literal(format!("{:?}", #var).into())} } );
112                    } else {
113                        panic!("Wrong log value type: {:?}", i.right.as_ref())
114                    }
115                }
116                "log_after" => {
117                    if let Ok(right) = syn::parse::<syn::Ident>(right_arg.clone().into()) {
118                        // log variable
119                        if right == "ret" {
120                            // TODO: optimise basing on function return type
121                            additional_items_after.push( quote!{ OutputParam { name: "ret".into(), value: OutputParamValue::Literal(format!("{:?}", ret).into())} } );
122                        } else {
123                            log_ident_after.push(right);
124                        }
125                    } else if let Ok(right) =
126                        syn::parse::<syn::ExprMethodCall>(right_arg.clone().into())
127                    {
128                        // lob method call result
129                        let var = syn::Ident::new(&format!("arg{}", idx), Span::call_site());
130                        log_expr_after_quote.push(quote! { let #var = #right; });
131                        let var_s = var.to_string();
132                        additional_items_after.push( quote!{ OutputParam { name: #var_s.into(), value: OutputParamValue::Literal(format!("{:?}", #var).into())} } );
133                    } else if let Ok(right) = syn::parse::<syn::ExprBlock>(right_arg.clone().into())
134                    {
135                        // lob block result
136                        let var = syn::Ident::new(&format!("arg{}", idx), Span::call_site());
137                        log_expr_after_quote.push(quote! { let #var = #right; });
138                        let var_s = var.to_string();
139                        additional_items_after.push( quote!{ OutputParam { name: #var_s.into(), value: OutputParamValue::Literal(format!("{:?}", #var).into())} } );
140                    } else if let Ok(right) = syn::parse::<syn::ExprUnary>(right_arg.into()) {
141                        // log deref result
142                        let var_name = match right.expr.as_ref() {
143                            syn::Expr::Path(p) => p.path.segments.last().unwrap().ident.clone(),
144                            _ => panic!("Not supported path expression: {:?}", right.expr),
145                        };
146                        let var =
147                            syn::Ident::new(&format!("{}_deref", var_name), Span::call_site());
148                        log_expr_quote.push(quote! { let #var = #right; });
149                        let var_s = var_name.to_string();
150                        additional_items_after.push( quote!{ OutputParam { name: #var_s.into(), value: OutputParamValue::Literal(format!("{:?}", #var).into())} } );
151                    } else {
152                        panic!("Wrong log_after value type: {:?}", i.right.as_ref())
153                    }
154                }
155                s => panic!("Wrong argument: {}", s),
156            }
157        }
158    }
159
160    // prepare output
161    let output = if let Ok(mut item) = syn::Item::parse.parse(input.clone()) {
162        match item {
163            syn::Item::Fn(ref mut item_fn) => {
164                let mut arg_evaluate = quote! {};
165                for i in log_expr_quote {
166                    arg_evaluate = quote! { #arg_evaluate #i };
167                }
168                let mut arg_evaluate_after = quote! {};
169                for i in log_expr_after_quote {
170                    arg_evaluate_after = quote! { #arg_evaluate_after #i };
171                }
172
173                let args_quote_array =
174                    create_params(&log_ident, item_fn.clone(), &additional_items);
175                let args_after_quote_array =
176                    create_params(&log_ident_after, item_fn.clone(), &additional_items_after);
177                let original_block = &mut item_fn.block;
178                let fn_signature = item_fn.sig.ident.to_string();
179
180                // new function block
181                item_fn.block = Box::new(parse_quote! {{
182                    extern crate radix_engine_profiling;
183                    use radix_engine_profiling::{QEMU_PLUGIN, data_analyzer::{OutputParam, OutputParamValue}};
184                    #arg_evaluate;
185                    #args_quote_array;
186                    QEMU_PLUGIN.with(|v| {
187                        v.borrow_mut().start_counting(#fn_signature, qemu_call_args.as_slice());
188                    });
189                    let ret = #original_block;
190                    QEMU_PLUGIN.with(|v| {
191                        v.borrow_mut().stop_counting(#fn_signature, qemu_call_args.as_slice());
192                    });
193                    #arg_evaluate_after;
194                    #args_after_quote_array;
195                    ret
196                }});
197                item.into_token_stream()
198            }
199            _ => syn::Error::new_spanned(item, "#[trace_resources] is not supported for this item")
200                .to_compile_error(),
201        }
202    } else {
203        let input2 = proc_macro2::TokenStream::from(input);
204        syn::Error::new_spanned(input2, "expected `fn` item").to_compile_error()
205    };
206
207    output.into()
208}
209
210/// Helper function
211#[cfg(all(target_family = "unix", feature = "resource_tracker"))]
212fn create_params(
213    ident: &Vec<syn::Ident>,
214    fn_sig: syn::ItemFn,
215    additional_items: &Vec<proc_macro2::TokenStream>,
216) -> proc_macro2::TokenStream {
217    let mut args_quote = Vec::new();
218    for arg_ident in ident {
219        let mut arg_found = false;
220        for fn_arg in fn_sig.sig.inputs.iter() {
221            match fn_arg {
222                FnArg::Typed(v) => {
223                    match v.pat.as_ref() {
224                        Ident(pi) => {
225                            if pi.ident == *arg_ident {
226                                arg_found = true;
227                                let var_name = pi.ident.to_string();
228                                match v.ty.as_ref() {
229                                    Path(tp) => {
230                                        // u8..i64, bool
231                                        if let Some(p) = tp.path.segments.last() {
232                                            if p.ident == "u8"
233                                                || p.ident == "u16"
234                                                || p.ident == "u32"
235                                                || p.ident == "u64"
236                                            {
237                                                args_quote.push( quote!{ OutputParam { name: #var_name.into(), value: OutputParamValue::NumberU64( #pi as u64 )} } );
238                                                break;
239                                            } else if p.ident == "i8"
240                                                || p.ident == "i16"
241                                                || p.ident == "i32"
242                                                || p.ident == "i64"
243                                            {
244                                                args_quote.push( quote!{ OutputParam { name: #var_name.into(), value: OutputParamValue::NumberI64( #pi as i64 )} } );
245                                                break;
246                                            } else if p.ident == "bool" {
247                                                args_quote.push( quote!{ OutputParam { name: #var_name.into(), value: OutputParamValue::NumberU64( #pi as u64 )} } );
248                                                break;
249                                            } else {
250                                                args_quote.push( quote!{ OutputParam { name: #var_name.into(), value: OutputParamValue::Literal(format!("{:?}", #pi).into())} } );
251                                                break;
252                                            }
253                                        }
254                                    }
255                                    Reference(tr) => {
256                                        //&str
257                                        match tr.elem.as_ref() {
258                                            Path(tp) => {
259                                                if let Some(p) = tp.path.segments.last() {
260                                                    if p.ident == "str" {
261                                                        args_quote.push( quote!{ OutputParam { name: #var_name.into(), value: OutputParamValue::Literal(#pi.into())} } );
262                                                        break;
263                                                    } else {
264                                                        panic!(
265                                                            "Not supported arg type: {}",
266                                                            p.ident
267                                                        );
268                                                    }
269                                                }
270                                            }
271                                            _ => (),
272                                        }
273                                    }
274                                    _ => (),
275                                }
276                            }
277                        }
278                        _ => (),
279                    }
280                }
281                _ => (),
282            }
283        }
284        if !arg_found {
285            panic!("Arg: {} not found", arg_ident.to_string());
286        }
287    }
288
289    let mut args_quote_len = args_quote.len();
290    let mut args_quote_array = quote! {};
291    if args_quote_len > 0 {
292        let q0 = &args_quote[0];
293        args_quote_array = quote! { #q0 };
294        if args_quote_len > 1 {
295            for i in 1..args_quote_len {
296                let q1 = &args_quote[i];
297                args_quote_array = quote! { #args_quote_array, #q1 };
298            }
299        }
300        for i in additional_items {
301            args_quote_len += 1;
302            args_quote_array = quote! { #args_quote_array, #i };
303        }
304    } else if additional_items.len() > 0 {
305        let q0 = &additional_items[0];
306        args_quote_array = quote! { #q0 };
307        if additional_items.len() > 1 {
308            for i in 1..additional_items.len() {
309                let q1 = &additional_items[i];
310                args_quote_array = quote! { #args_quote_array, #q1 };
311            }
312        }
313        args_quote_len = additional_items.len();
314    }
315
316    if args_quote_len > 0 {
317        args_quote_array =
318            quote! { let qemu_call_args: [OutputParam; #args_quote_len] = [#args_quote_array] };
319    } else {
320        args_quote_array = quote! { let qemu_call_args: [OutputParam; 0] = []; };
321    }
322
323    args_quote_array
324}
radix_engine_profiling_derive/lib.rs

radix_engine_profiling_derive/
lib.rs
