Override RepeatN::fold too

Author: scottmcm

library/alloc/src/vec/mod.rs

@@ -2887,7 +2887,7 @@ impl<T, A: Allocator> Vec<T, A> {
                     // Since the loop executes user code which can panic we have to update
                     // the length every step to correctly drop what we've written.
                     // NB can't overflow since we would have had to alloc the address space
-                    local_len.increment_len(1);
+                    local_len.increment_len_unchecked(1);
                 });
             }
         } else {

library/alloc/src/vec/set_len_on_drop.rs

@@ -14,9 +14,12 @@ impl<'a> SetLenOnDrop<'a> {
         SetLenOnDrop { local_len: *len, len }
     }
 
+    /// # Safety
+    /// `self.current_len() + increment` must not overflow.
     #[inline]
-    pub(super) fn increment_len(&mut self, increment: usize) {
-        self.local_len += increment;
+    pub(super) unsafe fn increment_len_unchecked(&mut self, increment: usize) {
+        // SAFETY: This is our precondition
+        self.local_len = unsafe { self.local_len.unchecked_add(increment) };
     }
 
     #[inline]

library/core/src/iter/sources/repeat_n.rs

@@ -112,6 +112,7 @@ impl<A> Drop for RepeatN<A> {
 
 trait SpecRepeatN<A> {
     unsafe fn spec_next_unchecked(&mut self) -> A;
+    fn spec_fold<B, F: FnMut(B, A) -> B>(self, b: B, f: F) -> B;
 }
 
 impl<A: Clone> SpecRepeatN<A> for RepeatN<A> {
@@ -126,6 +127,18 @@ impl<A: Clone> SpecRepeatN<A> for RepeatN<A> {
             A::clone(&self.element)
         }
     }
+
+    default fn spec_fold<B, F: FnMut(B, A) -> B>(mut self, mut b: B, mut f: F) -> B {
+        let mut count = self.count;
+        if let Some(element) = self.take_element() {
+            while count > 1 {
+                count -= 1;
+                b = f(b, element.clone());
+            }
+            b = f(b, element);
+        }
+        b
+    }
 }
 
 impl<A: Copy> SpecRepeatN<A> for RepeatN<A> {
@@ -136,6 +149,17 @@ impl<A: Copy> SpecRepeatN<A> for RepeatN<A> {
         // so skip having a branch that would need to be optimized out.
         *self.element
     }
+
+    fn spec_fold<B, F: FnMut(B, A) -> B>(self, mut b: B, mut f: F) -> B {
+        let mut count = self.count;
+        let element = *self.element;
+
+        while count > 0 {
+            count -= 1;
+            b = f(b, element);
+        }
+        b
+    }
 }
 
 #[unstable(feature = "iter_repeat_n", issue = "104434")]
@@ -175,6 +199,11 @@ impl<A: Clone> Iterator for RepeatN<A> {
         }
     }
 
+    #[inline]
+    fn fold<B, F: FnMut(B, A) -> B>(self, init: B, f: F) -> B {
+        self.spec_fold(init, f)
+    }
+
     #[inline]
     fn last(mut self) -> Option<A> {
         self.take_element()