Skip to content

Update slice::Iter::rfold to match slice::Iter::fold #119207

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions library/core/src/ops/index_range.rs
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,32 @@ impl Iterator for IndexRange {
let taken = self.take_prefix(n);
NonZeroUsize::new(n - taken.len()).map_or(Ok(()), Err)
}

#[inline]
fn fold<B, F>(self, mut init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
// Manually canonicalize, as in #106343.
// It's not clear why it's needed -- maybe avoiding the `Option`
// checks simplifies something? -- but it does seem to help, as without
// this override the `slice_iter_fold` codegen test no longer passes.
let Self { mut start, end } = self;
if start < end {
loop {
init = f(init, start);
// SAFETY: The if above ensures this is in-range for the first
// iteration, then we break below on getting to end which will
// ensure it never tries to go above usize::MAX.
start = unsafe { unchecked_add(start, 1) };
if start == end {
break;
}
}
}
init
}
}

impl DoubleEndedIterator for IndexRange {
Expand Down
1 change: 1 addition & 0 deletions library/core/src/slice/iter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ use crate::iter::{
use crate::marker::PhantomData;
use crate::mem::{self, SizedTypeProperties};
use crate::num::NonZeroUsize;
use crate::ops::IndexRange;
use crate::ptr::{self, invalid, invalid_mut, NonNull};

use super::{from_raw_parts, from_raw_parts_mut};
Expand Down
63 changes: 29 additions & 34 deletions library/core/src/slice/iter/macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,18 @@ macro_rules! iterator {
},
)
}

#[inline]
fn via_indices(self) -> impl DoubleEndedIterator<Item = $elem> {
let len = len!(self);
let ptr = self.ptr;
IndexRange::zero_to(len)
.map(move |i| {
// SAFETY: the loop iterates `i in 0..len`, which always
// is in bounds of the slice allocation
unsafe { ptr.add(i).$into_ref() }
})
}
}

#[stable(feature = "rust1", since = "1.0.0")]
Expand Down Expand Up @@ -209,50 +221,24 @@ macro_rules! iterator {
}

#[inline]
fn fold<B, F>(self, init: B, mut f: F) -> B
fn fold<B, F>(self, init: B, f: F) -> B
where
F: FnMut(B, Self::Item) -> B,
{
// this implementation consists of the following optimizations compared to the
// default implementation:
// - do-while loop, as is llvm's preferred loop shape,
// see https://releases.llvm.org/16.0.0/docs/LoopTerminology.html#more-canonical-loops
// - bumps an index instead of a pointer since the latter case inhibits
// some optimizations, see #111603
// - avoids Option wrapping/matching
if is_empty!(self) {
return init;
}
let mut acc = init;
let mut i = 0;
let len = len!(self);
loop {
// SAFETY: the loop iterates `i in 0..len`, which always is in bounds of
// the slice allocation
acc = f(acc, unsafe { & $( $mut_ )? *self.ptr.add(i).as_ptr() });
// SAFETY: `i` can't overflow since it'll only reach usize::MAX if the
// slice had that length, in which case we'll break out of the loop
// after the increment
i = unsafe { i.unchecked_add(1) };
if i == len {
break;
}
}
acc
// Apparently it's better to loop over indexes instead of
// pointers, see #111603. Maybe it's easier on SCEV?
// This also makes sure we're never doing `Option` checks.
self.via_indices().fold(init, f)
}

// We override the default implementation, which uses `try_fold`,
// because this simple implementation generates less LLVM IR and is
// faster to compile.
#[inline]
fn for_each<F>(mut self, mut f: F)
fn for_each<F>(self, f: F)
where
Self: Sized,
F: FnMut(Self::Item),
{
while let Some(x) = self.next() {
f(x);
}
// See `fold`.
self.via_indices().for_each(f)
}

// We override the default implementation, which uses `try_fold`,
Expand Down Expand Up @@ -427,6 +413,15 @@ macro_rules! iterator {
unsafe { self.pre_dec_end(advance) };
NonZeroUsize::new(n - advance).map_or(Ok(()), Err)
}

#[inline]
fn rfold<B, F>(self, init: B, f: F) -> B
where
F: FnMut(B, Self::Item) -> B,
{
// See `fold`.
self.via_indices().rfold(init, f)
}
}

#[stable(feature = "fused", since = "1.26.0")]
Expand Down
22 changes: 21 additions & 1 deletion tests/codegen/slice-iter-fold.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,10 @@
// compile-flags: -O
#![crate_type = "lib"]

// CHECK-LABEL: @slice_fold_to_last
// CHECK-LABEL: @slice_for_each_to_last =
// CHECK-SAME: alias{{.*}}@slice_fold_to_last

// CHECK-LABEL: @slice_fold_to_last(
#[no_mangle]
pub fn slice_fold_to_last(slice: &[i32]) -> Option<&i32> {
// CHECK-NOT: loop
Expand All @@ -11,3 +14,20 @@ pub fn slice_fold_to_last(slice: &[i32]) -> Option<&i32> {
// CHECK: ret
slice.iter().fold(None, |_, i| Some(i))
}

#[no_mangle]
pub fn slice_for_each_to_last(slice: &[i32]) -> Option<&i32> {
let mut last = None;
slice.iter().for_each(|i| last = Some(i));
last
}

// CHECK-LABEL: @slice_rfold_to_first(
#[no_mangle]
pub fn slice_rfold_to_first(slice: &[i32]) -> Option<&i32> {
// CHECK-NOT: loop
// CHECK-NOT: br
// CHECK-NOT: call
// CHECK: ret
slice.iter().rfold(None, |_, i| Some(i))
}