feat: add Bencher::iter_with_setup_wrapper (#49)

Add `Bencher::iter_with_setup_wrapper`. This API allows setup to be
performed before each iteration of the benchmark, and the setup and
routine functions can mutably borrow values from outside their closures.

This enables us to benchmark Oxc's components in isolation e.g.
oxc-project/oxc#5193.

README.md

@@ -37,6 +37,7 @@ This is fork is updated with:
 * `clap` replaced with [`bpaf`](https://github.com/pacak/bpaf) to reduce binary size and compilation time
 * merged the `criterion-plot` crate into `criterion2`
 * remove regex filter support to reduce compilation time
+* added `Bencher::iter_with_setup_wrapper` method
 
 ## Table of Contents
 

src/bencher.rs

@@ -360,6 +360,80 @@ impl<'a, M: Measurement> Bencher<'a, M> {
         self.elapsed_time = time_start.elapsed();
     }
 
+    /// Times a routine that requires some setup which mutably borrows data from outside the setup
+    /// function.
+    ///
+    /// The setup function is passed a [`WrapperRunner`]. It should perform whatever setup is required
+    /// and then call `run` with the `routine` function. Only the execution time of the `routine`
+    /// function is measured.
+    ///
+    /// Each iteration of the benchmark is executed in series. So `setup` can mutably borrow data from
+    /// outside its closure mutably and know that it has exclusive access to that data throughout each
+    /// `setup` + `routine` iteration.
+    /// i.e. equivalent to [`BatchSize::PerIteration`].
+    ///
+    /// Value returned by `routine` is returned from `run`. If you do not wish include drop time of
+    /// a value in the measurement, return it from `routine` so it is dropped outside of the measured
+    /// section.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// use criterion::*;
+    ///
+    /// fn create_global_data() -> Vec<u64> {
+    ///     # vec![]
+    ///     // ...
+    /// }
+    ///
+    /// fn reset_global_data(data: &mut Vec<u64>) {
+    ///     // ...
+    /// }
+    ///
+    /// // The algorithm to test
+    /// fn do_something_with(data: &mut [u64]) -> Vec<u64> {
+    ///     # vec![]
+    ///     // ...
+    /// }
+    ///
+    /// fn bench(c: &mut Criterion) {
+    ///     let mut data = create_global_data();
+    ///
+    ///     c.bench_function("with_setup_wrapper", |b| {
+    ///         b.iter_with_setup_wrapper(|runner| {
+    ///             // Perform setup on each iteration. Not included in measurement.
+    ///             reset_global_data(&mut data);
+    ///
+    ///             runner.run(|| {
+    ///                 // Code in this closure is measured
+    ///                 let result = do_something_with(&mut data);
+    ///                 // Return result if do not want to include time dropping it in measure
+    ///                 result
+    ///             });
+    ///         });
+    ///     });
+    /// }
+    ///
+    /// criterion_group!(benches, bench);
+    /// criterion_main!(benches);
+    /// ```
+    ///
+    #[inline(never)]
+    pub fn iter_with_setup_wrapper<S>(&mut self, mut setup: S)
+    where
+        S: FnMut(&mut WrapperRunner<'a, '_, M>),
+    {
+        self.iterated = true;
+        let time_start = Instant::now();
+        self.value = self.measurement.zero();
+
+        for _ in 0..self.iters {
+            WrapperRunner::execute(self, &mut setup);
+        }
+
+        self.elapsed_time = time_start.elapsed();
+    }
+
     // Benchmarks must actually call one of the iter methods. This causes benchmarks to fail loudly
     // if they don't.
     pub(crate) fn assert_iterated(&mut self) {
@@ -374,6 +448,37 @@ impl<'a, M: Measurement> Bencher<'a, M> {
     }
 }
 
+/// Runner used by [`Bencher::iter_with_setup_wrapper`].
+pub struct WrapperRunner<'a, 'b, M: Measurement> {
+    bencher: &'b mut Bencher<'a, M>,
+    has_run: bool,
+}
+
+impl<'a, 'b, M: Measurement> WrapperRunner<'a, 'b, M> {
+    fn execute<S>(bencher: &'b mut Bencher<'a, M>, setup: &mut S)
+    where
+        S: FnMut(&mut Self),
+    {
+        let mut runner = Self { bencher, has_run: false };
+        setup(&mut runner);
+        assert!(runner.has_run, "setup function must call `WrapperRunner::run`");
+    }
+
+    pub fn run<O, R: FnOnce() -> O>(&mut self, routine: R) -> O {
+        assert!(!self.has_run, "setup function must call `WrapperRunner::run` only once");
+        self.has_run = true;
+
+        let bencher = &mut self.bencher;
+
+        let start: <M as Measurement>::Intermediate = bencher.measurement.start();
+        let output = routine();
+        let end = bencher.measurement.end(start);
+        bencher.value = bencher.measurement.add(&bencher.value, &end);
+
+        black_box(output)
+    }
+}
+
 /// Async/await variant of the Bencher struct.
 #[cfg(feature = "async")]
 pub struct AsyncBencher<'a, 'b, A: AsyncExecutor, M: Measurement = WallTime> {
@@ -802,4 +907,11 @@ impl<'a, 'b, A: AsyncExecutor, M: Measurement> AsyncBencher<'a, 'b, A, M> {
             b.elapsed_time = time_start.elapsed();
         });
     }
+
+    pub fn iter_with_setup_wrapper<S>(&mut self, mut setup: S)
+    where
+        S: FnMut(&mut WrapperRunner<'a, '_, M>),
+    {
+        unimplemented!("Unsupported at present");
+    }
 }

