Use TrustedRandomAccess for loop desugaring

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

@@ -1,5 +1,5 @@
 use core::fmt;
-use core::iter::{FusedIterator, TrustedLen, TrustedRandomAccess, TrustedRandomAccessNoCoerce};
+use core::iter::{FusedIterator, TrustedLen, TrustedRandomAccess};
 use core::mem::MaybeUninit;
 use core::ops::Try;
 
@@ -205,10 +205,12 @@ unsafe impl<T> TrustedLen for Iter<'_, T> {}
 
 #[doc(hidden)]
 #[unstable(feature = "trusted_random_access", issue = "none")]
-unsafe impl<T> TrustedRandomAccess for Iter<'_, T> {}
+unsafe impl<T> TrustedRandomAccess for Iter<'_, T> {
+    fn cleanup_front(&mut self, num: usize) {
+        let _ = self.advance_by(num);
+    }
 
-#[doc(hidden)]
-#[unstable(feature = "trusted_random_access", issue = "none")]
-unsafe impl<T> TrustedRandomAccessNoCoerce for Iter<'_, T> {
-    const MAY_HAVE_SIDE_EFFECT: bool = false;
+    fn cleanup_back(&mut self, num: usize) {
+        let _ = self.advance_back_by(num);
+    }
 }

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

@@ -1,5 +1,5 @@
 use core::fmt;
-use core::iter::{FusedIterator, TrustedLen, TrustedRandomAccess, TrustedRandomAccessNoCoerce};
+use core::iter::{FusedIterator, TrustedLen, TrustedRandomAccess};
 use core::marker::PhantomData;
 
 use super::{count, wrap_index, RingSlices};
@@ -154,10 +154,12 @@ unsafe impl<T> TrustedLen for IterMut<'_, T> {}
 
 #[doc(hidden)]
 #[unstable(feature = "trusted_random_access", issue = "none")]
-unsafe impl<T> TrustedRandomAccess for IterMut<'_, T> {}
+unsafe impl<T> TrustedRandomAccess for IterMut<'_, T> {
+    fn cleanup_front(&mut self, num: usize) {
+        let _ = self.advance_by(num);
+    }
 
-#[doc(hidden)]
-#[unstable(feature = "trusted_random_access", issue = "none")]
-unsafe impl<T> TrustedRandomAccessNoCoerce for IterMut<'_, T> {
-    const MAY_HAVE_SIDE_EFFECT: bool = false;
+    fn cleanup_back(&mut self, num: usize) {
+        let _ = self.advance_back_by(num);
+    }
 }

library/alloc/src/vec/in_place_collect.rs

@@ -60,7 +60,7 @@
 //! # O(1) collect
 //!
 //! The main iteration itself is further specialized when the iterator implements
-//! [`TrustedRandomAccessNoCoerce`] to let the optimizer see that it is a counted loop with a single
+//! [`TrustedRandomAccess`] to let the optimizer see that it is a counted loop with a single
 //! [induction variable]. This can turn some iterators into a noop, i.e. it reduces them from O(n) to
 //! O(1). This particular optimization is quite fickle and doesn't always work, see [#79308]
 //!
@@ -70,7 +70,7 @@
 //! Since unchecked accesses through that trait do not advance the read pointer of `IntoIter`
 //! this would interact unsoundly with the requirements about dropping the tail described above.
 //! But since the normal `Drop` implementation of `IntoIter` would suffer from the same problem it
-//! is only correct for `TrustedRandomAccessNoCoerce` to be implemented when the items don't
+//! is only correct for `TrustedRandomAccess` to be implemented when the items don't
 //! have a destructor. Thus that implicit requirement also makes the specialization safe to use for
 //! in-place collection.
 //! Note that this safety concern is about the correctness of `impl Drop for IntoIter`,
@@ -134,7 +134,7 @@
 //! }
 //! vec.truncate(write_idx);
 //! ```
