Remove fake IntoIterator trait, invoke ForLoopDesugar directly and run cleanup on drop

The on-drop cleanup will enable using TRA in try_fold(&mut self) and implementing TRA for &mut Iterator

Author: the8472

compiler/rustc_hir/src/lang_items.rs

@@ -314,8 +314,10 @@ language_item_table! {
     ControlFlowBreak,        sym::Break,               cf_break_variant,           Target::Variant,        GenericRequirement::None;
 
     IntoFutureIntoFuture,    sym::into_future,         into_future_fn,             Target::Method(MethodKind::Trait { body: false }), GenericRequirement::None;
-    IntoIterIntoIter,        sym::into_iter,           into_iter_fn,               Target::Method(MethodKind::Trait { body: false }), GenericRequirement::None;
-    IteratorNext,            sym::next,                next_fn,                    Target::Method(MethodKind::Trait { body: false}), GenericRequirement::None;
+
+    IntoIterIntoIter,        sym::into_iter,           into_iter_fn,               Target::Method(MethodKind::Inherent), GenericRequirement::None;
+    IteratorNext,            sym::next,                next_fn,                    Target::Method(MethodKind::Inherent), GenericRequirement::None;
+
 
     PinNewUnchecked,         sym::new_unchecked,       new_unchecked_fn,           Target::Method(MethodKind::Inherent), GenericRequirement::None;
 

library/alloc/src/collections/vec_deque/iter.rs

@@ -211,4 +211,13 @@ unsafe impl<T> TrustedRandomAccess for Iter<'_, T> {}
 #[unstable(feature = "trusted_random_access", issue = "none")]
 unsafe impl<T> TrustedRandomAccessNoCoerce for Iter<'_, T> {
     const MAY_HAVE_SIDE_EFFECT: bool = false;
+    const NEEDS_CLEANUP: bool = false;
+
+    fn cleanup(&mut self, num: usize, forward: bool) {
+        if forward {
+            let _ = self.advance_by(num);
+        } else {
+            let _ = self.advance_back_by(num);
+        }
+    }
 }

library/alloc/src/collections/vec_deque/iter_mut.rs

@@ -160,4 +160,13 @@ unsafe impl<T> TrustedRandomAccess for IterMut<'_, T> {}
 #[unstable(feature = "trusted_random_access", issue = "none")]
 unsafe impl<T> TrustedRandomAccessNoCoerce for IterMut<'_, T> {
     const MAY_HAVE_SIDE_EFFECT: bool = false;
+    const NEEDS_CLEANUP: bool = false;
+
+    fn cleanup(&mut self, num: usize, forward: bool) {
+        if forward {
+            let _ = self.advance_by(num);
+        } else {
+            let _ = self.advance_back_by(num);
+        }
+    }
 }

library/alloc/src/vec/in_place_collect.rs

@@ -275,6 +275,9 @@ where
                 drop_guard.dst = dst.add(1);
             }
         }
+        // FIXME: this should run the TrustedRandomAccess cleanup code.
+        //  currently not running it is ok because IntoIter only implements TRA for Copy types
+        //  but if we relax that the cleanup will be needed
         mem::forget(drop_guard);
         len
     }

library/alloc/src/vec/into_iter.rs

@@ -291,6 +291,31 @@ where
     T: NonDrop,
 {
     const MAY_HAVE_SIDE_EFFECT: bool = false;
+    const NEEDS_CLEANUP: bool = true;
+
+    fn cleanup(&mut self, num: usize, forward: bool) {
+        if forward {
+            if mem::size_of::<T>() == 0 {
+                // SAFETY: due to unchecked casts of unsigned amounts to signed offsets the wraparound
+                // effectively results in unsigned pointers representing positions 0..usize::MAX,
+                // which is valid for ZSTs.
+                self.ptr = unsafe { arith_offset(self.ptr as *const i8, num as isize) as *mut T }
+            } else {
+                // SAFETY: the caller must guarantee that `num` is in bounds
+                self.ptr = unsafe { self.ptr.add(num) };
+            }
+        } else {
+            if mem::size_of::<T>() == 0 {
+                // SAFETY: same as above
+                self.end = unsafe {
+                    arith_offset(self.end as *const i8, num.wrapping_neg() as isize) as *mut T
+                }
+            } else {
+                // SAFETY: same as above
+                self.end = unsafe { self.end.offset(num.wrapping_neg() as isize) };
+            }
+        }
+    }
 }
 
 #[cfg(not(no_global_oom_handling))]

library/core/src/iter/adapters/cloned.rs

