From d75741f05e1a7507283a83e6be44835cb2b08cb8 Mon Sep 17 00:00:00 2001 From: Scott McMurray Date: Thu, 21 Dec 2023 21:54:20 -0800 Subject: [PATCH] Update slice::Iter::rfold to match slice::Iter::fold Like in 106343, including an equivalent codegen test. --- library/core/src/ops/index_range.rs | 26 +++++++++++ library/core/src/slice/iter.rs | 1 + library/core/src/slice/iter/macros.rs | 63 ++++++++++++--------------- tests/codegen/slice-iter-fold.rs | 22 +++++++++- 4 files changed, 77 insertions(+), 35 deletions(-) diff --git a/library/core/src/ops/index_range.rs b/library/core/src/ops/index_range.rs index 743799c4b3edf..652f910c75a03 100644 --- a/library/core/src/ops/index_range.rs +++ b/library/core/src/ops/index_range.rs @@ -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(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 { diff --git a/library/core/src/slice/iter.rs b/library/core/src/slice/iter.rs index fc54ea2377096..5199f6a8d56eb 100644 --- a/library/core/src/slice/iter.rs +++ b/library/core/src/slice/iter.rs @@ -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}; diff --git a/library/core/src/slice/iter/macros.rs b/library/core/src/slice/iter/macros.rs index 95bcd123b824f..4a4091d995658 100644 --- a/library/core/src/slice/iter/macros.rs +++ b/library/core/src/slice/iter/macros.rs @@ -133,6 +133,18 @@ macro_rules! iterator { }, ) } + + #[inline] + fn via_indices(self) -> impl DoubleEndedIterator { + 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")] @@ -209,50 +221,24 @@ macro_rules! iterator { } #[inline] - fn fold(self, init: B, mut f: F) -> B + fn fold(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(mut self, mut f: F) + fn for_each(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`, @@ -427,6 +413,15 @@ macro_rules! iterator { unsafe { self.pre_dec_end(advance) }; NonZeroUsize::new(n - advance).map_or(Ok(()), Err) } + + #[inline] + fn rfold(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")] diff --git a/tests/codegen/slice-iter-fold.rs b/tests/codegen/slice-iter-fold.rs index a55425cb6bbc6..d4967fc0aa612 100644 --- a/tests/codegen/slice-iter-fold.rs +++ b/tests/codegen/slice-iter-fold.rs @@ -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 @@ -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)) +}