1use proc_macro::TokenStream;
2use quote::{quote, quote_spanned};
3use syn::spanned::Spanned;
4
5fn token_stream_with_error(mut tokens: TokenStream, error: syn::Error) -> TokenStream {
6 tokens.extend(TokenStream::from(error.into_compile_error()));
7 tokens
8}
9
10enum TimeFormat {
11 Seconds,
12 Milliseconds,
13 Nanoseconds,
14}
15
16#[proc_macro_attribute]
17pub fn time(attrs: TokenStream, item: TokenStream) -> TokenStream {
18 let time_format: u8 = match attrs.to_string().as_str() {
19 "s" | "seconds" | "" => TimeFormat::Seconds,
20 "ms" | "milliseconds" => TimeFormat::Milliseconds,
21 "ns" | "nanoseconds" => TimeFormat::Nanoseconds,
22 _ => return quote! { compile_error!("attributes can only be s/ms/ns for seconds, milliseconds and nanoseconds respectively") }.into(),
23 } as u8;
24
25 let input: syn::ItemFn = match syn::parse(item.clone()) {
26 Ok(input) => input,
27 Err(error) => return token_stream_with_error(item, error),
28 };
29
30 let name = &input.sig.ident;
31 let inputs = &input.sig.inputs;
32 let body = &*input.block;
33 let ret = &input.sig.output;
34
35 if name != "main" {
36 return quote_spanned! {
37 name.span() => compile_error!("#[time] can only be applied to the main function");
38 }
39 .into();
40 }
41
42 if !inputs.is_empty() {
43 return quote_spanned! {
44 inputs.span() => compile_error!("the main function cannot have any arguments");
45 }
46 .into();
47 }
48
49 let output = quote! {
50 fn main() #ret {
51 use std::time::{Duration, Instant};
52 let start = Instant::now();
53 let ret = {
54 #body
55 };
56
57 let elapsed = start.elapsed();
58
59 match #time_format {
60 1 => {
61 println!("{}ms", elapsed.as_millis());
62 },
63 2 => {
64 println!("{}ns", elapsed.as_nanos());
65 },
66 _ => {
67 println!("{}.{}s", elapsed.as_secs(), elapsed.subsec_millis());
68 },
69 };
70
71 ret
72 }
73 };
74
75 output.into()
76}