Auto merge of #119207 - scottmcm:manual-index-range-fold, r=<try>

Update `slice::Iter::rfold` to match `slice::Iter::fold`

Adds a new codegen test for `rfold`, like the one from #106343, and makes a similar fix, updating `rfold` to work via indices too.

Author: bors

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<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 {

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};

library/core/src/slice/iter/macros.rs

@@ -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")]
@@ -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`,
@@ -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")]

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))
+}