@@ -132,6 +132,13 @@ where
     I: TrustedRandomAccessNoCoerce,
 {
     const MAY_HAVE_SIDE_EFFECT: bool = true;
+
+    const NEEDS_CLEANUP: bool = I::NEEDS_CLEANUP;
+
+    #[inline]
+    fn cleanup(&mut self, num: usize, forward: bool) {
+        self.it.cleanup(num, forward);
+    }
 }
 
 #[unstable(feature = "trusted_len", issue = "37572")]

library/core/src/iter/adapters/copied.rs

@@ -158,6 +158,13 @@ where
     I: TrustedRandomAccessNoCoerce,
 {
     const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT;
+
+    const NEEDS_CLEANUP: bool = I::NEEDS_CLEANUP;
+
+    #[inline]
+    fn cleanup(&mut self, num: usize, forward: bool) {
+        self.it.cleanup(num, forward);
+    }
 }
 
 #[stable(feature = "iter_copied", since = "1.36.0")]

library/core/src/iter/adapters/enumerate.rs

@@ -241,6 +241,13 @@ where
     I: TrustedRandomAccessNoCoerce,
 {
     const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT;
+
+    const NEEDS_CLEANUP: bool = I::NEEDS_CLEANUP;
+
+    #[inline]
+    fn cleanup(&mut self, num: usize, forward: bool) {
+        self.iter.cleanup(num, forward);
+    }
 }
 
 #[stable(feature = "fused", since = "1.26.0")]

library/core/src/iter/adapters/fuse.rs

@@ -231,6 +231,15 @@ where
     I: TrustedRandomAccessNoCoerce,
 {
     const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT;
+
+    const NEEDS_CLEANUP: bool = I::NEEDS_CLEANUP;
+
+    #[inline]
+    fn cleanup(&mut self, num: usize, forward: bool) {
+        if let Some(iter) = self.iter.as_mut() {
+            iter.cleanup(num, forward);
+        }
+    }
 }
 
 /// Fuse specialization trait

library/core/src/iter/adapters/map.rs

@@ -199,6 +199,13 @@ where
     I: TrustedRandomAccessNoCoerce,
 {
     const MAY_HAVE_SIDE_EFFECT: bool = true;
+
+    const NEEDS_CLEANUP: bool = I::NEEDS_CLEANUP;
+
+    #[inline]
+    fn cleanup(&mut self, num: usize, forward: bool) {
+        self.iter.cleanup(num, forward);
+    }
 }
 
 #[unstable(issue = "none", feature = "inplace_iteration")]

library/core/src/iter/adapters/zip.rs

@@ -400,6 +400,13 @@ where
     B: TrustedRandomAccessNoCoerce,
 {
     const MAY_HAVE_SIDE_EFFECT: bool = A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT;
+
+    const NEEDS_CLEANUP: bool = A::NEEDS_CLEANUP || B::NEEDS_CLEANUP;
+
+    fn cleanup(&mut self, num: usize, forward: bool) {
+        self.a.cleanup(num, forward);
+        self.b.cleanup(num, forward);
+    }
 }
 
 #[stable(feature = "fused", since = "1.26.0")]
@@ -542,9 +549,14 @@ pub unsafe trait TrustedRandomAccessNoCoerce: Sized {
     {
         self.size_hint().0
     }
+
     /// `true` if getting an iterator element may have side effects.
     /// Remember to take inner iterators into account.
     const MAY_HAVE_SIDE_EFFECT: bool;
+
+    const NEEDS_CLEANUP: bool;
+
+    fn cleanup(&mut self, num: usize, forward: bool);
 }
 
 /// Like `Iterator::__iterator_get_unchecked`, but doesn't require the compiler to

library/core/src/iter/loop_desugar.rs

@@ -1,64 +1,72 @@
-use crate::iter::IntoIterator as RealIntoIterator;
+use crate::iter::IntoIterator;
 use crate::iter::TrustedRandomAccessNoCoerce;
 
-#[unstable(feature = "trusted_random_access", issue = "none")]
-#[doc(hidden)]
-
-pub trait IntoIterator {
-    type IntoIter: Iterator;
-
-    #[unstable(feature = "trusted_random_access", issue = "none")]
-    // #[cfg_attr(not(bootstrap), lang = "loop_desugar")]
-    #[cfg_attr(not(bootstrap), lang = "into_iter")]
-    fn into_iter(self) -> Self::IntoIter;
-}
-
-impl<C: RealIntoIterator> IntoIterator for C {
-    type IntoIter = ForLoopDesugar<C::IntoIter>;
-
-    #[inline]
-    fn into_iter(self) -> Self::IntoIter {
-        ForLoopDesugar { iter: <Self as RealIntoIterator>::into_iter(self), idx: 0 }
-    }
-}
-
 #[derive(Debug)]
 #[doc(hidden)]
 #[unstable(feature = "trusted_random_access", issue = "none")]
