Specialize Iterator for &mut I where I: Sized

library/alloc/src/vec/drain.rs

@@ -1,7 +1,7 @@
 use crate::alloc::{Allocator, Global};
 use core::fmt;
 use core::iter::{FusedIterator, TrustedLen};
-use core::mem::{self};
+use core::mem::{self, MaybeUninit};
 use core::ptr::{self, NonNull};
 use core::slice::{self};
 
@@ -102,16 +102,11 @@ impl<T, A: Allocator> DoubleEndedIterator for Drain<'_, T, A> {
 #[stable(feature = "drain", since = "1.6.0")]
 impl<T, A: Allocator> Drop for Drain<'_, T, A> {
     fn drop(&mut self) {
-        /// Continues dropping the remaining elements in the `Drain`, then moves back the
-        /// un-`Drain`ed elements to restore the original `Vec`.
+        /// Moves back the un-`Drain`ed elements to restore the original `Vec`.
         struct DropGuard<'r, 'a, T, A: Allocator>(&'r mut Drain<'a, T, A>);
 
         impl<'r, 'a, T, A: Allocator> Drop for DropGuard<'r, 'a, T, A> {
             fn drop(&mut self) {
-                // Continue the same loop we have below. If the loop already finished, this does
-                // nothing.
-                self.0.for_each(drop);
-
                 if self.0.tail_len > 0 {
                     unsafe {
                         let source_vec = self.0.vec.as_mut();
@@ -129,15 +124,43 @@ impl<T, A: Allocator> Drop for Drain<'_, T, A> {
             }
         }
 
-        // exhaust self first
-        while let Some(item) = self.next() {
-            let guard = DropGuard(self);
-            drop(item);
-            mem::forget(guard);
+        let iter = mem::replace(&mut self.iter, (&mut []).iter());
+        let drop_len = iter.len();
+        let drop_ptr = iter.as_slice().as_ptr();
+
+        // forget iter so there's no aliasing reference
+        drop(iter);
+
+        let mut vec = self.vec;
+
+        if mem::size_of::<T>() == 0 {
+            // ZSTs have no identity, so we don't need to move them around, we only need to drop the correct amount.
+            // this can be achieved by manipulating the Vec length instead of moving values out from `iter`.
+            unsafe {
+                let vec = vec.as_mut();
+                let old_len = vec.len();
+                vec.set_len(old_len + drop_len + self.tail_len);
+                vec.truncate(old_len + self.tail_len);
+            }
+
+            return;
+        }
+
+        // ensure elements are moved back into their appropriate places, even when drop_in_place panics
+        let _guard = DropGuard(self);
+
+        if drop_len == 0 {
+            return;
         }
 
-        // Drop a `DropGuard` to move back the non-drained tail of `self`.
-        DropGuard(self);
+        unsafe {
+            let vec = vec.as_mut();
+            let spare_capacity = vec.spare_capacity_mut();
+            let drop_offset = drop_ptr.offset_from(spare_capacity.as_ptr() as *const _) as usize;
+            let drop_range = drop_offset..(drop_offset + drop_len);
+            let to_drop = &mut spare_capacity[drop_range];
+            ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(to_drop));
+        }
     }
 }
 

library/core/src/iter/traits/double_ended.rs

@@ -214,17 +214,13 @@ pub trait DoubleEndedIterator: Iterator {
     /// ```
     #[inline]
     #[stable(feature = "iterator_try_fold", since = "1.27.0")]
-    fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
+    fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R
     where
         Self: Sized,
         F: FnMut(B, Self::Item) -> R,
         R: Try<Ok = B>,
     {
-        let mut accum = init;
-        while let Some(x) = self.next_back() {
-            accum = f(accum, x)?;
-        }
-        try { accum }
+        default_try_rfold(self, init, f)
     }
 
     /// An iterator method that reduces the iterator's elements to a single,
@@ -278,16 +274,12 @@ pub trait DoubleEndedIterator: Iterator {
     /// ```
     #[inline]
     #[stable(feature = "iter_rfold", since = "1.27.0")]
-    fn rfold<B, F>(mut self, init: B, mut f: F) -> B
+    fn rfold<B, F>(self, init: B, f: F) -> B
     where
         Self: Sized,
         F: FnMut(B, Self::Item) -> B,
     {
-        let mut accum = init;
-        while let Some(x) = self.next_back() {
-            accum = f(accum, x);
-        }
-        accum
+        default_rfold(self, init, f)
     }
 
     /// Searches for an element of an iterator from the back that satisfies a predicate.
@@ -349,6 +341,33 @@ pub trait DoubleEndedIterator: Iterator {
     }
 }
 
+#[inline]
+fn default_try_rfold<I, B, F, R>(iter: &mut I, init: B, mut f: F) -> R
+where
+    I: DoubleEndedIterator + ?Sized,
+    F: FnMut(B, I::Item) -> R,
+    R: Try<Ok = B>,
+{
+    let mut accum = init;
+    while let Some(x) = iter.next_back() {
+        accum = f(accum, x)?;
+    }
+    try { accum }
+}
+
+#[inline]
+fn default_rfold<I, B, F>(mut iter: I, init: B, mut f: F) -> B
+where
+    I: DoubleEndedIterator,
+    F: FnMut(B, I::Item) -> B,
+{
+    let mut accum = init;
+    while let Some(x) = iter.next_back() {
+        accum = f(accum, x);
+    }
+    accum
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'a, I: DoubleEndedIterator + ?Sized> DoubleEndedIterator for &'a mut I {
     fn next_back(&mut self) -> Option<I::Item> {
@@ -360,4 +379,69 @@ impl<'a, I: DoubleEndedIterator + ?Sized> DoubleEndedIterator for &'a mut I {
     fn nth_back(&mut self, n: usize) -> Option<I::Item> {
         (**self).nth_back(n)
     }
+    #[inline]
+    fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Ok = B>,
+    {
+        SpecSizedDoubleEndedIterator::try_rfold(self, init, f)
+    }
+    #[inline]
+    fn rfold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        SpecSizedDoubleEndedIterator::rfold(self, init, f)
+    }
+}
+
+trait SpecSizedDoubleEndedIterator: DoubleEndedIterator {
+    fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Ok = B>;
+    fn rfold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B;
+}
+
+impl<I: DoubleEndedIterator + ?Sized> SpecSizedDoubleEndedIterator for &mut I {
+    #[inline]
+    default fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Ok = B>,
+    {
+        default_try_rfold(self, init, f)
+    }
+    #[inline]
+    default fn rfold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        default_rfold(self, init, f)
+    }
+}
+
+impl<I: DoubleEndedIterator + Sized> SpecSizedDoubleEndedIterator for &mut I {
+    #[inline]
+    fn try_rfold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Ok = B>,
+    {
+        (**self).try_rfold(init, f)
+    }
+    #[inline]
+    fn rfold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        #[inline]
+        fn ok<T, U>(mut f: impl FnMut(T, U) -> T) -> impl FnMut(T, U) -> Result<T, !> {
+            move |acc, x| Ok(f(acc, x))
+        }
+        self.try_rfold(init, ok(f)).unwrap()
+    }
 }

library/core/src/iter/traits/iterator.rs

@@ -1995,17 +1995,13 @@ pub trait Iterator {
     /// ```
     #[inline]
     #[stable(feature = "iterator_try_fold", since = "1.27.0")]
-    fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
+    fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R
     where
         Self: Sized,
         F: FnMut(B, Self::Item) -> R,
         R: Try<Ok = B>,
     {
-        let mut accum = init;
-        while let Some(x) = self.next() {
-            accum = f(accum, x)?;
-        }
-        try { accum }
+        default_try_fold(self, init, f)
     }
 
     /// An iterator method that applies a fallible function to each item in the
@@ -2136,16 +2132,12 @@ pub trait Iterator {
     #[doc(alias = "inject")]
     #[inline]
     #[stable(feature = "rust1", since = "1.0.0")]
-    fn fold<B, F>(mut self, init: B, mut f: F) -> B
+    fn fold<B, F>(self, init: B, f: F) -> B
     where
         Self: Sized,
         F: FnMut(B, Self::Item) -> B,
     {
-        let mut accum = init;
-        while let Some(x) = self.next() {
-            accum = f(accum, x);
-        }
-        accum
+        default_fold(self, init, f)
     }
 
     /// Reduces the elements to a single one, by repeatedly applying a reducing
@@ -3415,6 +3407,33 @@ pub trait Iterator {
     }
 }
 
+#[inline]
+fn default_try_fold<I, B, F, R>(iter: &mut I, init: B, mut f: F) -> R
+where
+    I: Iterator + ?Sized,
+    F: FnMut(B, I::Item) -> R,
+    R: Try<Ok = B>,
+{
+    let mut accum = init;
+    while let Some(x) = iter.next() {
+        accum = f(accum, x)?;
+    }
+    try { accum }
+}
+
+#[inline]
+fn default_fold<I, B, F>(mut iter: I, init: B, mut f: F) -> B
+where
+    I: Iterator,
+    F: FnMut(B, I::Item) -> B,
+{
+    let mut accum = init;
+    while let Some(x) = iter.next() {
+        accum = f(accum, x);
+    }
+    accum
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<I: Iterator + ?Sized> Iterator for &mut I {
     type Item = I::Item;
@@ -3430,4 +3449,69 @@ impl<I: Iterator + ?Sized> Iterator for &mut I {
     fn nth(&mut self, n: usize) -> Option<Self::Item> {
         (**self).nth(n)
     }
+    #[inline]
+    fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Ok = B>,
+    {
+        SpecSizedIterator::try_fold(self, init, f)
+    }
+    #[inline]
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        SpecSizedIterator::fold(self, init, f)
+    }
+}
+
+trait SpecSizedIterator: Iterator {
+    fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Ok = B>;
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B;
+}
+
+impl<I: Iterator + ?Sized> SpecSizedIterator for &mut I {
+    #[inline]
+    default fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Ok = B>,
+    {
+        default_try_fold(self, init, f)
+    }
+    #[inline]
+    default fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        default_fold(self, init, f)
+    }
+}
+
+impl<I: Iterator + Sized> SpecSizedIterator for &mut I {
+    #[inline]
+    fn try_fold<B, F, R>(&mut self, init: B, f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Ok = B>,
+    {
+        (**self).try_fold(init, f)
+    }
+    #[inline]
+    fn fold<B, F>(self, init: B, f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        #[inline]
+        fn ok<T, U>(mut f: impl FnMut(T, U) -> T) -> impl FnMut(T, U) -> Result<T, !> {
+            move |acc, x| Ok(f(acc, x))
+        }
+        self.try_fold(init, ok(f)).unwrap()
+    }
 }

src/test/ui/issues/issue-3044.stderr

@@ -11,7 +11,7 @@ LL | |     });
 note: associated function defined here
   --> $SRC_DIR/core/src/iter/traits/iterator.rs:LL:COL
    |
-LL |     fn fold<B, F>(mut self, init: B, mut f: F) -> B
+LL |     fn fold<B, F>(self, init: B, f: F) -> B
    |        ^^^^
 
 error: aborting due to previous error