Auto merge of #120682 - the8472:indexed-access, r=<try>

[WIP] rewrite TrustedRandomAccess into two directional variants

r? `@ghost`

Author: bors

library/alloc/src/lib.rs

@@ -160,6 +160,7 @@
 #![feature(str_internals)]
 #![feature(strict_provenance)]
 #![feature(trusted_fused)]
+#![feature(trusted_indexed_access)]
 #![feature(trusted_len)]
 #![feature(trusted_random_access)]
 #![feature(try_trait_v2)]

library/alloc/src/vec/in_place_collect.rs

@@ -80,7 +80,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
+//! [`UncheckedIndexedIterator`] 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]
 //!
@@ -157,8 +157,10 @@
 use crate::alloc::{handle_alloc_error, Global};
 use core::alloc::Allocator;
 use core::alloc::Layout;
-use core::iter::{InPlaceIterable, SourceIter, TrustedRandomAccessNoCoerce};
+use core::iter::UncheckedIndexedIterator;
+use core::iter::{InPlaceIterable, SourceIter};
 use core::marker::PhantomData;
+use core::mem::needs_drop;
 use core::mem::{self, ManuallyDrop, SizedTypeProperties};
 use core::num::NonZero;
 use core::ptr::{self, NonNull};
@@ -257,8 +259,9 @@ where
         // caveat: if they weren't we might not even make it to this point
         debug_assert_eq!(src_buf, src.buf.as_ptr());
         // check InPlaceIterable contract. This is only possible if the iterator advanced the
-        // source pointer at all. If it uses unchecked access via TrustedRandomAccess
-        // then the source pointer will stay in its initial position and we can't use it as reference
+        // source pointer at all. If it uses unchecked access via UncheckedIndexedIterator
+        // and doesn't perform cleanup then the source pointer will stay in its initial position
+        // and we can't use it as reference.
         if src.ptr != src_ptr {
             debug_assert!(
                 unsafe { dst_buf.add(len) as *const _ } <= src.ptr.as_ptr(),
@@ -369,28 +372,96 @@ where
     }
 }
 
+// impl<T, I> SpecInPlaceCollect<T, I> for I
+// where
+//     I: Iterator<Item = T> + TrustedRandomAccessNoCoerce,
+// {
+//     #[inline]
+//     unsafe fn collect_in_place(&mut self, dst_buf: *mut T, end: *const T) -> usize {
+//         let len = self.size();
+//         let mut drop_guard = InPlaceDrop { inner: dst_buf, dst: dst_buf };
+//         for i in 0..len {
+//             // Safety: InplaceIterable contract guarantees that for every element we read
+//             // one slot in the underlying storage will have been freed up and we can immediately
+//             // write back the result.
+//             unsafe {
+//                 let dst = dst_buf.add(i);
+//                 debug_assert!(dst as *const _ <= end, "InPlaceIterable contract violation");
+//                 ptr::write(dst, self.__iterator_get_unchecked(i));
+//                 // Since this executes user code which can panic we have to bump the pointer
+//                 // after each step.
+//                 drop_guard.dst = dst.add(1);
+//             }
+//         }
+//         mem::forget(drop_guard);
+//         len
+//     }
+// }
+
 impl<T, I> SpecInPlaceCollect<T, I> for I
 where
