1  // Copyright (C) 2019-2021 Alpha-Delta Network Inc.
  2  // This file is part of the alphastd library.
  3  
  4  // The alphastd 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 alphastd 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 alphastd 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]
 21  extern crate quote;
 22  #[cfg(feature = "time")]
 23  #[macro_use]
 24  extern crate syn;
 25  extern crate proc_macro;
 26  
 27  #[cfg(feature = "time")]
 28  const DEFAULT_LEVEL: &str = "debug";
 29  #[cfg(feature = "time")]
 30  const DEFAULT_NAME_PATTERN: &str = "{}";
 31  
 32  #[cfg(feature = "time")]
 33  fn 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")]
 49  fn 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")]
109  fn 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]
134  pub 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 = ::alpha_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]
163  pub fn time(_metadata: proc_macro::TokenStream, input: proc_macro::TokenStream) -> proc_macro::TokenStream {
164      input
165  }