Skip to main content

aleo_std_time/
lib.rs

1// Copyright (C) 2019-2021 Aleo Systems Inc.
2// This file is part of the aleo-std library.
3
4// The aleo-std library is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// The aleo-std library is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with the aleo-std library. If not, see <https://www.gnu.org/licenses/>.
16
17// With credits to PhilipDaniels/logging_timer.
18
19#[cfg(feature = "time")]
20#[macro_use]
21extern crate quote;
22#[cfg(feature = "time")]
23#[macro_use]
24extern crate syn;
25extern crate proc_macro;
26
27#[cfg(feature = "time")]
28const DEFAULT_LEVEL: &str = "debug";
29#[cfg(feature = "time")]
30const DEFAULT_NAME_PATTERN: &str = "{}";
31
32#[cfg(feature = "time")]
33fn extract_literal(token_tree: &proc_macro::TokenTree) -> String {
34    let s = match token_tree {
35        proc_macro::TokenTree::Literal(literal) => literal.to_string(),
36        _ => panic!(
37            "Invalid argument. Specify at most two string literal arguments, for log level and name pattern, in that order."
38        ),
39    };
40
41    // String literals seem to come through including their double quotes. Trim them off.
42    let s = s.trim().trim_matches('"').trim().to_string();
43    s
44}
45
46// We also allow 'Never' to mean disable timer instrumentation
47// altogether. Any casing is allowed.
48#[cfg(feature = "time")]
49fn get_log_level_and_name_pattern(metadata: proc_macro::TokenStream) -> (String, String) {
50    // Grab everything into a vector and filter out any intervening punctuation
51    // (commas come through as TokenTree::Punct(_)).
52    let macro_args: Vec<proc_macro::TokenTree> = metadata
53        .into_iter()
54        .filter(|token| match token {
55            proc_macro::TokenTree::Literal(_) => true,
56            _ => false,
57        })
58        .collect();
59
60    if macro_args.is_empty() {
61        return (DEFAULT_LEVEL.to_string(), DEFAULT_NAME_PATTERN.to_string());
62    }
63
64    if macro_args.len() > 2 {
65        panic!("Specify at most two string literal arguments, for log level and name pattern");
66    }
67
68    let first_arg = extract_literal(&macro_args[0]);
69
70    if first_arg.contains("{}") && macro_args.len() == 2 {
71        panic!("Invalid first argument. Specify the log level as the first argument and the pattern as the second.");
72    }
73
74    let first_arg_lower = first_arg.to_ascii_lowercase();
75    if macro_args.len() == 1 {
76        match first_arg_lower.as_str() {
77            "error" | "warn" | "info" | "debug" | "trace" | "never" => {
78                // User specified a valid log level as their only argument.
79                return (first_arg_lower, DEFAULT_NAME_PATTERN.to_string());
80            }
81            _ => {
82                // User specified something that doesn't look like a log-level.
83                // It may be a pattern with "{}", or it may just be a string.
84                // In any case, consider it to be the pattern and return it
85                // n.b. the original, not the lowered version.
86                return (DEFAULT_LEVEL.to_string(), first_arg.to_string());
87            }
88        }
89    }
90
91    // We have two arguments. We are stricter on the first now, we require
92    // that to be a valid log level.
93    match first_arg_lower.as_str() {
94        "error" | "warn" | "info" | "debug" | "trace" | "never" => {
95            let mut second_arg = extract_literal(&macro_args[1]);
96            if second_arg.is_empty() {
97                second_arg += DEFAULT_NAME_PATTERN;
98            }
99
100            return (first_arg_lower, second_arg.to_string());
101        }
102        _ => {
103            panic!("Invalid first argument. Specify the log level as the first argument and the pattern as the second.")
104        }
105    }
106}
107
108#[cfg(feature = "time")]
109fn get_timer_name(name_pattern: &str, function_name: &str) -> String {
110    let function_name_with_parenthesis = format!("{}()", function_name);
111    let timer_name = name_pattern.replacen("{}", &function_name_with_parenthesis, 1);
112    timer_name
113}
114
115/// Instruments the function with an `timer!`, which logs two messages, one at the start
116/// of the function and one at the end of execution stating the elapsed time.
117///
118/// The attribute accepts two string literals as arguments. The first is the log level,
119/// valid values of which are "error", "warn", "info", "debug", "trace" or "never".
120/// The default value is "debug". "never" can be used to temporarily disable instrumentation
121/// of the function without deleting the attribute.
122///
123/// The second argument is the function name pattern. The pattern is helpful to
124/// disambiguate functions when you have many functions in the same module with the same
125/// name: `new` might occur many times on different structs, for example. In the pattern,
126/// "{}" will be replaced with the name of the function.
127///
128/// Examples:
129///     #[time]                                 // Use default print level
130///     #[time("FirstStruct::{}")]              // Prints "FirstStruct::new()"
131///     #[time("never")]                        // Turn off instrumentation at compile time
132#[cfg(feature = "time")]
133#[proc_macro_attribute]
134pub fn time(metadata: proc_macro::TokenStream, input: proc_macro::TokenStream) -> proc_macro::TokenStream {
135    let (level, name_pattern) = get_log_level_and_name_pattern(metadata);
136
137    if level != "never" {
138        let input_fn: syn::ItemFn = parse_macro_input!(input as syn::ItemFn);
139        let visibility = input_fn.vis;
140        let ident = input_fn.sig.ident;
141        let inputs = input_fn.sig.inputs;
142        let output = input_fn.sig.output;
143        let generics = &input_fn.sig.generics;
144        let where_clause = &input_fn.sig.generics.where_clause;
145        let block = input_fn.block;
146
147        let timer_name = get_timer_name(&name_pattern, &ident.to_string());
148
149        (quote!(
150            #visibility fn #ident #generics (#inputs) #output #where_clause {
151                let _tmr = ::aleo_std::prelude::timer!(#timer_name);
152                #block
153            }
154        ))
155        .into()
156    } else {
157        proc_macro::TokenStream::from(input).into()
158    }
159}
160
161#[cfg(not(feature = "time"))]
162#[proc_macro_attribute]
163pub fn time(_metadata: proc_macro::TokenStream, input: proc_macro::TokenStream) -> proc_macro::TokenStream {
164    input
165}