Skip to content

Commit 884205a

Browse files
committed
Fix array::IntoIter::fold to use the optimized Range::fold
It was using `Iterator::by_ref` in the implementation, which ended up pessimizing it enough that, for example, it didn't vectorize when we tried it in the <https://rust-lang.zulipchat.com/#narrow/stream/257879-project-portable-simd/topic/Reducing.20sum.20into.20wider.20types> conversation. Demonstration that the codegen test doesn't pass on the current nightly: <https://rust.godbolt.org/z/Taxev5eMn>
1 parent 8ec1d8a commit 884205a

File tree

3 files changed

+76
-1
lines changed

3 files changed

+76
-1
lines changed

core/src/array/iter.rs

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -266,7 +266,7 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
266266
Fold: FnMut(Acc, Self::Item) -> Acc,
267267
{
268268
let data = &mut self.data;
269-
self.alive.by_ref().fold(init, |acc, idx| {
269+
iter::ByRefSized(&mut self.alive).fold(init, |acc, idx| {
270270
// SAFETY: idx is obtained by folding over the `alive` range, which implies the
271271
// value is currently considered alive but as the range is being consumed each value
272272
// we read here will only be read once and then considered dead.
@@ -323,6 +323,20 @@ impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
323323
})
324324
}
325325

326+
#[inline]
327+
fn rfold<Acc, Fold>(mut self, init: Acc, mut rfold: Fold) -> Acc
328+
where
329+
Fold: FnMut(Acc, Self::Item) -> Acc,
330+
{
331+
let data = &mut self.data;
332+
iter::ByRefSized(&mut self.alive).rfold(init, |acc, idx| {
333+
// SAFETY: idx is obtained by folding over the `alive` range, which implies the
334+
// value is currently considered alive but as the range is being consumed each value
335+
// we read here will only be read once and then considered dead.
336+
rfold(acc, unsafe { data.get_unchecked(idx).assume_init_read() })
337+
})
338+
}
339+
326340
fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
327341
let len = self.len();
328342

core/src/iter/adapters/by_ref_sized.rs

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,3 +40,32 @@ impl<I: Iterator> Iterator for ByRefSized<'_, I> {
4040
self.0.try_fold(init, f)
4141
}
4242
}
43+
44+
impl<I: DoubleEndedIterator> DoubleEndedIterator for ByRefSized<'_, I> {
45+
fn next_back(&mut self) -> Option<Self::Item> {
46+
self.0.next_back()
47+
}
48+
49+
fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
50+
self.0.advance_back_by(n)
51+
}
52+
53+
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
54+
self.0.nth_back(n)
55+
}
56+
57+
fn rfold<B, F>(self, init: B, f: F) -> B
58+
where
59+
F: FnMut(B, Self::Item) -> B,
60+
{
61+
self.0.rfold(init, f)
62+
}
63+
64+
fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R
65+
where
66+
F: FnMut(B, Self::Item) -> R,
67+
R: Try<Output = B>,
68+
{
69+
self.0.try_rfold(init, f)
70+
}
71+
}

core/tests/array.rs

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -668,3 +668,35 @@ fn array_mixed_equality_nans() {
668668
assert!(!(mut3 == array3));
669669
assert!(mut3 != array3);
670670
}
671+
672+
#[test]
673+
fn array_into_iter_fold() {
674+
// Strings to help MIRI catch if we double-free or something
675+
let a = ["Aa".to_string(), "Bb".to_string(), "Cc".to_string()];
676+
let mut s = "s".to_string();
677+
a.into_iter().for_each(|b| s += &b);
678+
assert_eq!(s, "sAaBbCc");
679+
680+
let a = [1, 2, 3, 4, 5, 6];
681+
let mut it = a.into_iter();
682+
it.advance_by(1).unwrap();
683+
it.advance_back_by(2).unwrap();
684+
let s = it.fold(10, |a, b| 10 * a + b);
685+
assert_eq!(s, 10234);
686+
}
687+
688+
#[test]
689+
fn array_into_iter_rfold() {
690+
// Strings to help MIRI catch if we double-free or something
691+
let a = ["Aa".to_string(), "Bb".to_string(), "Cc".to_string()];
692+
let mut s = "s".to_string();
693+
a.into_iter().rev().for_each(|b| s += &b);
694+
assert_eq!(s, "sCcBbAa");
695+
696+
let a = [1, 2, 3, 4, 5, 6];
697+
let mut it = a.into_iter();
698+
it.advance_by(1).unwrap();
699+
it.advance_back_by(2).unwrap();
700+
let s = it.rfold(10, |a, b| 10 * a + b);
701+
assert_eq!(s, 10432);
702+
}

0 commit comments

Comments
 (0)