Skip to main content

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}