-pub struct ForLoopDesugar<I> {
+pub struct ForLoopDesugar<I: Iterator> {
     iter: I,
     idx: usize,
 }
 
-#[unstable(feature = "trusted_random_access", issue = "none")]
-impl<I: Iterator> Iterator for ForLoopDesugar<I> {
-    type Item = I::Item;
+impl<I: Iterator> ForLoopDesugar<I> {
+    #[inline]
+    #[cfg_attr(not(bootstrap), lang = "into_iter")]
+    #[cfg_attr(bootstrap, allow(dead_code))]
+    pub fn new(it: impl IntoIterator<Item = I::Item, IntoIter = I>) -> Self {
+        let mut desugar = ForLoopDesugar { iter: it.into_iter(), idx: 0 };
+        desugar.setup();
+        desugar
+    }
 
     #[inline]
-    fn next(&mut self) -> Option<I::Item> {
-        // self.iter.next_spec(&mut self.idx)
+    #[cfg_attr(not(bootstrap), lang = "next")]
+    #[cfg_attr(bootstrap, allow(dead_code))]
+    pub fn next(&mut self) -> Option<I::Item> {
         self.next_spec()
     }
 }
 
+unsafe impl<#[may_dangle] I: Iterator> Drop for ForLoopDesugar<I> {
+    #[inline]
+    fn drop(&mut self) {
+        self.cleanup();
+    }
+}
+
 trait DesugarSpec<T> {
+    fn setup(&mut self);
+
     fn next_spec(&mut self) -> Option<T>;
+
+    fn cleanup(&mut self);
 }
 
 impl<I, T> DesugarSpec<T> for ForLoopDesugar<I>
 where
     I: Iterator<Item = T>,
 {
+    #[inline]
+    default fn setup(&mut self) {}
+
     #[inline]
     default fn next_spec(&mut self) -> Option<I::Item> {
         self.iter.next()
     }
+
+    #[inline]
+    default fn cleanup(&mut self) {}
 }
 
 impl<I, T> DesugarSpec<T> for ForLoopDesugar<I>
 where
     I: TrustedRandomAccessNoCoerce + Iterator<Item = T>,
 {
+    #[inline]
+    fn setup(&mut self) {
+        let _ = self.iter.advance_by(0);
+    }
+
     #[inline]
     fn next_spec(&mut self) -> Option<I::Item> {
         let idx = self.idx;
@@ -73,4 +81,11 @@ where
             None
         }
     }
+
+    #[inline]
+    fn cleanup(&mut self) {
+        if I::NEEDS_CLEANUP {
+            self.iter.cleanup(self.idx, true);
+        }
+    }
 }

library/core/src/iter/range.rs

@@ -502,7 +502,18 @@ macro_rules! unsafe_range_trusted_random_access_impl {
         #[doc(hidden)]
         #[unstable(feature = "trusted_random_access", issue = "none")]
         unsafe impl TrustedRandomAccessNoCoerce for ops::Range<$t> {
+
+            fn cleanup(&mut self, num: usize, forward: bool) {
+                if forward {
+                    let _ = self.advance_by(num);
+                } else {
+                    let _ = self.advance_back_by(num);
+                }
+            }
+
             const MAY_HAVE_SIDE_EFFECT: bool = false;
+
+            const NEEDS_CLEANUP: bool = false;
         }
     )*)
 }

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

@@ -94,8 +94,8 @@ pub trait Iterator {
     /// assert_eq!(None, iter.next());
     /// assert_eq!(None, iter.next());
     /// ```
-    #[lang = "next"]
     #[stable(feature = "rust1", since = "1.0.0")]
+    #[cfg_attr(bootstrap, lang = "next")]
     fn next(&mut self) -> Option<Self::Item>;
 
     /// Returns the bounds on the remaining length of the iterator.

library/core/src/lib.rs

@@ -174,6 +174,7 @@
 #![feature(deprecated_suggestion)]
 #![feature(doc_cfg)]
 #![feature(doc_notable_trait)]
+#![feature(dropck_eyepatch)]
 #![feature(rustdoc_internals)]
 #![feature(exhaustive_patterns)]
 #![feature(doc_cfg_hide)]