-    I: Iterator<Item = T> + TrustedRandomAccessNoCoerce,
+    I: Iterator<Item = T> + UncheckedIndexedIterator,
 {
     #[inline]
     unsafe fn collect_in_place(&mut self, dst_buf: *mut T, end: *const T) -> usize {
-        let len = self.size();
-        let mut drop_guard = InPlaceDrop { inner: dst_buf, dst: dst_buf };
-        for i in 0..len {
-            // Safety: InplaceIterable contract guarantees that for every element we read
-            // one slot in the underlying storage will have been freed up and we can immediately
-            // write back the result.
+        let len = self.size_hint().0;
+
+        if len == 0 {
+            return 0;
+        }
+
+        struct LoopGuard<'a, I>
+        where
+            I: Iterator + UncheckedIndexedIterator,
+        {
+            it: &'a mut I,
+            len: usize,
+            idx: usize,
+            dst_buf: *mut I::Item,
+        }
+
+        impl<I> Drop for LoopGuard<'_, I>
+        where
+            I: Iterator + UncheckedIndexedIterator,
+        {
+            #[inline]
+            fn drop(&mut self) {
+                unsafe {
+                    let new_len = self.len - self.idx;
+                    if I::CLEANUP_ON_DROP {
+                        self.it.set_front_index_from_end_unchecked(new_len, self.len);
+                    }
+                    if needs_drop::<I::Item>() && self.idx != self.len {
+                        let raw_slice =
+                            ptr::slice_from_raw_parts_mut::<I::Item>(self.dst_buf, self.idx);
+                        ptr::drop_in_place(raw_slice);
+                    }
+                }
+            }
+        }
+
+        let mut state = LoopGuard { it: self, len, idx: 0, dst_buf };
+
+        loop {
             unsafe {
-                let dst = dst_buf.add(i);
+                let idx = state.idx;
+                state.idx = idx.unchecked_add(1);
+                let dst = state.dst_buf.add(idx);
                 debug_assert!(dst as *const _ <= end, "InPlaceIterable contract violation");
-                ptr::write(dst, self.__iterator_get_unchecked(i));
-                // Since this executes user code which can panic we have to bump the pointer
-                // after each step.
-                drop_guard.dst = dst.add(1);
+                dst.write(state.it.index_from_end_unchecked(len - idx));
+            }
+            if state.idx == len {
+                break;
             }
         }
-        mem::forget(drop_guard);
+
+        // disarm guard, we don't want the front elements to get dropped
+        mem::forget(state);
+        // since the guard was disarmed, update the iterator state
+        if Self::CLEANUP_ON_DROP {
+            unsafe { self.set_front_index_from_end_unchecked(0, len) };
+        }
+
         len
     }
 }

library/alloc/src/vec/into_iter.rs

@@ -8,10 +8,10 @@ use core::array;
 use core::fmt;
 use core::iter::{
     FusedIterator, InPlaceIterable, SourceIter, TrustedFused, TrustedLen,
-    TrustedRandomAccessNoCoerce,
+    TrustedRandomAccessNoCoerce, UncheckedIndexedIterator,
 };
 use core::marker::PhantomData;
-use core::mem::{ManuallyDrop, MaybeUninit, SizedTypeProperties};
+use core::mem::{self, ManuallyDrop, MaybeUninit, SizedTypeProperties};
 use core::num::NonZero;
 #[cfg(not(no_global_oom_handling))]
 use core::ops::Deref;
@@ -224,7 +224,8 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
         let exact = if T::IS_ZST {
             self.end.addr().wrapping_sub(self.ptr.as_ptr().addr())
         } else {
-            unsafe { non_null!(self.end, T).sub_ptr(self.ptr) }
+            // FIXME(#121239): this should use sub_ptr but llvm == 18 doesn't optimize as it should
+            unsafe { non_null!(self.end, T).offset_from(self.ptr) as usize }
         };
         (exact, Some(exact))
     }
@@ -303,6 +304,57 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
         // them for `Drop`.
         unsafe { self.ptr.add(i).read() }
     }
