feat: add Bencher::iter_with_setup_wrapper

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();