-use core::iter::{InPlaceIterable, SourceIter, TrustedRandomAccessNoCoerce};
+use core::iter::{InPlaceIterable, SourceIter, TrustedRandomAccess};
 use core::mem::{self, ManuallyDrop};
 use core::ptr::{self};
 
@@ -195,7 +195,7 @@ where
         // itself once IntoIter goes out of scope.
         // If the drop panics then we also leak any elements collected into dst_buf.
         //
-        // Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce
+        // Note: This access to the source wouldn't be allowed by the TrustedRandomIterator
         // contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the
         // module documenttation why this is ok anyway.
         src.forget_allocation_drop_remaining();
@@ -230,7 +230,7 @@ trait SpecInPlaceCollect<T, I>: Iterator<Item = T> {
     /// collected. `end` is the last writable element of the allocation and used for bounds checks.
     ///
     /// This method is specialized and one of its implementations makes use of
-    /// `Iterator::__iterator_get_unchecked` calls with a `TrustedRandomAccessNoCoerce` bound
+    /// `Iterator::__iterator_get_unchecked` calls with a `TrustedRandomAccess` bound
     /// on `I` which means the caller of this method must take the safety conditions
     /// of that trait into consideration.
     fn collect_in_place(&mut self, dst: *mut T, end: *const T) -> usize;
@@ -256,7 +256,7 @@ where
 
 impl<T, I> SpecInPlaceCollect<T, I> for I
 where
-    I: Iterator<Item = T> + TrustedRandomAccessNoCoerce,
+    I: Iterator<Item = T> + TrustedRandomAccess,
 {
     #[inline]
     fn collect_in_place(&mut self, dst_buf: *mut T, end: *const T) -> usize {
@@ -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

@@ -5,7 +5,8 @@ use crate::raw_vec::RawVec;
 use core::fmt;
 use core::intrinsics::arith_offset;
 use core::iter::{
-    FusedIterator, InPlaceIterable, SourceIter, TrustedLen, TrustedRandomAccessNoCoerce,
+    FusedIterator, InPlaceIterable, SourceIter, TrustedLen, TrustedRandomAccess,
+    TrustedRandomAccessNeedsCleanup,
 };
 use core::marker::PhantomData;
 use core::mem::{self, ManuallyDrop};
@@ -200,7 +201,7 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
     #[doc(hidden)]
     unsafe fn __iterator_get_unchecked(&mut self, i: usize) -> Self::Item
     where
-        Self: TrustedRandomAccessNoCoerce,
+        Self: TrustedRandomAccess,
     {
         // SAFETY: the caller must guarantee that `i` is in bounds of the
         // `Vec<T>`, so `i` cannot overflow an `isize`, and the `self.ptr.add(i)`
@@ -284,15 +285,39 @@ impl<T: Copy> NonDrop for T {}
 
 #[doc(hidden)]
 #[unstable(issue = "none", feature = "std_internals")]
-// TrustedRandomAccess (without NoCoerce) must not be implemented because
-// subtypes/supertypes of `T` might not be `NonDrop`
-unsafe impl<T, A: Allocator> TrustedRandomAccessNoCoerce for IntoIter<T, A>
+unsafe impl<T, A: Allocator> TrustedRandomAccess for IntoIter<T, A>
 where
     T: NonDrop,
 {
-    const MAY_HAVE_SIDE_EFFECT: bool = false;
+    fn cleanup_front(&mut self, num: usize) {
+        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) };
+        }
+    }
+
+    fn cleanup_back(&mut self, num: usize) {
+        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) };
+        }
+    }
 }
 
+#[doc(hidden)]
+#[unstable(issue = "none", feature = "std_internals")]
+unsafe impl<T, A: Allocator> TrustedRandomAccessNeedsCleanup for IntoIter<T, A> where T: NonDrop {}
+
 #[cfg(not(no_global_oom_handling))]
 #[stable(feature = "vec_into_iter_clone", since = "1.8.0")]
 impl<T: Clone, A: Allocator + Clone> Clone for IntoIter<T, A> {