1//! Macros for [Divan](https://github.com/nvzqz/divan), a statistically-comfy
2//! benchmarking library brought to you by [Nikolai Vazquez](https://hachyderm.io/@nikolai).
3//!
4//! See [`divan`](https://docs.rs/divan) crate for documentation.
5
6use proc_macro::TokenStream;
7use quote::{quote, ToTokens};
8
9mod attr_options;
10mod tokens;
11
12use attr_options::*;
13use syn::{Expr, FnArg};
14
15#[derive(Clone, Copy)]
16enum Macro<'a> {
17 Bench { fn_sig: &'a syn::Signature },
18 BenchGroup,
19}
20
21impl Macro<'_> {
22 fn name(&self) -> &'static str {
23 match self {
24 Self::Bench { .. } => "bench",
25 Self::BenchGroup => "bench_group",
26 }
27 }
28}
29
30/// Lists of comma-separated `#[cfg]` parameters.
31mod systems {
32 use super::*;
33
34 pub fn elf() -> proc_macro2::TokenStream {
35 quote! {
36 target_os = "android",
37 target_os = "dragonfly",
38 target_os = "freebsd",
39 target_os = "fuchsia",
40 target_os = "haiku",
41 target_os = "illumos",
42 target_os = "linux",
43 target_os = "netbsd",
44 target_os = "openbsd",
45 target_os = "wasi",
46 target_os = "emscripten"
47 }
48 }
49
50 pub fn mach_o() -> proc_macro2::TokenStream {
51 quote! {
52 target_os = "ios",
53 target_os = "macos",
54 target_os = "tvos",
55 target_os = "watchos"
56 }
57 }
58}
59
60/// Attributes applied to a `static` containing a pointer to a function to run
61/// before `main`.
62fn pre_main_attrs() -> proc_macro2::TokenStream {
63 let elf: TokenStream = systems::elf();
64 let mach_o: TokenStream = systems::mach_o();
65
66 quote! {
67 #[used]
68 #[cfg_attr(windows, link_section = ".CRT$XCU")]
69 #[cfg_attr(any(#elf), link_section = ".init_array")]
70 #[cfg_attr(any(#mach_o), link_section = "__DATA,__mod_init_func,mod_init_funcs")]
71 }
72}
73
74fn unsupported_error(attr_name: &str) -> proc_macro2::TokenStream {
75 let elf: TokenStream = systems::elf();
76 let mach_o: TokenStream = systems::mach_o();
77
78 let error: String = format!("Unsupported target OS for `#[divan::{attr_name}]`");
79
80 quote! {
81 #[cfg(not(any(windows, #elf, #mach_o)))]
82 ::std::compile_error!(#error);
83 }
84}
85
86#[proc_macro_attribute]
87pub fn bench(options: TokenStream, item: TokenStream) -> TokenStream {
88 let option_none = tokens::option_none();
89 let option_some = tokens::option_some();
90
91 let fn_item = item.clone();
92 let fn_item = syn::parse_macro_input!(fn_item as syn::ItemFn);
93 let fn_sig = &fn_item.sig;
94
95 let attr = Macro::Bench { fn_sig };
96 let attr_name = attr.name();
97
98 let options = match AttrOptions::parse(options, attr) {
99 Ok(options) => options,
100 Err(compile_error) => return compile_error,
101 };
102
103 // Items needed by generated code.
104 let AttrOptions { private_mod, .. } = &options;
105
106 let fn_ident = &fn_sig.ident;
107 let fn_name = fn_ident.to_string();
108 let fn_name_pretty = fn_name.strip_prefix("r#").unwrap_or(&fn_name);
109
110 // Find any `#[ignore]` attribute so that we can use its span to help
111 // compiler diagnostics.
112 let ignore_attr_ident =
113 fn_item.attrs.iter().map(|attr| attr.meta.path()).find(|path| path.is_ident("ignore"));
114
115 // If the function is `extern "ABI"`, it is wrapped in a Rust-ABI function.
116 let is_extern_abi = fn_sig.abi.is_some();
117
118 let fn_args = &fn_sig.inputs;
119
120 let type_param: Option<(usize, &syn::TypeParam)> = fn_sig
121 .generics
122 .params
123 .iter()
124 .enumerate()
125 .filter_map(|(i, param)| match param {
126 syn::GenericParam::Type(param) => Some((i, param)),
127 _ => None,
128 })
129 .next();
130
131 let const_param: Option<(usize, &syn::ConstParam)> = fn_sig
132 .generics
133 .params
134 .iter()
135 .enumerate()
136 .filter_map(|(i, param)| match param {
137 syn::GenericParam::Const(param) => Some((i, param)),
138 _ => None,
139 })
140 .next();
141
142 let is_type_before_const = match (type_param, const_param) {
143 (Some((t, _)), Some((c, _))) => t < c,
144 _ => false,
145 };
146
147 // Prefixed with "__" to prevent IDEs from recommending using this symbol.
148 //
149 // The static is local to intentionally cause a compile error if this
150 // attribute is used multiple times on the same function.
151 let static_ident = syn::Ident::new(
152 &format!("__DIVAN_BENCH_{}", fn_name_pretty.to_uppercase()),
153 fn_ident.span(),
154 );
155
156 let meta = entry_meta_expr(&fn_name, &options, ignore_attr_ident);
157
158 let bench_entry_runner = quote! { #private_mod::BenchEntryRunner };
159
160 // Creates a `__DIVAN_ARGS` global variable to be used in the entry.
161 let bench_args_global = if options.args_expr.is_some() {
162 quote! {
163 static __DIVAN_ARGS: #private_mod::BenchArgs = #private_mod::BenchArgs::new();
164 }
165 } else {
166 Default::default()
167 };
168
169 // The last argument type is used as the only `args` item type because we
170 // currently only support one runtime argument.
171 let last_arg_type = if options.args_expr.is_some() {
172 fn_args.last().map(|arg| match arg {
173 FnArg::Receiver(arg) => &*arg.ty,
174 FnArg::Typed(arg) => &*arg.ty,
175 })
176 } else {
177 None
178 };
179
180 let last_arg_type_tokens = last_arg_type
181 .map(|ty| match ty {
182 // Remove lifetime from references to not use the lifetime outside
183 // of its declaration. This allows benchmarks to take arguments with
184 // lifetimes.
185 syn::Type::Reference(ty) if ty.lifetime.is_some() => {
186 let mut ty = ty.clone();
187 ty.lifetime = None;
188 ty.to_token_stream()
189 }
190
191 _ => ty.to_token_stream(),
192 })
193 .unwrap_or_default();
194
195 // Some argument literals need an explicit type.
196 let arg_return_tokens = options
197 .args_expr
198 .as_ref()
199 .map(|args| match args {
200 // Empty array.
201 Expr::Array(args) if args.elems.is_empty() => quote! {
202 -> [#last_arg_type_tokens; 0]
203 },
204
205 _ => Default::default(),
206 })
207 .unwrap_or_default();
208
209 // Creates a function expr for the benchmarking function, optionally
210 // monomorphized with generic parameters.
211 let make_bench_fn = |generics: &[&dyn ToTokens]| {
212 let mut fn_expr = if generics.is_empty() {
213 // Use identifier as-is.
214 fn_ident.to_token_stream()
215 } else {
216 // Apply generic arguments.
217 quote! { #fn_ident::< #(#generics),* > }
218 };
219
220 // Handle function arguments.
221 match (fn_args.len(), &options.args_expr) {
222 // Simple benchmark with no arguments provided.
223 (0, None) => {
224 // Wrap in Rust ABI.
225 if is_extern_abi {
226 fn_expr = quote! { || #fn_expr() };
227 }
228
229 quote! {
230 #bench_entry_runner::Plain(|divan /* Bencher */| divan.bench(#fn_expr))
231 }
232 }
233
234 // `args` option used without function arguments; handled earlier in
235 // `AttrOptions::parse`.
236 (0, Some(_)) => unreachable!(),
237
238 // `Bencher` function argument.
239 (1, None) => {
240 // Wrap in Rust ABI.
241 if is_extern_abi {
242 fn_expr = quote! { |divan /* Bencher */| #fn_expr(divan) };
243 }
244
245 quote! { #bench_entry_runner::Plain(#fn_expr) }
246 }
247
248 // Function argument comes from `args` option.
249 (1, Some(args)) => quote! {
250 #bench_entry_runner::Args(|| __DIVAN_ARGS.runner(
251 || #arg_return_tokens { #args },
252
253 |arg| #private_mod::ToStringHelper(arg).to_string(),
254
255 |divan, __divan_arg| divan.bench(|| #fn_expr(
256 #private_mod::Arg::<#last_arg_type_tokens>::get(__divan_arg)
257 )),
258 ))
259 },
260
261 // `Bencher` and `args` option function arguments.
262 (2, Some(args)) => quote! {
263 #bench_entry_runner::Args(|| __DIVAN_ARGS.runner(
264 || #arg_return_tokens { #args },
265
266 |arg| #private_mod::ToStringHelper(arg).to_string(),
267
268 |divan, __divan_arg| #fn_expr(
269 divan,
270 #private_mod::Arg::<#last_arg_type_tokens>::get(__divan_arg),
271 ),
272 ))
273 },
274
275 // Ensure `args` is set if arguments are provided after `Bencher`.
276 (_, None) => quote! {
277 ::std::compile_error!(::std::concat!(
278 "expected 'args' option containing '",
279 ::std::stringify!(#last_arg_type_tokens),
280 "'",
281 ))
282 },
283
284 // `args` option used with unsupported number of arguments; handled
285 // earlier in `AttrOptions::parse`.
286 (_, Some(_)) => unreachable!(),
287 }
288 };
289
290 let pre_main_attrs = pre_main_attrs();
291 let unsupported_error = unsupported_error(attr_name);
292
293 // Creates a `GroupEntry` static for generic benchmarks.
294 let make_generic_group = |generic_benches: proc_macro2::TokenStream| {
295 let entry = quote! {
296 #private_mod::GroupEntry {
297 meta: #meta,
298 generic_benches: #option_some({ #generic_benches }),
299 }
300 };
301
302 quote! {
303 #unsupported_error
304
305 // Push this static into `GROUP_ENTRIES` before `main` is called.
306 static #static_ident: #private_mod::GroupEntry = {
307 {
308 // Add `push` to the initializer section.
309 #pre_main_attrs
310 static PUSH: extern "C" fn() = push;
311
312 extern "C" fn push() {
313 static NODE: #private_mod::EntryList<#private_mod::GroupEntry>
314 = #private_mod::EntryList::new(&#static_ident);
315
316 #private_mod::GROUP_ENTRIES.push(&NODE);
317 }
318 }
319
320 // All generic entries share the same `BenchArgs` instance for
321 // efficiency and to ensure all entries use the same values, or
322 // at least the same names in the case of interior mutability.
323 #bench_args_global
324
325 #entry
326 };
327 }
328 };
329
330 // Creates a `GenericBenchEntry` expr for a generic benchmark instance.
331 let make_generic_bench_entry =
332 |ty: Option<&dyn ToTokens>, const_value: Option<&dyn ToTokens>| {
333 let generic_const_value = const_value.map(|const_value| quote!({ #const_value }));
334
335 let generics: Vec<&dyn ToTokens> = {
336 let mut generics = Vec::new();
337
338 generics.extend(generic_const_value.as_ref().map(|t| t as &dyn ToTokens));
339 generics.extend(ty);
340
341 if is_type_before_const {
342 generics.reverse();
343 }
344
345 generics
346 };
347
348 let bench_fn = make_bench_fn(&generics);
349
350 let type_value = match ty {
351 Some(ty) => quote! {
352 #option_some(#private_mod::EntryType::new::<#ty>())
353 },
354 None => option_none.clone(),
355 };
356
357 let const_value = match const_value {
358 Some(const_value) => quote! {
359 #option_some(#private_mod::EntryConst::new(&#const_value))
360 },
361 None => option_none.clone(),
362 };
363
364 quote! {
365 #private_mod::GenericBenchEntry {
366 group: &#static_ident,
367 bench: #bench_fn,
368 ty: #type_value,
369 const_value: #const_value,
370 }
371 }
372 };
373
374 let generated_items: proc_macro2::TokenStream = match &options.generic.consts {
375 // Only specified `types = []` or `consts = []`; generate nothing.
376 _ if options.generic.is_empty() => Default::default(),
377
378 None => match &options.generic.types {
379 // No generics; generate a simple benchmark entry.
380 None => {
381 let bench_fn = make_bench_fn(&[]);
382
383 let entry = quote! {
384 #private_mod::BenchEntry {
385 meta: #meta,
386 bench: #bench_fn,
387 }
388 };
389
390 quote! {
391 // Push this static into `BENCH_ENTRIES` before `main` is
392 // called.
393 static #static_ident: #private_mod::BenchEntry = {
394 {
395 // Add `push` to the initializer section.
396 #pre_main_attrs
397 static PUSH: extern "C" fn() = push;
398
399 extern "C" fn push() {
400 static NODE: #private_mod::EntryList<#private_mod::BenchEntry>
401 = #private_mod::EntryList::new(&#static_ident);
402
403 #private_mod::BENCH_ENTRIES.push(&NODE);
404 }
405 }
406
407 #bench_args_global
408
409 #entry
410 };
411 }
412 }
413
414 // Generate a benchmark group entry with generic benchmark entries.
415 Some(GenericTypes::List(generic_types)) => {
416 let generic_benches =
417 generic_types.iter().map(|ty| make_generic_bench_entry(Some(&ty), None));
418
419 make_generic_group(quote! {
420 &[&[#(#generic_benches),*]]
421 })
422 }
423 },
424
425 // Generate a benchmark group entry with generic benchmark entries.
426 Some(Expr::Array(generic_consts)) => {
427 let consts_count = generic_consts.elems.len();
428 let const_type = &const_param.unwrap().1.ty;
429
430 let generic_benches = options.generic.types_iter().map(|ty| {
431 let generic_benches = (0..consts_count).map(move |i| {
432 let const_value = quote! { __DIVAN_CONSTS[#i] };
433 make_generic_bench_entry(ty, Some(&const_value))
434 });
435
436 // `static` is necessary because `EntryConst` uses interior
437 // mutability to cache the `ToString` result.
438 quote! {
439 static __DIVAN_GENERIC_BENCHES: [#private_mod::GenericBenchEntry; #consts_count] = [#(#generic_benches),*];
440 &__DIVAN_GENERIC_BENCHES
441 }
442 });
443
444 make_generic_group(quote! {
445 // We refer to our own slice because it:
446 // - Type-checks values, even if `generic_benches` is empty
447 // because the user set `types = []`
448 // - Prevents re-computing constants, which can slightly improve
449 // compile time given that Miri is slow
450 const __DIVAN_CONSTS: &[#const_type] = &#generic_consts;
451
452 &[#({ #generic_benches }),*]
453 })
454 }
455
456 // Generate a benchmark group entry with generic benchmark entries over
457 // an expression of constants.
458 //
459 // This is limited to a maximum of 20 because we need some constant to
460 // instantiate each function instance.
461 Some(generic_consts) => {
462 // The maximum number of elements for non-array expressions.
463 const MAX_EXTERN_COUNT: usize = 20;
464
465 let const_type = &const_param.unwrap().1.ty;
466
467 let generic_benches = options.generic.types_iter().map(|ty| {
468 let generic_benches = (0..MAX_EXTERN_COUNT).map(move |i| {
469 let const_value = quote! {
470 // Fallback to the first constant if out of bounds.
471 __DIVAN_CONSTS[if #i < __DIVAN_CONST_COUNT { #i } else { 0 }]
472 };
473 make_generic_bench_entry(ty, Some(&const_value))
474 });
475
476 // `static` is necessary because `EntryConst` uses interior
477 // mutability to cache the `ToString` result.
478 quote! {
479 static __DIVAN_GENERIC_BENCHES: [#private_mod::GenericBenchEntry; __DIVAN_CONST_COUNT]
480 = match #private_mod::shrink_array([#(#generic_benches),*]) {
481 Some(array) => array,
482 _ => panic!("external 'consts' cannot contain more than 20 values"),
483 };
484
485 &__DIVAN_GENERIC_BENCHES
486 }
487 });
488
489 make_generic_group(quote! {
490 const __DIVAN_CONST_COUNT: usize = __DIVAN_CONSTS.len();
491 const __DIVAN_CONSTS: &[#const_type] = &#generic_consts;
492
493 &[#({ #generic_benches }),*]
494 })
495 }
496 };
497
498 // Append our generated code to the existing token stream.
499 let mut result = item;
500 result.extend(TokenStream::from(generated_items));
501 result
502}
503
504#[proc_macro_attribute]
505pub fn bench_group(options: TokenStream, item: TokenStream) -> TokenStream {
506 let attr = Macro::BenchGroup;
507 let attr_name = attr.name();
508
509 let options = match AttrOptions::parse(options, attr) {
510 Ok(options) => options,
511 Err(compile_error) => return compile_error,
512 };
513
514 // Items needed by generated code.
515 let AttrOptions { private_mod, .. } = &options;
516
517 let option_none = tokens::option_none();
518
519 // TODO: Make module parsing cheaper by parsing only the necessary parts.
520 let mod_item = item.clone();
521 let mod_item = syn::parse_macro_input!(mod_item as syn::ItemMod);
522
523 let mod_ident = &mod_item.ident;
524 let mod_name = mod_ident.to_string();
525 let mod_name_pretty = mod_name.strip_prefix("r#").unwrap_or(&mod_name);
526
527 // Find any `#[ignore]` attribute so that we can use its span to help
528 // compiler diagnostics.
529 //
530 // TODO: Fix `unused_attributes` warning when using `#[ignore]` on a module.
531 let ignore_attr_ident =
532 mod_item.attrs.iter().map(|attr| attr.meta.path()).find(|path| path.is_ident("ignore"));
533
534 // Prefixed with "__" to prevent IDEs from recommending using this symbol.
535 //
536 // By having the static be local, we cause a compile error if this attribute
537 // is used multiple times on the same function.
538 let static_ident = syn::Ident::new(
539 &format!("__DIVAN_GROUP_{}", mod_name_pretty.to_uppercase()),
540 mod_ident.span(),
541 );
542
543 let meta = entry_meta_expr(&mod_name, &options, ignore_attr_ident);
544
545 let pre_main_attrs = pre_main_attrs();
546 let unsupported_error = unsupported_error(attr_name);
547
548 let generated_items = quote! {
549 #unsupported_error
550
551 // Push this static into `GROUP_ENTRIES` before `main` is called.
552 static #static_ident: #private_mod::EntryList<#private_mod::GroupEntry> = {
553 {
554 // Add `push` to the initializer section.
555 #pre_main_attrs
556 static PUSH: extern "C" fn() = push;
557
558 extern "C" fn push() {
559 #private_mod::GROUP_ENTRIES.push(&#static_ident);
560 }
561 }
562
563 #private_mod::EntryList::new({
564 static #static_ident: #private_mod::GroupEntry = #private_mod::GroupEntry {
565 meta: #meta,
566 generic_benches: #option_none,
567 };
568
569 &#static_ident
570 })
571 };
572 };
573
574 // Append our generated code to the existing token stream.
575 let mut result = item;
576 result.extend(TokenStream::from(generated_items));
577 result
578}
579
580/// Constructs an `EntryMeta` expression.
581fn entry_meta_expr(
582 raw_name: &str,
583 options: &AttrOptions,
584 ignore_attr_ident: Option<&syn::Path>,
585) -> proc_macro2::TokenStream {
586 let AttrOptions { private_mod, .. } = &options;
587
588 let raw_name_pretty = raw_name.strip_prefix("r#").unwrap_or(raw_name);
589
590 let display_name: &dyn ToTokens = match &options.name_expr {
591 Some(name) => name,
592 None => &raw_name_pretty,
593 };
594
595 let bench_options = options.bench_options_fn(ignore_attr_ident);
596
597 quote! {
598 #private_mod::EntryMeta {
599 raw_name: #raw_name,
600 display_name: #display_name,
601 bench_options: #bench_options,
602 module_path: ::std::module_path!(),
603
604 // `Span` location info is nightly-only, so use macros.
605 location: #private_mod::EntryLocation {
606 file: ::std::file!(),
607 line: ::std::line!(),
608 col: ::std::column!(),
609 },
610 }
611 }
612}
613