+
+    #[inline]
+    unsafe fn index_from_end_unchecked(&mut self, idx: usize) -> Self::Item {
+        if T::IS_ZST {
+            // SAFETY: conjuring a ZST
+            unsafe { NonNull::dangling().read() }
+        } else {
+            let end = non_null!(self.end, T);
+            unsafe { end.sub(idx).read() }
+            //self.ptr = unsafe { end.sub(idx).add(1) };
+            //unsafe { self.end.sub(idx).read() }
+        }
+    }
+
+    #[inline]
+    unsafe fn index_from_start_unchecked(&mut self, idx: usize) -> Self::Item {
+        if T::IS_ZST {
+            // SAFETY: conjuring a ZST
+            unsafe { NonNull::dangling().read() }
+        } else {
+            //let end = non_null!(mut self.end, T);
+            //*end = unsafe { self.ptr.add(idx) };
+            unsafe { self.ptr.add(idx).read() }
+        }
+    }
+}
+
+#[unstable(feature = "trusted_indexed_access", issue = "none")]
+impl<T, A: Allocator> UncheckedIndexedIterator for IntoIter<T, A> {
+    const MAY_HAVE_SIDE_EFFECT: bool = false;
+    const CLEANUP_ON_DROP: bool = mem::needs_drop::<T>();
+
+    #[inline]
+    unsafe fn set_front_index_from_end_unchecked(&mut self, new_len: usize, _old_len: usize) {
+        if T::IS_ZST {
+            self.end = self.ptr.as_ptr().cast_const().wrapping_byte_add(new_len);
+        } else {
+            let end = non_null!(self.end, T);
+            self.ptr = unsafe { end.sub(new_len) };
+        }
+    }
+
+    #[inline]
+    unsafe fn set_end_index_from_start_unchecked(&mut self, new_len: usize, _old_len: usize) {
+        if T::IS_ZST {
+            self.end = self.ptr.as_ptr().cast_const().wrapping_byte_add(new_len);
+        } else {
+            let end = non_null!(mut self.end, T);
+            *end = unsafe { self.ptr.add(new_len) };
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

library/core/benches/iter.rs

@@ -457,11 +457,12 @@ fn bench_trusted_random_access_adapters(b: &mut Bencher) {
             .map(|(a, b)| a.wrapping_add(b))
             .fuse();
         let mut acc: usize = 0;
-        let size = iter.size();
+        let size = iter.size_hint().0;
         for i in 0..size {
             // SAFETY: TRA requirements are satisfied by 0..size iteration and then dropping the
             // iterator.
-            acc = acc.wrapping_add(unsafe { iter.__iterator_get_unchecked(i) });
+            // The iterator is not owning, so we skip cleanup.
+            acc = acc.wrapping_add(unsafe { iter.index_from_start_unchecked(i) });
         }
         acc
     })

library/core/benches/lib.rs

@@ -3,6 +3,7 @@
 #![feature(flt2dec)]
 #![feature(test)]
 #![feature(trusted_random_access)]
+#![feature(trusted_indexed_access)]
 #![feature(iter_array_chunks)]
 #![feature(iter_next_chunk)]
 #![feature(iter_advance_by)]

library/core/src/array/iter.rs

@@ -4,8 +4,10 @@ use crate::num::NonZero;
 use crate::{
     fmt,
     intrinsics::transmute_unchecked,
-    iter::{self, FusedIterator, TrustedLen, TrustedRandomAccessNoCoerce},
-    mem::MaybeUninit,
+    iter::{
+        self, FusedIterator, TrustedLen, TrustedRandomAccessNoCoerce, UncheckedIndexedIterator,
+    },
+    mem::{self, MaybeUninit},
     ops::{IndexRange, Range},
     ptr,
 };
@@ -300,6 +302,24 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
         // SAFETY: The caller must provide an idx that is in bound of the remainder.
         unsafe { self.data.as_ptr().add(self.alive.start()).add(idx).cast::<T>().read() }
     }
+
+    #[inline]
+    unsafe fn index_from_end_unchecked(&mut self, idx: usize) -> Self::Item
+    where
+        Self: UncheckedIndexedIterator,
+    {
+        // SAFETY: The caller must provide an idx that is in bound of the remainder.
+        unsafe { self.data.as_ptr().add(self.alive.end()).sub(idx).cast::<T>().read() }
+    }
+
+    #[inline]
+    unsafe fn index_from_start_unchecked(&mut self, idx: usize) -> Self::Item
+    where
+        Self: UncheckedIndexedIterator,
+    {
+        // SAFETY: The caller must provide an idx that is in bound of the remainder.
+        unsafe { self.data.as_ptr().add(self.alive.start()).add(idx).cast::<T>().read() }
+    }
 }
 
 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
