diff --git a/library/core/benches/iter.rs b/library/core/benches/iter.rs index fb6b4b7837941..5be6b1c86c99c 100644 --- a/library/core/benches/iter.rs +++ b/library/core/benches/iter.rs @@ -345,3 +345,96 @@ fn bench_partial_cmp(b: &mut Bencher) { fn bench_lt(b: &mut Bencher) { b.iter(|| (0..100000).map(black_box).lt((0..100000).map(black_box))) } + +#[bench] +fn bench_fold_fold_mut_vec(b: &mut Bencher) { + b.iter(|| { + (0..100000).map(black_box).fold(Vec::new(), |mut v, n| { + if n % 2 == 0 { + v.push(n * 3); + } + v + }) + }); +} + +#[bench] +fn bench_fold_fold_mut_vec_mut(b: &mut Bencher) { + b.iter(|| { + (0..100000).map(black_box).fold_mut(Vec::new(), |v, n| { + if n % 2 == 0 { + v.push(n * 3); + } + }) + }); +} + +#[bench] +fn bench_fold_fold_mut_hashmap(b: &mut Bencher) { + use std::collections::HashMap; + + b.iter(|| { + (0..100000).map(black_box).fold(HashMap::new(), |mut hm, n| { + *hm.entry(n % 3).or_insert(0) += 1; + hm + }) + }); +} + +#[bench] +fn bench_fold_fold_mut_hashmap_mut(b: &mut Bencher) { + use std::collections::HashMap; + + b.iter(|| { + (0..100000).map(black_box).fold_mut(HashMap::new(), |hm, n| { + *hm.entry(n % 3).or_insert(0) += 1; + }) + }); +} + +#[bench] +fn bench_fold_fold_mut_num(b: &mut Bencher) { + b.iter(|| (0..100000).map(black_box).fold(0, |sum, n| sum + n)); +} + +#[bench] +fn bench_fold_fold_mut_num_mut(b: &mut Bencher) { + b.iter(|| { + (0..100000).map(black_box).fold_mut(0, |sum, n| { + *sum += n; + }) + }); +} + +#[bench] +fn bench_fold_fold_mut_chain(b: &mut Bencher) { + b.iter(|| (0i64..1000000).chain(0..1000000).map(black_box).fold(0, |sum, n| sum + n)); +} + +#[bench] +fn bench_fold_fold_mut_chain_mut(b: &mut Bencher) { + b.iter(|| { + (0i64..1000000).chain(0..1000000).map(black_box).fold_mut(0, |sum, n| { + *sum += n; + }) + }); +} + +#[bench] +fn bench_fold_fold_mut_chain_flat_map(b: &mut Bencher) { + b.iter(|| { + (0i64..1000000) + .flat_map(|x| once(x).chain(once(x))) + .map(black_box) + .fold(0, |sum, n| sum + n) + }); +} + +#[bench] +fn bench_fold_fold_mut_chain_flat_map_mut(b: &mut Bencher) { + b.iter(|| { + (0i64..1000000).flat_map(|x| once(x).chain(once(x))).map(black_box).fold_mut(0, |sum, n| { + *sum += n; + }) + }); +} diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs index b8a09f822b6da..9ba6900603980 100644 --- a/library/core/src/iter/traits/iterator.rs +++ b/library/core/src/iter/traits/iterator.rs @@ -1990,6 +1990,44 @@ pub trait Iterator { accum } + /// An iterator method that applies a function, producing a single, final value. + /// + /// `fold_mut()` is very similar to [`fold()`] except that the closure + /// takes a `&mut` to the 'accumulator' and does not need to return a new value. + /// + /// [`fold()`]: Iterator::fold + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(iterator_fold_mut)] + /// use std::collections::HashMap; + /// + /// let word = "abracadabra"; + /// + /// // the count of each letter in a HashMap + /// let counts = word.chars().fold_mut(HashMap::new(), |map, c| { + /// *map.entry(c).or_insert(0) += 1; + /// }); + /// + /// assert_eq!(counts[&'a'], 5); + /// ``` + #[inline] + #[unstable(feature = "iterator_fold_mut", issue = "76751")] + fn fold_mut(mut self, init: B, mut f: F) -> B + where + Self: Sized, + F: FnMut(&mut B, Self::Item), + { + let mut accum = init; + while let Some(x) = self.next() { + f(&mut accum, x); + } + accum + } + /// The same as [`fold()`], but uses the first element in the /// iterator as the initial value, folding every subsequent element into it. /// If the iterator is empty, return [`None`]; otherwise, return the result diff --git a/library/core/tests/iter.rs b/library/core/tests/iter.rs index 00e3972c42f9d..d290d487938ee 100644 --- a/library/core/tests/iter.rs +++ b/library/core/tests/iter.rs @@ -3222,3 +3222,15 @@ fn test_flatten_non_fused_inner() { assert_eq!(iter.next(), Some(1)); assert_eq!(iter.next(), None); } + +#[test] +fn test_fold_mut() { + let nums = [1, 2, 3, 4, 5]; + let result = nums.iter().fold_mut(Vec::new(), |v, i| { + if i % 2 == 0 { + v.push(i * 3); + } + }); + + assert_eq!(result, vec![6, 12]); +}