minibench/lib.rs
1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8//! Run function repetitively and prints out aggregated result.
9//!
10//! Differences from other benchmark library:
11//! - Do not enforce measuring "wall time". It could be changed to "CPU time", "memory usage", etc.
12//! - Do not run only the benchmark part repetitively. For example, a benchmark needs some complex
13//! setup that cannot be reused across runs. That setup cost needs to be excluded from benchmark
14//! result cleanly.
15//! - Minimalism. Without fancy features.
16
17use std::env::args;
18use std::sync::OnceLock;
19
20pub mod measure;
21pub use measure::Measure;
22
23/// Measure the best wall clock time.
24pub fn elapsed(func: impl FnMut()) -> Result<self::measure::WallClock, String> {
25 self::measure::WallClock::measure(func)
26}
27
28/// Run a function repeatably. Print the measurement result.
29///
30/// The actual measurement is dependent on the return value of the function.
31/// For example,
32/// - [`WallClock::measure`] (or [`elapsed`]) measures the wall clock time,
33/// and the function being measured does not need to provide an output.
34/// - [`Bytes::measure`] might expect the function to return a [`usize`]
35/// in bytes. The function will be only run once.
36///
37/// If `std::env::args` (excluding the first item and flags) is not empty, and none of them is a
38/// substring of `name`, skip running and return directly.
39///
40/// Example:
41///
42/// ```
43/// use minibench::*;
44/// bench("example", || {
45/// // prepare
46/// elapsed(|| {
47/// // measure
48/// })
49/// })
50/// ```
51pub fn bench<T: Measure, F: FnMut() -> Result<T, String>>(name: impl ToString, mut func: F) {
52 let name = name.to_string();
53 if bench_enabled(&name) {
54 let mut try_func = || -> Result<T, String> {
55 let mut measured = func()?;
56 while measured.need_more() {
57 measured = measured.merge(func()?);
58 }
59 Ok(measured)
60 };
61 let text = match try_func() {
62 Ok(measured) => measured.to_string(),
63 Err(text) => text,
64 };
65 println!("{:50}{}", name, text);
66 }
67}
68
69/// Test if a benchmark is enabled.
70pub fn bench_enabled(name: &str) -> bool {
71 static ARGS: OnceLock<Vec<String>> = OnceLock::new();
72 let args = ARGS.get_or_init(|| {
73 // The first arg is the program name. Skip it and flag-like arguments (ex. --bench).
74 args().skip(1).filter(|a| !a.starts_with('-')).collect()
75 });
76 args.is_empty() || args.iter().any(|a| name.contains(a))
77}