From fc0d3aeeddc93086383e7d6829103ebeb2513cd7 Mon Sep 17 00:00:00 2001 From: The 8472 Date: Wed, 6 Mar 2024 18:22:35 +0100 Subject: [PATCH 1/3] workaround missed optimization in llvm18 --- library/alloc/src/vec/into_iter.rs | 3 ++- tests/codegen/vec-iter.rs | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs index dfd42ca06193a..53c76b94a8c83 100644 --- a/library/alloc/src/vec/into_iter.rs +++ b/library/alloc/src/vec/into_iter.rs @@ -224,7 +224,8 @@ impl Iterator for IntoIter { 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)) } diff --git a/tests/codegen/vec-iter.rs b/tests/codegen/vec-iter.rs index 310680969c4fe..339f4cee711dd 100644 --- a/tests/codegen/vec-iter.rs +++ b/tests/codegen/vec-iter.rs @@ -13,7 +13,7 @@ pub fn vec_iter_len_nonnull(it: &vec::IntoIter) -> usize { // CHECK: load ptr // CHECK-SAME: !nonnull // CHECK-SAME: !noundef - // CHECK: sub nuw + // CHECK: sub // CHECK: ret it.len() } From 36d8523dc6286d115516c1641f4e0b2f522b1acb Mon Sep 17 00:00:00 2001 From: The 8472 Date: Mon, 5 Feb 2024 21:02:32 +0100 Subject: [PATCH 2/3] [WIP] rewrite TrustedRandomAccess into two directional variants --- library/alloc/src/lib.rs | 1 + library/alloc/src/vec/in_place_collect.rs | 105 +++++- library/alloc/src/vec/into_iter.rs | 55 ++- library/core/benches/iter.rs | 5 +- library/core/benches/lib.rs | 1 + library/core/src/array/iter.rs | 42 ++- library/core/src/iter/adapters/cloned.rs | 44 ++- library/core/src/iter/adapters/copied.rs | 42 ++- library/core/src/iter/adapters/enumerate.rs | 48 ++- library/core/src/iter/adapters/fuse.rs | 61 +++ library/core/src/iter/adapters/map.rs | 51 ++- library/core/src/iter/adapters/zip.rs | 357 +++++++++++++++++- library/core/src/iter/mod.rs | 2 + library/core/src/iter/range.rs | 40 ++ library/core/src/iter/traits/iterator.rs | 19 + library/core/src/iter/traits/mod.rs | 4 + .../src/iter/traits/unchecked_iterator.rs | 71 ++++ library/core/src/ops/index_range.rs | 16 + library/core/src/slice/iter.rs | 3 +- library/core/src/slice/iter/macros.rs | 54 +++ library/core/tests/iter/adapters/cloned.rs | 3 +- library/core/tests/iter/adapters/zip.rs | 14 +- library/core/tests/lib.rs | 1 + tests/codegen/vec-in-place.rs | 8 + 24 files changed, 1006 insertions(+), 41 deletions(-) diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index 28695ade5bf55..ec0921d492c3f 100644 --- a/library/alloc/src/lib.rs +++ b/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)] diff --git a/library/alloc/src/vec/in_place_collect.rs b/library/alloc/src/vec/in_place_collect.rs index 07eb91c900597..d04ac3eada0c4 100644 --- a/library/alloc/src/vec/in_place_collect.rs +++ b/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 SpecInPlaceCollect for I +// where +// I: Iterator + 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 SpecInPlaceCollect for I where - I: Iterator + TrustedRandomAccessNoCoerce, + I: Iterator + 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 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::() && self.idx != self.len { + let raw_slice = + ptr::slice_from_raw_parts_mut::(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 } } diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs index 53c76b94a8c83..e8c3eb3f48146 100644 --- a/library/alloc/src/vec/into_iter.rs +++ b/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; @@ -304,6 +304,57 @@ impl Iterator for IntoIter { // 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 UncheckedIndexedIterator for IntoIter { + const MAY_HAVE_SIDE_EFFECT: bool = false; + const CLEANUP_ON_DROP: bool = mem::needs_drop::(); + + #[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")] diff --git a/library/core/benches/iter.rs b/library/core/benches/iter.rs index c1cec5e6d3c8c..0591bcb3d4cc5 100644 --- a/library/core/benches/iter.rs +++ b/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 }) diff --git a/library/core/benches/lib.rs b/library/core/benches/lib.rs index 4d14b930e4171..99188f7a3232a 100644 --- a/library/core/benches/lib.rs +++ b/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)] diff --git a/library/core/src/array/iter.rs b/library/core/src/array/iter.rs index e3d2cd2a31fbc..d0b0f00215ac4 100644 --- a/library/core/src/array/iter.rs +++ b/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 Iterator for IntoIter { // 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::().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::().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::().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 UncheckedIndexedIterator for IntoIter { + const MAY_HAVE_SIDE_EFFECT: bool = false; + const CLEANUP_ON_DROP: bool = mem::needs_drop::(); + + #[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 Clone for IntoIter { fn clone(&self) -> Self { diff --git a/library/core/src/iter/adapters/cloned.rs b/library/core/src/iter/adapters/cloned.rs index 1a106ef97942b..671a527a15711 100644 --- a/library/core/src/iter/adapters/cloned.rs +++ b/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 UncheckedIndexedIterator for Cloned +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 where diff --git a/library/core/src/iter/adapters/copied.rs b/library/core/src/iter/adapters/copied.rs index 6d82d1581f79d..99441d20183b2 100644 --- a/library/core/src/iter/adapters/copied.rs +++ b/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::NonZero; @@ -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,27 @@ where { } +#[unstable(feature = "trusted_indexed_access", issue = "none")] +impl UncheckedIndexedIterator for Copied +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 = I::MAY_HAVE_SIDE_EFFECT; + const CLEANUP_ON_DROP: bool = I::CLEANUP_ON_DROP; +} + #[doc(hidden)] #[unstable(feature = "trusted_random_access", issue = "none")] unsafe impl TrustedRandomAccess for Copied where I: TrustedRandomAccess {} diff --git a/library/core/src/iter/adapters/enumerate.rs b/library/core/src/iter/adapters/enumerate.rs index ef46040f0a703..0bee84e6aaf33 100644 --- a/library/core/src/iter/adapters/enumerate.rs +++ b/library/core/src/iter/adapters/enumerate.rs @@ -1,7 +1,10 @@ use crate::iter::adapters::{ zip::try_get_unchecked, SourceIter, TrustedRandomAccess, TrustedRandomAccessNoCoerce, }; -use crate::iter::{FusedIterator, InPlaceIterable, TrustedFused, TrustedLen}; +use crate::iter::traits::SpecIndexedAccess as _; +use crate::iter::{ + FusedIterator, InPlaceIterable, TrustedFused, TrustedLen, UncheckedIndexedIterator, +}; use crate::num::NonZero; use crate::ops::Try; @@ -136,6 +139,27 @@ where let value = unsafe { try_get_unchecked(&mut self.iter, idx) }; (self.count + idx, value) } + + #[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 + let val = unsafe { self.iter.index_from_end_unchecked_inner(idx) }; + let unadjusted_len = self.iter.size_hint().0; + (unadjusted_len - idx + self.count, val) + } + + #[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 + let val = unsafe { self.iter.index_from_start_unchecked_inner(idx) }; + (self.count + idx, val) + } } #[stable(feature = "rust1", since = "1.0.0")] @@ -240,6 +264,28 @@ where const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT; } +#[unstable(feature = "trusted_indexed_access", issue = "none")] +impl UncheckedIndexedIterator for Enumerate +where + I: UncheckedIndexedIterator, +{ + const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT; + const CLEANUP_ON_DROP: bool = I::CLEANUP_ON_DROP; + + #[inline] + unsafe fn set_front_index_from_end_unchecked(&mut self, new_len: usize, old_len: usize) { + self.count += old_len - new_len; + // SAFETY: forwarding to unsafe function with the same preconditions + unsafe { self.iter.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.iter.set_end_index_from_start_unchecked(new_len, old_len) } + } +} + #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Enumerate where I: FusedIterator {} diff --git a/library/core/src/iter/adapters/fuse.rs b/library/core/src/iter/adapters/fuse.rs index 7781ed088b76c..1b23625cbc939 100644 --- a/library/core/src/iter/adapters/fuse.rs +++ b/library/core/src/iter/adapters/fuse.rs @@ -1,8 +1,10 @@ use crate::intrinsics; use crate::iter::adapters::zip::try_get_unchecked; use crate::iter::adapters::SourceIter; +use crate::iter::traits::SpecIndexedAccess; use crate::iter::{ FusedIterator, TrustedFused, TrustedLen, TrustedRandomAccess, TrustedRandomAccessNoCoerce, + UncheckedIndexedIterator, }; use crate::ops::Try; @@ -121,6 +123,34 @@ where None => unsafe { intrinsics::unreachable() }, } } + + #[inline] + unsafe fn index_from_end_unchecked(&mut self, idx: usize) -> Self::Item + where + Self: UncheckedIndexedIterator, + { + match self.iter { + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + Some(ref mut iter) => unsafe { iter.index_from_end_unchecked_inner(idx) }, + // SAFETY: the caller asserts there is an item at `i`, so we're not exhausted. + None => unsafe { intrinsics::unreachable() }, + } + } + + #[inline] + unsafe fn index_from_start_unchecked(&mut self, idx: usize) -> Self::Item + where + Self: UncheckedIndexedIterator, + { + match self.iter { + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + Some(ref mut iter) => unsafe { iter.index_from_start_unchecked_inner(idx) }, + // SAFETY: the caller asserts there is an item at `i`, so we're not exhausted. + None => unsafe { intrinsics::unreachable() }, + } + } } #[stable(feature = "rust1", since = "1.0.0")] @@ -227,6 +257,37 @@ where const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT; } +#[unstable(feature = "trusted_indexed_access", issue = "none")] +impl UncheckedIndexedIterator for Fuse +where + I: UncheckedIndexedIterator, +{ + const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT; + const CLEANUP_ON_DROP: bool = I::CLEANUP_ON_DROP; + + #[inline] + unsafe fn set_front_index_from_end_unchecked(&mut self, new_len: usize, old_len: usize) { + // SAFETY: FIXME: is this unwrap ok? + unsafe { + self.iter + .as_mut() + .unwrap_unchecked() + .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: FIXME: is this unwrap ok? + unsafe { + self.iter + .as_mut() + .unwrap_unchecked() + .set_end_index_from_start_unchecked(new_len, old_len) + } + } +} + /// Fuse specialization trait /// /// We only need to worry about `&mut self` methods, which diff --git a/library/core/src/iter/adapters/map.rs b/library/core/src/iter/adapters/map.rs index 6e163e20d8ec4..459b8259d7305 100644 --- a/library/core/src/iter/adapters/map.rs +++ b/library/core/src/iter/adapters/map.rs @@ -2,7 +2,11 @@ use crate::fmt; use crate::iter::adapters::{ zip::try_get_unchecked, SourceIter, TrustedRandomAccess, TrustedRandomAccessNoCoerce, }; -use crate::iter::{FusedIterator, InPlaceIterable, TrustedFused, TrustedLen, UncheckedIterator}; +use crate::iter::traits::SpecIndexedAccess; +use crate::iter::{ + FusedIterator, InPlaceIterable, TrustedFused, TrustedLen, UncheckedIndexedIterator, + UncheckedIterator, +}; use crate::num::NonZero; use crate::ops::Try; @@ -138,6 +142,51 @@ where // `Iterator::__iterator_get_unchecked`. unsafe { (self.f)(try_get_unchecked(&mut self.iter, 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 + let val = unsafe { self.iter.index_from_end_unchecked_inner(idx) }; + (self.f)(val) + } + + #[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 + let val = unsafe { self.iter.index_from_start_unchecked_inner(idx) }; + (self.f)(val) + } +} + +#[unstable(feature = "trusted_indexed_access", issue = "none")] +impl UncheckedIndexedIterator for Map +where + I: UncheckedIndexedIterator, +{ + const MAY_HAVE_SIDE_EFFECT: bool = true; + const CLEANUP_ON_DROP: bool = I::CLEANUP_ON_DROP; + + #[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.iter.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.iter.set_end_index_from_start_unchecked(new_len, old_len); + } + } } #[stable(feature = "rust1", since = "1.0.0")] diff --git a/library/core/src/iter/adapters/zip.rs b/library/core/src/iter/adapters/zip.rs index 2e885f06b5272..711af710ee6dd 100644 --- a/library/core/src/iter/adapters/zip.rs +++ b/library/core/src/iter/adapters/zip.rs @@ -1,7 +1,9 @@ use crate::cmp; use crate::fmt::{self, Debug}; -use crate::iter::{FusedIterator, TrustedFused}; +use crate::hint::assert_unchecked; +use crate::iter::{FusedIterator, TrustedFused, UncheckedIndexedIterator}; use crate::iter::{InPlaceIterable, SourceIter, TrustedLen, UncheckedIterator}; +use crate::mem; use crate::num::NonZero; /// An iterator that iterates two other iterators simultaneously. @@ -15,9 +17,9 @@ pub struct Zip { a: A, b: B, // index, len and a_len are only used by the specialized version of zip - index: usize, - len: usize, - a_len: usize, + remaining: usize, + a_offset: usize, + b_offset: usize, } impl Zip { pub(in crate::iter) fn new(a: A, b: B) -> Zip { @@ -112,6 +114,26 @@ where // requirements as `Iterator::__iterator_get_unchecked`. unsafe { ZipImpl::get_unchecked(self, idx) } } + + #[inline] + unsafe fn index_from_end_unchecked(&mut self, idx: usize) -> Self::Item + where + Self: UncheckedIndexedIterator, + { + // SAFETY: `ZipImpl::get_unchecked_from_end` has same safety + // requirements as `Iterator::get_unchecked_from_end`. + unsafe { ZipImpl::get_unchecked_from_end(self, idx) } + } + + #[inline] + unsafe fn index_from_start_unchecked(&mut self, idx: usize) -> Self::Item + where + Self: UncheckedIndexedIterator, + { + // SAFETY: `ZipImpl::get_unchecked_from_start` has same safety + // requirements as `Iterator::get_unchecked_from_start`. + unsafe { ZipImpl::get_unchecked_from_start(self, idx) } + } } #[stable(feature = "rust1", since = "1.0.0")] @@ -145,6 +167,14 @@ trait ZipImpl { unsafe fn get_unchecked(&mut self, idx: usize) -> ::Item where Self: Iterator + TrustedRandomAccessNoCoerce; + + unsafe fn get_unchecked_from_end(&mut self, idx: usize) -> ::Item + where + Self: Iterator + UncheckedIndexedIterator; + + unsafe fn get_unchecked_from_start(&mut self, idx: usize) -> ::Item + where + Self: Iterator + UncheckedIndexedIterator; } // Work around limitations of specialization, requiring `default` impls to be repeated @@ -155,9 +185,9 @@ macro_rules! zip_impl_general_defaults { Zip { a, b, - index: 0, // unused - len: 0, // unused - a_len: 0, // unused + remaining: 0, // unused + a_offset: 0, // unused + b_offset: 0, // unused } } @@ -248,8 +278,23 @@ where { SpecFold::spec_fold(self, init, f) } + + default unsafe fn get_unchecked_from_end(&mut self, _idx: usize) -> ::Item + where + Self: Iterator + UncheckedIndexedIterator, + { + unreachable!("Always specialized"); + } + + default unsafe fn get_unchecked_from_start(&mut self, _idx: usize) -> ::Item + where + Self: Iterator + UncheckedIndexedIterator, + { + unreachable!("Always specialized"); + } } +/* #[doc(hidden)] impl ZipImpl for Zip where @@ -411,6 +456,218 @@ where } } } +*/ + +impl Zip +where + A: UncheckedIndexedIterator + DoubleEndedIterator + ExactSizeIterator, + B: UncheckedIndexedIterator + DoubleEndedIterator + ExactSizeIterator, +{ + fn equalize_lengths(&mut self) { + if self.a_offset > 0 { + let _ = self.a.advance_back_by(self.a_offset); + self.a_offset = 0; + } + if self.b_offset > 0 { + let _ = self.b.advance_back_by(self.b_offset); + self.b_offset = 0; + } + } +} + +#[doc(hidden)] +impl ZipImpl for Zip +where + A: UncheckedIndexedIterator + Iterator, + B: UncheckedIndexedIterator + Iterator, +{ + fn new(a: A, b: B) -> Self { + let a_len = a.size_hint().0; + let b_len = b.size_hint().0; + let len = a_len.min(b_len); + Zip { + a, + b, + remaining: len, + a_offset: a_len.saturating_sub(len), + b_offset: b_len.saturating_sub(len), + } + } + + #[inline] + fn next(&mut self) -> Option { + let remaining = self.remaining; + if remaining > 0 { + self.remaining = remaining - 1; + + // SAFETY: ... + unsafe { + let idx_a = remaining + self.a_offset; + let idx_b = remaining + self.b_offset; + + let guard = + IndexGuardForward { it: &mut self.a, new_len: idx_a - 1, old_len: idx_a }; + let a = guard.it.index_from_end_unchecked(idx_a); + drop(guard); + let guard = + IndexGuardForward { it: &mut self.b, new_len: idx_b - 1, old_len: idx_b }; + let b = guard.it.index_from_end_unchecked(idx_b); + drop(guard); + + return Some((a, b)); + } + } + + None + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) + } + + #[inline] + fn nth(&mut self, n: usize) -> Option { + let to_skip = self.remaining.min(n); + self.remaining -= to_skip; + // SAFETY: UncheckedIndexedIterator contract requires that the length + // and thus the remaining count is correct. So we are sure + // we can skip this amount. + unsafe { + if A::MAY_HAVE_SIDE_EFFECT { + assert_unchecked(self.a.advance_by(to_skip).is_ok()); + } + if B::MAY_HAVE_SIDE_EFFECT { + assert_unchecked(self.b.advance_by(to_skip).is_ok()); + } + } + + ZipImpl::next(self) + } + + #[inline] + fn next_back(&mut self) -> Option + where + A: DoubleEndedIterator + ExactSizeIterator, + B: DoubleEndedIterator + ExactSizeIterator, + { + if Self::MAY_HAVE_SIDE_EFFECT && self.a_offset != self.b_offset { + self.equalize_lengths() + } + + let remaining = self.remaining; + if remaining > 0 { + self.remaining = remaining - 1; + + // SAFETY: ... + // don't have to use offsets here, equalized_lengths will have taken care of that. + unsafe { + let idx = remaining - 1; + let guard = IndexGuardBack { it: &mut self.a, new_len: idx, old_len: remaining }; + let a = guard.it.index_from_start_unchecked(idx); + drop(guard); + let guard = IndexGuardBack { it: &mut self.b, new_len: idx, old_len: remaining }; + let b = guard.it.index_from_start_unchecked(idx); + + return Some((a, b)); + } + } + + None + } + + fn fold(mut self, init: Acc, mut f: F) -> Acc + where + F: FnMut(Acc, Self::Item) -> Acc, + { + struct Guard<'a, A, B> + where + A: UncheckedIndexedIterator + Iterator, + B: UncheckedIndexedIterator + Iterator, + { + zip: &'a mut Zip, + old_len: usize, + remaining_a: usize, + remaining_b: usize, + } + + impl Drop for Guard<'_, A, B> + where + A: UncheckedIndexedIterator + Iterator, + B: UncheckedIndexedIterator + Iterator, + { + fn drop(&mut self) { + let offset_a = self.zip.a_offset; + let offset_b = self.zip.b_offset; + // SAFETY: ... + unsafe { + self.zip.a.set_front_index_from_end_unchecked( + self.remaining_a + offset_a, + self.old_len + offset_a, + ); + self.zip.b.set_front_index_from_end_unchecked( + self.remaining_b + offset_b, + self.old_len + offset_b, + ); + } + } + } + + let mut acc = init; + let len = self.remaining; + let mut guard = Guard { zip: &mut self, old_len: len, remaining_a: len, remaining_b: len }; + + for i in 0..len { + // SAFETY: ... + let (a, b) = unsafe { + guard.remaining_a = len - i - 1; + let a = guard.zip.a.index_from_end_unchecked(len - i + guard.zip.a_offset); + guard.remaining_b = len - i - 1; + let b = guard.zip.b.index_from_end_unchecked(len - i + guard.zip.b_offset); + + (a, b) + }; + acc = f(acc, (a, b)); + } + + if !Self::CLEANUP_ON_DROP { + mem::forget(guard); + } + + acc + } + + unsafe fn get_unchecked(&mut self, _idx: usize) -> ::Item + where + Self: Iterator + TrustedRandomAccessNoCoerce, + { + unreachable!("Wrong specialization") + } + + #[inline] + unsafe fn get_unchecked_from_end(&mut self, idx: usize) -> ::Item { + // SAFETY: the offsets represent how much longer each inner iterator is + // compared to the length we told the caller. So we adjust the index accordingly. + // Other than that the caller must uphold the same preconditions. + unsafe { + let a = self.a.index_from_end_unchecked(idx + self.a_offset); + let b = self.b.index_from_end_unchecked(idx + self.b_offset); + + (a, b) + } + } + + #[inline] + unsafe fn get_unchecked_from_start(&mut self, idx: usize) -> ::Item { + // SAFETY: forwarding to unsafe function with the same preconditions + unsafe { + let a = self.a.index_from_start_unchecked(idx); + let b = self.b.index_from_start_unchecked(idx); + + (a, b) + } + } +} #[stable(feature = "rust1", since = "1.0.0")] impl ExactSizeIterator for Zip @@ -420,6 +677,7 @@ where { } +/* #[doc(hidden)] #[unstable(feature = "trusted_random_access", issue = "none")] unsafe impl TrustedRandomAccess for Zip @@ -438,6 +696,53 @@ where { const MAY_HAVE_SIDE_EFFECT: bool = A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT; } +*/ + +#[unstable(feature = "trusted_indexed_access", issue = "none")] +impl UncheckedIndexedIterator for Zip +where + A: UncheckedIndexedIterator, + B: UncheckedIndexedIterator, +{ + const MAY_HAVE_SIDE_EFFECT: bool = A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT; + const CLEANUP_ON_DROP: bool = A::CLEANUP_ON_DROP || B::CLEANUP_ON_DROP; + + #[inline] + unsafe fn set_front_index_from_end_unchecked(&mut self, new_len: usize, old_len: usize) { + debug_assert_eq!(self.remaining, old_len); + self.remaining = new_len; + // SAFETY: the offsets represent how much longer each inner iterator is + // compared to the length we told the caller. So we adjust the index accordingly. + // Other than that the caller must uphold the same preconditions. + unsafe { + self.a.set_front_index_from_end_unchecked( + new_len + self.a_offset, + old_len + self.a_offset, + ); + self.b.set_front_index_from_end_unchecked( + new_len + self.b_offset, + old_len + self.b_offset, + ); + } + } + + #[inline] + unsafe fn set_end_index_from_start_unchecked(&mut self, new_len: usize, old_len: usize) { + debug_assert_eq!(self.remaining, old_len); + self.remaining = new_len; + // SAFETY: forwarding to unsafe function with the same preconditions + unsafe { + self.a.set_end_index_from_start_unchecked( + new_len + self.a_offset, + old_len + self.a_offset, + ); + self.b.set_end_index_from_start_unchecked( + new_len + self.b_offset, + old_len + self.b_offset, + ); + } + } +} #[stable(feature = "fused", since = "1.26.0")] impl FusedIterator for Zip @@ -692,3 +997,41 @@ impl SpecFold for Zip { accum } } + +struct IndexGuardForward<'a, I> +where + I: UncheckedIndexedIterator, +{ + new_len: usize, + old_len: usize, + it: &'a mut I, +} + +struct IndexGuardBack<'a, I> +where + I: UncheckedIndexedIterator, +{ + new_len: usize, + old_len: usize, + it: &'a mut I, +} + +impl Drop for IndexGuardForward<'_, I> { + #[inline] + fn drop(&mut self) { + // SAFETY: ... + unsafe { + self.it.set_front_index_from_end_unchecked(self.new_len, self.old_len); + } + } +} + +impl Drop for IndexGuardBack<'_, I> { + #[inline] + fn drop(&mut self) { + // SAFETY: ... + unsafe { + self.it.set_end_index_from_start_unchecked(self.new_len, self.old_len); + } + } +} diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs index 44fef3e145b78..bdf2156486975 100644 --- a/library/core/src/iter/mod.rs +++ b/library/core/src/iter/mod.rs @@ -461,6 +461,8 @@ pub use self::adapters::{ pub use self::adapters::{Intersperse, IntersperseWith}; pub(crate) use self::adapters::try_process; +#[unstable(feature = "trusted_indexed_access", issue = "none")] +pub use self::traits::UncheckedIndexedIterator; pub(crate) use self::traits::UncheckedIterator; mod adapters; diff --git a/library/core/src/iter/range.rs b/library/core/src/iter/range.rs index 68937161e046a..1cd4757a27a73 100644 --- a/library/core/src/iter/range.rs +++ b/library/core/src/iter/range.rs @@ -1,5 +1,6 @@ use crate::ascii::Char as AsciiChar; use crate::convert::TryFrom; +use crate::iter::UncheckedIndexedIterator; use crate::mem; use crate::net::{Ipv4Addr, Ipv6Addr}; use crate::num::NonZero; @@ -612,6 +613,25 @@ macro_rules! unsafe_range_trusted_random_access_impl { unsafe impl TrustedRandomAccessNoCoerce for ops::Range<$t> { const MAY_HAVE_SIDE_EFFECT: bool = false; } + + #[unstable(feature = "trusted_indexed_access", issue = "none")] + impl UncheckedIndexedIterator for ops::Range<$t> { + // all impls are Copy + const MAY_HAVE_SIDE_EFFECT: bool = false; + const CLEANUP_ON_DROP: bool = false; + + #[inline] + unsafe fn set_front_index_from_end_unchecked(&mut self, new_len: usize, _old_len: usize) { + // SAFETY: ... + self.start = unsafe { Step::backward_unchecked(self.end.clone(), new_len) } + } + + #[inline] + unsafe fn set_end_index_from_start_unchecked(&mut self, new_len: usize, _old_len: usize) { + // SAFETY: ... + self.start = unsafe { Step::forward_unchecked(self.start.clone(), new_len) } + } + } )*) } @@ -886,6 +906,26 @@ impl Iterator for ops::Range { // which means even repeated reads of the same index would be safe. unsafe { Step::forward_unchecked(self.start.clone(), idx) } } + + #[inline] + unsafe fn index_from_end_unchecked(&mut self, idx: usize) -> Self::Item + where + Self: UncheckedIndexedIterator, + { + // SAFETY: The TrustedRandomAccess contract requires that callers only pass an index + // that is in bounds. + unsafe { Step::backward_unchecked(self.end.clone(), idx) } + } + + #[inline] + unsafe fn index_from_start_unchecked(&mut self, idx: usize) -> Self::Item + where + Self: UncheckedIndexedIterator, + { + // SAFETY: The TrustedRandomAccess contract requires that callers only pass an index + // that is in bounds. + unsafe { Step::forward_unchecked(self.start.clone(), idx) } + } } // These macros generate `ExactSizeIterator` impls for various range types. diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs index e1904ed220cb4..383ce0b4942e5 100644 --- a/library/core/src/iter/traits/iterator.rs +++ b/library/core/src/iter/traits/iterator.rs @@ -13,6 +13,7 @@ use super::super::{ Inspect, Map, MapWhile, MapWindows, Peekable, Rev, Scan, Skip, SkipWhile, StepBy, Take, TakeWhile, }; +use super::UncheckedIndexedIterator; fn _assert_is_object_safe(_: &dyn Iterator) {} @@ -4056,6 +4057,24 @@ pub trait Iterator { { unreachable!("Always specialized"); } + + #[doc(hidden)] + #[unstable(feature = "trusted_indexed_access", issue = "none")] + unsafe fn index_from_end_unchecked(&mut self, _idx: usize) -> Self::Item + where + Self: UncheckedIndexedIterator, + { + unreachable!("Always specialized"); + } + + #[doc(hidden)] + #[unstable(feature = "trusted_indexed_access", issue = "none")] + unsafe fn index_from_start_unchecked(&mut self, _idx: usize) -> Self::Item + where + Self: UncheckedIndexedIterator, + { + unreachable!("Always specialized"); + } } /// Compares two iterators element-wise using the given function. diff --git a/library/core/src/iter/traits/mod.rs b/library/core/src/iter/traits/mod.rs index d4c9cc4b16037..2dffd931e67b9 100644 --- a/library/core/src/iter/traits/mod.rs +++ b/library/core/src/iter/traits/mod.rs @@ -23,4 +23,8 @@ pub use self::marker::TrustedFused; #[unstable(feature = "trusted_step", issue = "85731")] pub use self::marker::TrustedStep; +#[unstable(feature = "trusted_indexed_access", issue = "none")] +pub use self::unchecked_iterator::UncheckedIndexedIterator; pub(crate) use self::unchecked_iterator::UncheckedIterator; + +pub(in crate::iter) use self::unchecked_iterator::SpecIndexedAccess; diff --git a/library/core/src/iter/traits/unchecked_iterator.rs b/library/core/src/iter/traits/unchecked_iterator.rs index ae4bfcad4e68f..add98d9bf7a54 100644 --- a/library/core/src/iter/traits/unchecked_iterator.rs +++ b/library/core/src/iter/traits/unchecked_iterator.rs @@ -34,3 +34,74 @@ pub(crate) trait UncheckedIterator: TrustedLen { unsafe { opt.unwrap_unchecked() } } } + +/// Specialization trait for unchecked iterator indexing +/// +/// # Safety and Use +/// +/// * `size_hint` must be exact +/// * the size must fit into an `usize` +/// * each position must only be accessed once +/// * no other iterator methods must be called between getting the size, +/// calling an unchecked getter one or more times and +/// finally comitting the index updates with an unchecked setter. +/// * changing iteration direction requires comitting the index updates +/// from the previous iteration direction +/// +/// The start-end range represents inclusive-exclusive bounds. +/// Meaning that the end index 0 represents a positon one-past the current +/// range of valid items. +/// While start index 0 represents the first item if the iterator is non-empty. +/// +/// +#[doc(hidden)] +#[unstable(feature = "trusted_indexed_access", issue = "none")] +#[rustc_specialization_trait] +pub trait UncheckedIndexedIterator: Sized { + /// `true` if getting an iterator element may have side effects. + /// Remember to take inner iterators into account. + /// `true` is the conservative choice. + const MAY_HAVE_SIDE_EFFECT: bool; + + /// `true` if updating interator indexes is needed + /// even if the iterator will no longer be used, e.g. to avoid double-drops or leaks. + /// `trus` is the conservative choice. + const CLEANUP_ON_DROP: bool; + + unsafe fn set_front_index_from_end_unchecked(&mut self, _new_len: usize, _old_len: usize); + unsafe fn set_end_index_from_start_unchecked(&mut self, _new_len: usize, _old_len: usize); +} + +pub(in crate::iter) unsafe trait SpecIndexedAccess: Iterator { + /// If `Self: TrustedRandomAccess`, it must be safe to call + /// `Iterator::__iterator_get_unchecked(self, index)`. + unsafe fn index_from_end_unchecked_inner(&mut self, index: usize) -> Self::Item; + + unsafe fn index_from_start_unchecked_inner(&mut self, index: usize) -> Self::Item; +} + +unsafe impl SpecIndexedAccess for I { + default unsafe fn index_from_end_unchecked_inner(&mut self, _: usize) -> Self::Item { + unreachable!("Should only be called on UncheckedIndexedIterator iterators"); + } + + default unsafe fn index_from_start_unchecked_inner(&mut self, _: usize) -> Self::Item { + unreachable!("Should only be called on UncheckedIndexedIterator iterators"); + } +} + +unsafe impl SpecIndexedAccess for I { + #[inline] + unsafe fn index_from_end_unchecked_inner(&mut self, idx: usize) -> Self::Item { + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + unsafe { self.index_from_end_unchecked(idx) } + } + + #[inline] + unsafe fn index_from_start_unchecked_inner(&mut self, idx: usize) -> Self::Item { + // SAFETY: the caller must uphold the contract for + // `Iterator::__iterator_get_unchecked`. + unsafe { self.index_from_start_unchecked(idx) } + } +} diff --git a/library/core/src/ops/index_range.rs b/library/core/src/ops/index_range.rs index 07ea2e930d57a..26425d2518afe 100644 --- a/library/core/src/ops/index_range.rs +++ b/library/core/src/ops/index_range.rs @@ -108,6 +108,22 @@ impl IndexRange { self.end = mid; suffix } + + /// # Safety + /// + /// `start` must be `<= end`. + #[inline] + pub unsafe fn set_start_unchecked(&mut self, start: usize) { + self.start = start; + } + + /// # Safety + /// + /// `end` must be ``>= start`. + #[inline] + pub unsafe fn set_end_unchecked(&mut self, end: usize) { + self.end = end; + } } impl Iterator for IndexRange { diff --git a/library/core/src/slice/iter.rs b/library/core/src/slice/iter.rs index d7d4f90c1a538..6aac43cf2a0ef 100644 --- a/library/core/src/slice/iter.rs +++ b/library/core/src/slice/iter.rs @@ -7,7 +7,8 @@ use crate::cmp; use crate::fmt; use crate::hint::assert_unchecked; use crate::iter::{ - FusedIterator, TrustedLen, TrustedRandomAccess, TrustedRandomAccessNoCoerce, UncheckedIterator, + FusedIterator, TrustedLen, TrustedRandomAccess, TrustedRandomAccessNoCoerce, + UncheckedIndexedIterator, UncheckedIterator, }; use crate::marker::PhantomData; use crate::mem::{self, SizedTypeProperties}; diff --git a/library/core/src/slice/iter/macros.rs b/library/core/src/slice/iter/macros.rs index 7910981d0f5ee..994ec6541d322 100644 --- a/library/core/src/slice/iter/macros.rs +++ b/library/core/src/slice/iter/macros.rs @@ -383,9 +383,63 @@ macro_rules! iterator { unsafe { & $( $mut_ )? * self.ptr.as_ptr().add(idx) } } + #[inline] + unsafe fn index_from_end_unchecked(&mut self, idx: usize) -> Self::Item { + if T::IS_ZST { + // SAFETY: for ZSTs we can always conjure references + unsafe { NonNull::dangling().$into_ref() } + } else { + // SAFETY: for non-ZSTs it's the field's invariant that it must be non-null + let end = unsafe { *ptr::addr_of!(self.end_or_len).cast::>() }; + // SAFETY: the caller promises that the index is in bounds + unsafe { & $( $mut_ )? *end.sub(idx).as_ptr() } + } + } + + #[inline] + unsafe fn index_from_start_unchecked(&mut self, idx: usize) -> Self::Item { + if T::IS_ZST { + // SAFETY: for ZSTs we can always conjure references + unsafe { NonNull::dangling().$into_ref() } + } else { + // SAFETY: the caller promises that the index is in bounds + unsafe { & $( $mut_ )? *self.ptr.add(idx).as_ptr() } + } + } + $($extra)* } + #[unstable(feature = "trusted_indexed_access", issue = "none")] + impl<'a, T> UncheckedIndexedIterator for $name<'a, T> { + const MAY_HAVE_SIDE_EFFECT: bool = false; + const CLEANUP_ON_DROP: bool = false; + + #[inline] + unsafe fn set_front_index_from_end_unchecked(&mut self, new_len: usize, _old_len: usize) { + if T::IS_ZST { + self.end_or_len = ptr::without_provenance_mut(new_len); + } else { + // SAFETY: for non-ZSTs it's the field's invariant that it must be non-null + let end = unsafe { *ptr::addr_of!(self.end_or_len).cast::>() }; + // SAFETY: the caller promises that the index is in bounds or one-past-the-end + 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_or_len = ptr::without_provenance_mut(new_len); + } else { + // SAFETY: for non-ZSTs it's the field's invariant that it must be non-null + let end = unsafe { &mut *ptr::addr_of_mut!(self.end_or_len).cast::>() }; + // SAFETY: the caller promises that the index is in bounds or one-past-the-end + *end = unsafe { self.ptr.add(new_len) }; + } + } + } + #[stable(feature = "rust1", since = "1.0.0")] impl<'a, T> DoubleEndedIterator for $name<'a, T> { #[inline] diff --git a/library/core/tests/iter/adapters/cloned.rs b/library/core/tests/iter/adapters/cloned.rs index 78babb7feab18..a17b6337febdf 100644 --- a/library/core/tests/iter/adapters/cloned.rs +++ b/library/core/tests/iter/adapters/cloned.rs @@ -31,7 +31,8 @@ fn test_cloned_side_effects() { .zip(&[1]); for _ in iter {} } - assert_eq!(count, 2); + // either result is permitted by the Iterator::zip contract + assert!([1, 2].contains(&count)); } #[test] diff --git a/library/core/tests/iter/adapters/zip.rs b/library/core/tests/iter/adapters/zip.rs index ba54de5822bf7..ab8c894477597 100644 --- a/library/core/tests/iter/adapters/zip.rs +++ b/library/core/tests/iter/adapters/zip.rs @@ -108,7 +108,8 @@ fn test_zip_next_back_side_effects_exhausted() { iter.next(); iter.next(); assert_eq!(iter.next_back(), None); - assert_eq!(a, vec![1, 2, 3, 4, 6, 5]); + // zip contract allows either outcome + assert!(a == [1, 2, 3, 4, 6, 5] || a == [1, 2, 3, 6, 5, 4]); assert_eq!(b, vec![200, 300, 400]); } @@ -119,7 +120,8 @@ fn test_zip_cloned_sideffectful() { for _ in xs.iter().cloned().zip(ys.iter().cloned()) {} - assert_eq!(&xs, &[1, 1, 1, 0][..]); + // zip contract allows either outcome + assert!(&xs == &[1, 1, 1, 0] || &xs == &[1, 1, 0, 0]); assert_eq!(&ys, &[1, 1][..]); let xs = [CountClone::new(), CountClone::new()]; @@ -138,7 +140,8 @@ fn test_zip_map_sideffectful() { for _ in xs.iter_mut().map(|x| *x += 1).zip(ys.iter_mut().map(|y| *y += 1)) {} - assert_eq!(&xs, &[1, 1, 1, 1, 1, 0]); + // zip contract allows either outcome + assert!(&xs == &[1, 1, 1, 1, 1, 0] || &xs == &[1, 1, 1, 1, 0, 0]); assert_eq!(&ys, &[1, 1, 1, 1]); let mut xs = [0; 4]; @@ -212,7 +215,8 @@ fn test_zip_nth_back_side_effects_exhausted() { iter.next(); iter.next(); assert_eq!(iter.nth_back(0), None); - assert_eq!(a, vec![1, 2, 3, 4, 6, 5]); + // zip contract allows either outcome + assert!(a == [1, 2, 3, 4, 6, 5] || a == [1, 2, 3, 6, 5, 4]); assert_eq!(b, vec![200, 300, 400]); } @@ -231,7 +235,7 @@ fn test_zip_trusted_random_access_composition() { assert_eq!(z1.next().unwrap(), (0, 0)); let mut z2 = z1.zip(c); - fn assert_trusted_random_access(_a: &T) {} + fn assert_trusted_random_access(_a: &T) {} assert_trusted_random_access(&z2); assert_eq!(z2.next().unwrap(), ((1, 1), 1)); } diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index fa0e9a979d060..ab576330212ae 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -107,6 +107,7 @@ #![feature(strict_provenance)] #![feature(strict_provenance_atomic_ptr)] #![feature(trusted_random_access)] +#![feature(trusted_indexed_access)] #![feature(unsize)] #![feature(const_array_from_ref)] #![feature(const_slice_from_ref)] diff --git a/tests/codegen/vec-in-place.rs b/tests/codegen/vec-in-place.rs index 7a175dc4f7e1c..2530f7869bbe4 100644 --- a/tests/codegen/vec-in-place.rs +++ b/tests/codegen/vec-in-place.rs @@ -57,6 +57,14 @@ pub fn vec_iterator_cast_unwrap(vec: Vec>) -> Vec { vec.into_iter().map(|e| e.0).collect() } +// CHECK-LABEL: @vec_iterator_cast_unwrap_droppable +#[no_mangle] +pub fn vec_iterator_cast_unwrap_droppable(vec: Vec>) -> Vec { + // CHECK-NOT: loop + // CHECK-NOT: call + vec.into_iter().map(|e| e.0).collect() +} + // CHECK-LABEL: @vec_iterator_cast_aggregate #[no_mangle] pub fn vec_iterator_cast_aggregate(vec: Vec<[u64; 4]>) -> Vec { From 29813544a60d3413469e45a346c3189ff5fae99c Mon Sep 17 00:00:00 2001 From: The 8472 Date: Thu, 7 Mar 2024 20:34:32 +0100 Subject: [PATCH 3/3] select functions in a const block --- library/alloc/src/vec/in_place_collect.rs | 172 +++++++++++----------- 1 file changed, 90 insertions(+), 82 deletions(-) diff --git a/library/alloc/src/vec/in_place_collect.rs b/library/alloc/src/vec/in_place_collect.rs index d04ac3eada0c4..b2a9220b202e4 100644 --- a/library/alloc/src/vec/in_place_collect.rs +++ b/library/alloc/src/vec/in_place_collect.rs @@ -231,97 +231,105 @@ where I: Iterator + InPlaceCollect, ::Source: AsVecIntoIter, { - default fn from_iter(mut iterator: I) -> Self { + default fn from_iter(iterator: I) -> Self { // See "Layout constraints" section in the module documentation. We rely on const // optimization here since these conditions currently cannot be expressed as trait bounds - if const { !in_place_collectible::(I::MERGE_BY, I::EXPAND_BY) } { - // fallback to more generic implementations - return SpecFromIterNested::from_iter(iterator); - } - - 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(), - inner.ptr, - inner.cap, - inner.buf.as_ptr() as *mut T, - inner.end as *const T, - inner.cap * mem::size_of::() / mem::size_of::(), - ) + let fun = const { + if !in_place_collectible::(I::MERGE_BY, I::EXPAND_BY) { + SpecFromIterNested::::from_iter + } else { + from_iter + } }; - // SAFETY: `dst_buf` and `dst_end` are the start and end of the buffer. - let len = unsafe { SpecInPlaceCollect::collect_in_place(&mut iterator, dst_buf, dst_end) }; - - let src = unsafe { iterator.as_inner().as_into_iter() }; - // 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 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(), - "InPlaceIterable contract violation, write pointer advanced beyond read pointer" - ); - } + fun(iterator) + } +} - // The ownership of the source allocation and the new `T` values is temporarily moved into `dst_guard`. - // This is safe because - // * `forget_allocation_drop_remaining` immediately forgets the allocation - // before any panic can occur in order to avoid any double free, and then proceeds to drop - // any remaining values at the tail of the source. - // * the shrink either panics without invalidating the allocation, aborts or - // succeeds. In the last case we disarm the guard. - // - // Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce - // contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the - // module documentation why this is ok anyway. - let dst_guard = - InPlaceDstDataSrcBufDrop { ptr: dst_buf, len, src_cap, src: PhantomData:: }; - src.forget_allocation_drop_remaining(); - - // Adjust the allocation if the source had a capacity in bytes that wasn't a multiple - // of the destination type size. - // Since the discrepancy should generally be small this should only result in some - // bookkeeping updates and no memmove. - if needs_realloc::(src_cap, dst_cap) { - let alloc = Global; - debug_assert_ne!(src_cap, 0); - debug_assert_ne!(dst_cap, 0); - unsafe { - // The old allocation exists, therefore it must have a valid layout. - let src_align = mem::align_of::(); - let src_size = mem::size_of::().unchecked_mul(src_cap); - let old_layout = Layout::from_size_align_unchecked(src_size, src_align); - - // The allocation must be equal or smaller for in-place iteration to be possible - // therefore the new layout must be ≤ the old one and therefore valid. - let dst_align = mem::align_of::(); - let dst_size = mem::size_of::().unchecked_mul(dst_cap); - let new_layout = Layout::from_size_align_unchecked(dst_size, dst_align); - - let result = alloc.shrink( - NonNull::new_unchecked(dst_buf as *mut u8), - old_layout, - new_layout, - ); - let Ok(reallocated) = result else { handle_alloc_error(new_layout) }; - dst_buf = reallocated.as_ptr() as *mut T; - } - } else { - debug_assert_eq!(src_cap * mem::size_of::(), dst_cap * mem::size_of::()); +fn from_iter(mut iterator: I) -> Vec +where + I: Iterator + InPlaceCollect, + ::Source: AsVecIntoIter, +{ + 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(), + inner.ptr, + inner.cap, + inner.buf.as_ptr() as *mut T, + inner.end as *const T, + inner.cap * mem::size_of::() / mem::size_of::(), + ) + }; + + // SAFETY: `dst_buf` and `dst_end` are the start and end of the buffer. + let len = unsafe { SpecInPlaceCollect::collect_in_place(&mut iterator, dst_buf, dst_end) }; + + let src = unsafe { iterator.as_inner().as_into_iter() }; + // 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 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(), + "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 + // * `forget_allocation_drop_remaining` immediately forgets the allocation + // before any panic can occur in order to avoid any double free, and then proceeds to drop + // any remaining values at the tail of the source. + // * the shrink either panics without invalidating the allocation, aborts or + // succeeds. In the last case we disarm the guard. + // + // Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce + // contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the + // module documentation why this is ok anyway. + let dst_guard = + InPlaceDstDataSrcBufDrop { ptr: dst_buf, len, src_cap, src: PhantomData:: }; + src.forget_allocation_drop_remaining(); + + // Adjust the allocation if the source had a capacity in bytes that wasn't a multiple + // of the destination type size. + // Since the discrepancy should generally be small this should only result in some + // bookkeeping updates and no memmove. + if needs_realloc::(src_cap, dst_cap) { + let alloc = Global; + debug_assert_ne!(src_cap, 0); + debug_assert_ne!(dst_cap, 0); + unsafe { + // The old allocation exists, therefore it must have a valid layout. + let src_align = mem::align_of::(); + let src_size = mem::size_of::().unchecked_mul(src_cap); + let old_layout = Layout::from_size_align_unchecked(src_size, src_align); + + // The allocation must be equal or smaller for in-place iteration to be possible + // therefore the new layout must be ≤ the old one and therefore valid. + let dst_align = mem::align_of::(); + let dst_size = mem::size_of::().unchecked_mul(dst_cap); + let new_layout = Layout::from_size_align_unchecked(dst_size, dst_align); + + let result = + alloc.shrink(NonNull::new_unchecked(dst_buf as *mut u8), old_layout, new_layout); + let Ok(reallocated) = result else { handle_alloc_error(new_layout) }; + dst_buf = reallocated.as_ptr() as *mut T; } + } else { + debug_assert_eq!(src_cap * mem::size_of::(), dst_cap * mem::size_of::()); + } - mem::forget(dst_guard); + mem::forget(dst_guard); - let vec = unsafe { Vec::from_raw_parts(dst_buf, len, dst_cap) }; + let vec = unsafe { Vec::from_raw_parts(dst_buf, len, dst_cap) }; - vec - } + vec } fn write_in_place_with_drop(