src/codspeed/bencher.rs

@@ -120,12 +120,62 @@ impl<'a> Bencher<'a> {
         }
     }
 
+    #[inline(never)]
+    pub fn iter_with_setup_wrapper<S>(&mut self, mut setup: S)
+    where
+        S: FnMut(&mut WrapperRunner),
+    {
+        let mut codspeed = self.codspeed.borrow_mut();
+        let name = self.uri.as_str();
+
+        for i in 0..codspeed::codspeed::WARMUP_RUNS + 1 {
+            if i < codspeed::codspeed::WARMUP_RUNS {
+                WrapperRunner::execute(None, name, &mut setup)
+            } else {
+                WrapperRunner::execute(Some(&mut *codspeed), name, &mut setup)
+            }
+        }
+    }
+
     #[cfg(feature = "async")]
     pub fn to_async<'b, A: AsyncExecutor>(&'b mut self, runner: A) -> AsyncBencher<'a, 'b, A> {
         AsyncBencher { b: self, runner }
     }
 }
 
+/// Runner used by [`Bencher::iter_with_setup_wrapper`].
+pub struct WrapperRunner<'c> {
+    codspeed: Option<&'c mut CodSpeed>,
+    name: &'c str,
+    has_run: bool,
+}
+
+impl<'c> WrapperRunner<'c> {
+    fn execute<S>(codspeed: Option<&'c mut CodSpeed>, name: &'c str, setup: &mut S)
+    where
+        S: FnMut(&mut Self),
+    {
+        let mut runner = Self { codspeed, name, has_run: false };
+        setup(&mut runner);
+        assert!(runner.has_run, "setup function must call `WrapperRunner::run`");
+    }
+
+    pub fn run<O, R: FnOnce() -> O>(&mut self, routine: R) -> O {
+        assert!(!self.has_run, "setup function must call `WrapperRunner::run` only once");
+        self.has_run = true;
+
+        let output = if let Some(codspeed) = self.codspeed.as_mut() {
+            codspeed.start_benchmark(self.name);
+            let output = black_box(routine());
+            codspeed.end_benchmark();
+            output
+        } else {
+            routine()
+        };
+        black_box(output)
+    }
+}
+
 #[cfg(feature = "async")]
 pub struct AsyncBencher<'a, 'b, A: AsyncExecutor> {
     b: &'b mut Bencher<'a>,
@@ -256,4 +306,11 @@ impl<'a, 'b, A: AsyncExecutor> AsyncBencher<'a, 'b, A> {
             }
         });
     }
+
+    pub fn iter_with_setup_wrapper<S>(&mut self, mut setup: S)
+    where
+        S: FnMut(&mut WrapperRunner),
+    {
+        unimplemented!("Unsupported at present");
+    }
 }

tests/criterion_tests.rs

@@ -288,6 +288,19 @@ fn test_timing_loops() {
     group.bench_function("iter_batched_ref_10_iterations", |b| {
         b.iter_batched_ref(|| vec![10], |v| v[0], BatchSize::NumIterations(10))
     });
+    let mut global_data: Vec<u8> = vec![];
+    group.bench_function("iter_with_setup_wrapper", |b| {
+        b.iter_with_setup_wrapper(|runner| {
+            global_data.clear();
+            global_data.extend(&[1, 2, 3, 4, 5]);
+            let len = runner.run(|| {
+                global_data.push(6);
+                global_data.len()
+            });
+            assert_eq!(len, 6);
+            assert_eq!(global_data.len(), 6);
+        })
+    });
 }
 
 #[test]
@@ -297,6 +310,27 @@ fn test_bench_with_no_iteration_panics() {
     short_benchmark(&dir).bench_function("no_iter", |_b| {});
 }
 
+#[test]
+#[should_panic(expected = "setup function must call `WrapperRunner::run`")]
+fn test_setup_wrapper_with_no_runner_call_panics() {
+    let dir = temp_dir();
+    short_benchmark(&dir).bench_function("no_run_call", |b| {
+        b.iter_with_setup_wrapper(|_runner| {});
+    });
+}
+
+#[test]
+#[should_panic(expected = "setup function must call `WrapperRunner::run` only once")]
+fn test_setup_wrapper_with_multiple_runner_calls_panics() {
+    let dir = temp_dir();
+    short_benchmark(&dir).bench_function("no_run_call", |b| {
+        b.iter_with_setup_wrapper(|runner| {
+            runner.run(|| {});
+            runner.run(|| {});
+        });
+    });
+}
+
 #[test]
 fn test_benchmark_group_with_input() {
     let dir = temp_dir();