diff --git a/README.md b/README.md index baf6ce1..e08558a 100644 --- a/README.md +++ b/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 diff --git a/src/bencher.rs b/src/bencher.rs index f47c7b2..b4bc6c9 100644 --- a/src/bencher.rs +++ b/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 { + /// # vec![] + /// // ... + /// } + /// + /// fn reset_global_data(data: &mut Vec) { + /// // ... + /// } + /// + /// // The algorithm to test + /// fn do_something_with(data: &mut [u64]) -> Vec { + /// # 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(&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(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>(&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: ::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(&mut self, mut setup: S) + where + S: FnMut(&mut WrapperRunner<'a, '_, M>), + { + unimplemented!("Unsupported at present"); + } } diff --git a/src/codspeed/bencher.rs b/src/codspeed/bencher.rs index 995fc63..d895ea6 100644 --- a/src/codspeed/bencher.rs +++ b/src/codspeed/bencher.rs @@ -120,12 +120,62 @@ impl<'a> Bencher<'a> { } } + #[inline(never)] + pub fn iter_with_setup_wrapper(&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(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>(&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(&mut self, mut setup: S) + where + S: FnMut(&mut WrapperRunner), + { + unimplemented!("Unsupported at present"); + } } diff --git a/tests/criterion_tests.rs b/tests/criterion_tests.rs index 5a4168c..ddf494b 100644 --- a/tests/criterion_tests.rs +++ b/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 = 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();