diff --git a/Cargo.toml b/Cargo.toml index 3ec0f56f8..54343bf4c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -66,3 +66,7 @@ harness = false [[bench]] name = "bench1" harness = false + +[[bench]] +name = "fold_mut" +harness = false diff --git a/benches/fold_mut.rs b/benches/fold_mut.rs new file mode 100644 index 000000000..5e5cc9fd8 --- /dev/null +++ b/benches/fold_mut.rs @@ -0,0 +1,109 @@ +use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use itertools::Itertools; + +fn bench_sum(c: &mut Criterion) { + let mut group = c.benchmark_group("fold sum accumulator"); + + group.bench_function("fold", |b| { + b.iter(|| { + (0i64..1_000_000) + .map(black_box) + .fold(0, |sum, n| sum + n) + }) + }); + + group.bench_function("fold_mut", |b| { + b.iter(|| { + (0i64..1_000_000).map(black_box).fold_mut(0, |sum, n| { + *sum += n; + }) + }) + }); + + group.finish(); +} + +fn bench_vec(c: &mut Criterion) { + let mut group = c.benchmark_group("fold vec accumulator"); + + group.bench_function("fold", |b| { + b.iter(|| { + (0i64..1_000_000) + .map(black_box) + .fold(Vec::new(), |mut v, n| { + v.push(n); + v + }) + }) + }); + + group.bench_function("fold_mut", |b| { + b.iter(|| { + (0i64..1_000_000) + .map(black_box) + .fold_mut(Vec::new(), |v, n| { + v.push(n); + }) + }) + }); + + group.finish(); +} + +fn bench_num_with_chain(c: &mut Criterion) { + let mut group = c.benchmark_group("fold chained iterator with num accumulator"); + + group.bench_function("fold", |b| { + b.iter(|| { + (0i64..1_000_000) + .chain(0i64..1_000_000) + .map(black_box) + .fold(0, |sum, n| sum + n) + }) + }); + + group.bench_function("fold_mut", |b| { + b.iter(|| { + (0i64..1_000_000) + .chain(0i64..1_000_000) + .map(black_box) + .fold_mut(0, |sum, n| { + *sum += n; + }) + }) + }); + + group.finish(); +} + +fn bench_vec_with_chain(c: &mut Criterion) { + let mut group = c.benchmark_group("fold chained iterator with vec accumulator"); + + group.bench_function("fold", |b| { + b.iter(|| { + (0i64..1_000_000) + .chain(0i64..1_000_000) + .map(black_box) + .fold(Vec::new(), |mut v, n| { + v.push(n); + v + }) + }) + }); + + group.bench_function("fold_mut", |b| { + b.iter(|| { + (0i64..1_000_000) + .chain(0i64..1_000_000) + .map(black_box) + .fold_mut(Vec::new(), |v, n| { + v.push(n); + }) + }) + }); + + group.finish(); +} + +criterion_group!(benches, bench_sum, bench_vec, bench_num_with_chain, bench_vec_with_chain); +criterion_main!(benches); diff --git a/src/lib.rs b/src/lib.rs index 4c10178fd..98a3c0e4d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2831,6 +2831,46 @@ pub trait Itertools : Iterator { self.for_each(|item| *counts.entry(item).or_default() += 1); counts } + + /// An alternative to [`fold()`] - `fold_mut()` also applies a function + /// producing a single value. The main difference is that the closure + /// passed to `fold_mut()` accepts a `&mut` to the accumulator instead + /// of consuming the accumulator. This can simplify some closures + /// that might otherwise be forced to return the accumulator awkwardly: + /// ``` + /// let evens = [1, 2, 3, 4, 5, 6].iter().fold(Vec::new(), |mut evens, num| { + /// if num % 2 == 0 { + /// evens.push(num); + /// } + /// evens // potentially awkward return + /// }); + /// ``` + /// + /// `fold_mut()` may also be more performant in situations where the + /// accumulator is "large" as passing it by `&mut` can be cheaper than moving it. + /// + /// # Examples + /// ``` + /// # use itertools::Itertools; + /// let evens = [1, 2, 3, 4, 5, 6].iter().fold_mut(Vec::new(), |evens, &num| { + /// if num % 2 == 0 { + /// evens.push(num); + /// } + /// }); + /// + /// assert_eq!(evens, [2, 4, 6]); + /// ``` + /// [`fold()`]: https://doc.rust-lang.org/std/iter/trait.Iterator.html#method.fold + #[cfg(feature = "use_std")] + fn fold_mut(self, init: B, mut f: F) -> B + where + Self: Sized, + F: FnMut(&mut B, Self::Item), + { + let mut accum = init; + self.for_each(|item| f(&mut accum, item)); + accum + } } impl Itertools for T where T: Iterator { } diff --git a/tests/quick.rs b/tests/quick.rs index ff05a478c..3a5b26d4a 100644 --- a/tests/quick.rs +++ b/tests/quick.rs @@ -1237,3 +1237,24 @@ quickcheck! { TestResult::from_bool(itertools::equal(x, y)) } } + +quickcheck! { + fn test_fold_mut(a: Vec) -> TestResult { + let with_fold_mut = + a.iter().fold_mut(Vec::new(), |v, &n| { + if n % 2 == 0 { + v.push(n); + } + }); + + let with_fold = + a.iter().fold(Vec::new(), |mut v, &n| { + if n % 2 == 0 { + v.push(n); + } + v + }); + + TestResult::from_bool(itertools::equal(with_fold_mut, with_fold)) + } +}