@@ -400,6 +420,24 @@ where
     const MAY_HAVE_SIDE_EFFECT: bool = false;
 }
 
+#[unstable(feature = "trusted_indexed_access", issue = "none")]
+impl<T, const N: usize> UncheckedIndexedIterator for IntoIter<T, N> {
+    const MAY_HAVE_SIDE_EFFECT: bool = false;
+    const CLEANUP_ON_DROP: bool = mem::needs_drop::<T>();
+
+    #[inline]
+    unsafe fn set_front_index_from_end_unchecked(&mut self, new_len: usize, _old_len: usize) {
+        // SAFETY: ...
+        unsafe { self.alive.set_start_unchecked(self.alive.end() - new_len) };
+    }
+
+    #[inline]
+    unsafe fn set_end_index_from_start_unchecked(&mut self, new_len: usize, _old_len: usize) {
+        // SAFETY: ...
+        unsafe { self.alive.set_end_unchecked(self.alive.start() + new_len) };
+    }
+}
+
 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
 impl<T: Clone, const N: usize> Clone for IntoIter<T, N> {
     fn clone(&self) -> Self {

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

@@ -1,7 +1,10 @@
 use crate::iter::adapters::{
     zip::try_get_unchecked, SourceIter, TrustedRandomAccess, TrustedRandomAccessNoCoerce,
 };
-use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen, UncheckedIterator};
+use crate::iter::traits::SpecIndexedAccess as _;
+use crate::iter::{
+    FusedIterator, InPlaceIterable, TrustedLen, UncheckedIndexedIterator, UncheckedIterator,
+};
 use crate::ops::Try;
 use core::num::NonZero;
 
@@ -69,6 +72,24 @@ where
         // `Iterator::__iterator_get_unchecked`.
         unsafe { try_get_unchecked(&mut self.it, idx).clone() }
     }
+
+    #[inline]
+    unsafe fn index_from_end_unchecked(&mut self, idx: usize) -> Self::Item
+    where
+        Self: UncheckedIndexedIterator,
+    {
+        // SAFETY: forwarding to unsafe function with the same preconditions
+        unsafe { self.it.index_from_end_unchecked_inner(idx) }.clone()
+    }
+
+    #[inline]
+    unsafe fn index_from_start_unchecked(&mut self, idx: usize) -> Self::Item
+    where
+        Self: UncheckedIndexedIterator,
+    {
+        // SAFETY: forwarding to unsafe function with the same preconditions
+        unsafe { self.it.index_from_start_unchecked_inner(idx) }.clone()
+    }
 }
 
 #[stable(feature = "iter_cloned", since = "1.1.0")]
@@ -134,6 +155,27 @@ where
     const MAY_HAVE_SIDE_EFFECT: bool = true;
 }
 
+#[unstable(feature = "trusted_indexed_access", issue = "none")]
+impl<I> UncheckedIndexedIterator for Cloned<I>
+where
+    I: UncheckedIndexedIterator,
+{
+    #[inline]
+    unsafe fn set_front_index_from_end_unchecked(&mut self, new_len: usize, old_len: usize) {
+        // SAFETY: forwarding to unsafe function with the same preconditions
+        unsafe { self.it.set_front_index_from_end_unchecked(new_len, old_len) }
+    }
+
+    #[inline]
+    unsafe fn set_end_index_from_start_unchecked(&mut self, new_len: usize, old_len: usize) {
+        // SAFETY: forwarding to unsafe function with the same preconditions
+        unsafe { self.it.set_end_index_from_start_unchecked(new_len, old_len) }
+    }
+
+    const MAY_HAVE_SIDE_EFFECT: bool = true;
+    const CLEANUP_ON_DROP: bool = I::CLEANUP_ON_DROP;
+}
+
 #[unstable(feature = "trusted_len", issue = "37572")]
 unsafe impl<'a, I, T: 'a> TrustedLen for Cloned<I>
 where