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

@@ -156,6 +156,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

@@ -157,7 +157,8 @@
 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::{self, ManuallyDrop, SizedTypeProperties};
 use core::num::NonZeroUsize;
@@ -237,7 +238,7 @@ where
             return SpecFromIterNested::from_iter(iterator);
         }
 
-        let (src_buf, src_ptr, src_cap, mut dst_buf, dst_end, dst_cap) = unsafe {
+        let (src_buf, _src_ptr, src_cap, mut dst_buf, dst_end, dst_cap) = unsafe {
             let inner = iterator.as_inner().as_into_iter();
             (
                 inner.buf.as_ptr(),
@@ -256,15 +257,11 @@ where
         // check if SourceIter contract was upheld
         // 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
-        if src.ptr != src_ptr {
-            debug_assert!(
-                unsafe { dst_buf.add(len) as *const _ } <= src.ptr.as_ptr(),
-                "InPlaceIterable contract violation, write pointer advanced beyond read pointer"
-            );
-        }
+        // check InPlaceIterable contract.
+        debug_assert!(
+            unsafe { dst_buf.add(len) as *const _ } <= src.ptr.as_ptr(),
+            "InPlaceIterable contract violation, write pointer advanced beyond read pointer"
+        );
 
         // The ownership of the source allocation and the new `T` values is temporarily moved into `dst_guard`.
         // This is safe because
@@ -369,28 +366,93 @@ 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;
+                    self.it.set_front_index_from_end_unchecked(new_len, self.len);
+                    if 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 = state.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(state.len - idx));
+                //drop_guard.dst = dst_buf.add(i).add(1);
+            }
+            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 is disarmed, update the iterator state
+        unsafe { self.set_front_index_from_end_unchecked(0, len) };
+
         len
     }
 }

library/alloc/src/vec/into_iter.rs

@@ -8,7 +8,7 @@ use core::array;
 use core::fmt;
 use core::iter::{
     FusedIterator, InPlaceIterable, SourceIter, TrustedFused, TrustedLen,
-    TrustedRandomAccessNoCoerce,
+    TrustedRandomAccessNoCoerce, UncheckedIndexedIterator,
 };
 use core::marker::PhantomData;
 use core::mem::{self, ManuallyDrop, MaybeUninit, SizedTypeProperties};
@@ -129,6 +129,7 @@ impl<T, A: Allocator> IntoIter<T, A> {
     /// This method is used by in-place iteration, refer to the vec::in_place_collect
     /// documentation for an overview.
     #[cfg(not(no_global_oom_handling))]
+    #[inline]
     pub(super) fn forget_allocation_drop_remaining(&mut self) {
         let remaining = self.as_raw_mut_slice();
 
@@ -307,6 +308,51 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
         // them for `Drop`.
         unsafe { if T::IS_ZST { mem::zeroed() } else { 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 { mem::zeroed() }
+        } else {
+            let end = non_null!(self.end, T);
+            //self.ptr = unsafe { end.sub(idx).add(1) };
+            unsafe { 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 { mem::zeroed() }
+        } 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> {
+    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) };
+        }
+    }
+
+    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/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::NonZeroUsize;
 
@@ -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,24 @@ 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) }
+    }
+}
+
 #[unstable(feature = "trusted_len", issue = "37572")]
 unsafe impl<'a, I, T: 'a> TrustedLen for Cloned<I>
 where

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

@@ -1,7 +1,8 @@
 use crate::iter::adapters::{
     zip::try_get_unchecked, SourceIter, TrustedRandomAccess, TrustedRandomAccessNoCoerce,
 };
-use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen};
+use crate::iter::traits::SpecIndexedAccess;
+use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen, UncheckedIndexedIterator};
 use crate::mem::MaybeUninit;
 use crate::mem::SizedTypeProperties;
 use crate::num::NonZeroUsize;
@@ -102,6 +103,24 @@ where
         // `Iterator::__iterator_get_unchecked`.
         *unsafe { try_get_unchecked(&mut self.it, idx) }
     }
+
+    #[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) }
+    }
+
+    #[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) }
+    }
 }
 
 #[stable(feature = "iter_copied", since = "1.36.0")]
@@ -159,6 +178,24 @@ where
 {
 }
 
+#[unstable(feature = "trusted_indexed_access", issue = "none")]
+impl<I> UncheckedIndexedIterator for Copied<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) }
+    }
+}
+
 #[doc(hidden)]
 #[unstable(feature = "trusted_random_access", issue = "none")]
 unsafe impl<I> TrustedRandomAccess for Copied<I> where I: TrustedRandomAccess {}