diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index f66dc648f809b..95dc40aa04a83 100644 --- a/library/alloc/src/lib.rs +++ b/library/alloc/src/lib.rs @@ -136,6 +136,7 @@ #![feature(legacy_receiver_trait)] #![feature(likely_unlikely)] #![feature(local_waker)] +#![feature(maybe_uninit_array_assume_init)] #![feature(maybe_uninit_uninit_array_transpose)] #![feature(panic_internals)] #![feature(pattern)] diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs index 97e4912596b94..b7fcecf11e059 100644 --- a/library/alloc/src/vec/into_iter.rs +++ b/library/alloc/src/vec/into_iter.rs @@ -445,6 +445,47 @@ impl DoubleEndedIterator for IntoIter { } } + #[inline] + fn next_chunk_back(&mut self) -> Result<[T; N], core::array::IntoIter> { + let mut raw_ary = [const { MaybeUninit::uninit() }; N]; + + let len = self.len(); + + if T::IS_ZST { + if len < N { + self.forget_remaining_elements(); + // Safety: ZSTs can be conjured ex nihilo, only the amount has to be correct + return Err(unsafe { array::IntoIter::new_unchecked(raw_ary, N - len..N) }); + } + + self.end = self.end.wrapping_byte_sub(N); + // Safety: ditto + return Ok(unsafe { MaybeUninit::array_assume_init(raw_ary) }); + } + + if len < N { + // Safety: `len` indicates that this many elements are available and we just checked that + // it fits into the array. + unsafe { + ptr::copy_nonoverlapping(self.ptr.as_ptr(), raw_ary.as_mut_ptr() as *mut T, len); + self.forget_remaining_elements(); + return Err(array::IntoIter::new_unchecked(raw_ary, 0..len)); + } + } + + // Safety: `len` is larger than the array size. Copy a fixed amount here to fully initialize + // the array. + unsafe { + ptr::copy_nonoverlapping( + self.ptr.add(len - N).as_ptr(), + raw_ary.as_mut_ptr() as *mut T, + N, + ); + self.end = self.end.sub(N); + Ok(MaybeUninit::array_assume_init(raw_ary)) + } + } + #[inline] fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero> { let step_size = self.len().min(n); diff --git a/library/alloctests/tests/vec.rs b/library/alloctests/tests/vec.rs index ccf9b8095b290..fc58e4364fe66 100644 --- a/library/alloctests/tests/vec.rs +++ b/library/alloctests/tests/vec.rs @@ -1027,6 +1027,15 @@ fn test_into_iter_next_chunk() { assert_eq!(iter.next_chunk::<4>().unwrap_err().as_slice(), &[]); // N is explicitly 4 } +#[test] +fn test_into_iter_next_chunk_back() { + let mut iter = b"lorem".to_vec().into_iter(); + + assert_eq!(iter.next_chunk_back().unwrap(), [b'e', b'm']); // N is inferred as 2 + assert_eq!(iter.next_chunk_back().unwrap(), [b'l', b'o', b'r']); // N is inferred as 3 + assert_eq!(iter.next_chunk_back::<4>().unwrap_err().as_slice(), &[]); // N is explicitly 4 +} + #[test] fn test_into_iter_clone() { fn iter_equal>(it: I, slice: &[i32]) { @@ -1131,6 +1140,14 @@ fn test_into_iter_zst() { let mut it = vec![C, C].into_iter(); it.next_chunk::<4>().unwrap_err(); drop(it); + + let mut it = vec![C, C].into_iter(); + it.next_chunk_back::<1>().unwrap(); + drop(it); + + let mut it = vec![C, C].into_iter(); + it.next_chunk_back::<4>().unwrap_err(); + drop(it); } #[test] diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs index 5784dbed03634..63a99349d7c9a 100644 --- a/library/core/src/array/mod.rs +++ b/library/core/src/array/mod.rs @@ -969,6 +969,54 @@ const impl Drop for Guard<'_, T> { } } +/// Panic guard for incremental initialization of arrays from the back. +/// +/// Elements of the array are populated starting from the end towards the beginning. +/// Disarm the guard with `mem::forget` once the array has been fully initialized. +/// +/// # Safety +/// +/// All write accesses to this structure are unsafe and must maintain a correct +/// count of `initialized` elements. +struct GuardBack<'a, T> { + /// The array to be initialized (will be filled from the end). + pub array_mut: &'a mut [MaybeUninit], + /// The number of items that have been initialized so far. + pub initialized: usize, +} + +impl GuardBack<'_, T> { + /// Adds an item to the array and updates the initialized item counter. + /// + /// # Safety + /// + /// No more than N elements must be initialized. + #[inline] + pub(crate) unsafe fn push_unchecked(&mut self, item: T) { + // SAFETY: If `initialized` was correct before and the caller does not + // invoke this method more than N times, then writes will be in-bounds + // and slots will not be initialized more than once. + unsafe { + let offset = self.initialized.unchecked_add(1); + let index = self.array_mut.len().unchecked_sub(offset); + self.array_mut.get_unchecked_mut(index).write(item); + self.initialized = offset; + } + } +} + +impl Drop for GuardBack<'_, T> { + #[inline] + fn drop(&mut self) { + debug_assert!(self.initialized <= self.array_mut.len()); + let len = self.array_mut.len(); + // SAFETY: this slice will contain only initialized objects. + unsafe { + self.array_mut.get_unchecked_mut(len - self.initialized..len).assume_init_drop(); + } + } +} + /// Pulls `N` items from `iter` and returns them as an array. If the iterator /// yields fewer than `N` items, `Err` is returned containing an iterator over /// the already yielded items. @@ -1076,3 +1124,109 @@ const fn iter_next_chunk_erased( mem::forget(guard); Ok(()) } + +/// Pulls `N` items from the back of `iter` and returns them as an array. +/// If the iterator yields fewer than `N` items, `Err` is returned containing +/// an iterator over the already yielded items. +/// +/// Since the iterator is passed as a mutable reference and this function calls +/// `next_back` at most `N` times, the iterator can still be used afterwards to +/// retrieve the remaining items. +/// +/// If `iter.next_back()` panics, all items already yielded by the iterator are +/// dropped. +/// +/// Used for [`DoubleEndedIterator::next_chunk_back`]. +#[inline] +pub(crate) fn iter_next_chunk_back( + iter: &mut impl DoubleEndedIterator, +) -> Result<[T; N], IntoIter> { + iter.spec_next_chunk_back() +} + +pub(crate) trait SpecNextChunkBack: + DoubleEndedIterator +{ + fn spec_next_chunk_back(&mut self) -> Result<[T; N], IntoIter>; +} + +impl, T, const N: usize> SpecNextChunkBack for I { + #[inline] + default fn spec_next_chunk_back(&mut self) -> Result<[T; N], IntoIter> { + let mut array = [const { MaybeUninit::uninit() }; N]; + let r = iter_next_chunk_back_erased(&mut array, self); + match r { + Ok(()) => { + // SAFETY: All elements of `array` were populated. + Ok(unsafe { MaybeUninit::array_assume_init(array) }) + } + Err(initialized) => { + // SAFETY: Only the last `initialized` elements were populated + Err(unsafe { IntoIter::new_unchecked(array, N - initialized..N) }) + } + } + } +} + +impl + TrustedLen, T, const N: usize> SpecNextChunkBack + for I +{ + fn spec_next_chunk_back(&mut self) -> Result<[T; N], IntoIter> { + let len = (*self).size_hint().0; + let mut array = [const { MaybeUninit::uninit() }; N]; + if len < N { + // SAFETY: `TrustedLen`, an unsafe trait, requires that i can get len items out of it. + unsafe { write_back(&mut array, self, len) }; + // SAFETY: Only the last `len` elements were populated + Err(unsafe { IntoIter::new_unchecked(array, N - len..N) }) + } else { + // SAFETY: `TrustedLen`, an unsafe trait, requires that i can get N items out of it. + unsafe { write_back(&mut array, self, N) }; + // SAFETY: All N items were populated + Ok(unsafe { MaybeUninit::array_assume_init(array) }) + } + } +} + +// SAFETY: `from` must have len items, and len items must be < N. +unsafe fn write_back( + to: &mut [MaybeUninit; N], + from: &mut impl DoubleEndedIterator, + len: usize, +) { + let mut guard = GuardBack { array_mut: to, initialized: 0 }; + while guard.initialized < len { + // SAFETY: caller has guaranteed, from has len items. + let item = unsafe { from.next_back().unwrap_unchecked() }; + // SAFETY: guard.initialized < len < N + unsafe { guard.push_unchecked(item) }; + } + crate::mem::forget(guard); +} + +/// Version of [`iter_next_chunk_back`] using a passed-in slice +/// in order to avoid needing to monomorphize for every array length. +/// +/// Unfortunately this loop has two exit conditions, the buffer filling up +/// or the iterator running out of items, making it tend to optimize poorly. +#[inline] +fn iter_next_chunk_back_erased( + buffer: &mut [MaybeUninit], + iter: &mut impl DoubleEndedIterator, +) -> Result<(), usize> { + // if `Iterator::next_back` panics, this guard will drop already initialized items + let mut guard = GuardBack { array_mut: buffer, initialized: 0 }; + while guard.initialized < guard.array_mut.len() { + let Some(item) = iter.next_back() else { + let initialized = guard.initialized; + mem::forget(guard); + return Err(initialized); + }; + + // SAFETY: The loop condition ensures we have space to push the item + unsafe { guard.push_unchecked(item) }; + } + + mem::forget(guard); + Ok(()) +} diff --git a/library/core/src/iter/traits/double_ended.rs b/library/core/src/iter/traits/double_ended.rs index ce80874a23897..a16732e1d8c74 100644 --- a/library/core/src/iter/traits/double_ended.rs +++ b/library/core/src/iter/traits/double_ended.rs @@ -1,3 +1,4 @@ +use crate::array; use crate::marker::Destruct; use crate::num::NonZero; use crate::ops::{ControlFlow, Try}; @@ -95,6 +96,53 @@ pub const trait DoubleEndedIterator: [const] Iterator { #[stable(feature = "rust1", since = "1.0.0")] fn next_back(&mut self) -> Option; + /// Advances from the back of the iterator and returns an array containing the next + /// `N` values in sequence. + /// + /// If there are not enough elements to fill the array then `Err` is returned + /// containing an iterator over the remaining elements. + /// + /// Note: This is not equivalent to doing `iter.rev().next_chunk()` as this method + /// takes elements from the back of the iterator and preserves the order that the + /// elements were seen in the original iterator. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// #![feature(iter_next_chunk)] + /// + /// let mut iter = "lorem".chars(); + /// + /// assert_eq!(iter.next_chunk_back().unwrap(), ['e', 'm']); // N is inferred as 2 + /// assert_eq!(iter.next_chunk_back().unwrap(), ['l', 'o', 'r']); // N is inferred as 3 + /// assert_eq!(iter.next_chunk_back::<4>().unwrap_err().as_slice(), &[]); // N is explicitly 4 + /// ``` + /// + /// Split a string and get the last three items in sequence. + /// + /// ``` + /// #![feature(iter_next_chunk)] + /// + /// let quote = "not all those who wander are lost"; + /// let [first, second, third] = quote.split_whitespace().next_chunk_back().unwrap(); + /// assert_eq!(first, "wander"); + /// assert_eq!(second, "are"); + /// assert_eq!(third, "lost"); + /// ``` + #[inline] + #[unstable(feature = "iter_next_chunk", issue = "98326")] + #[rustc_non_const_trait_method] + fn next_chunk_back( + &mut self, + ) -> Result<[Self::Item; N], array::IntoIter> + where + Self: Sized, + { + crate::array::iter_next_chunk_back(self) + } + /// Advances the iterator from the back by `n` elements. /// /// `advance_back_by` is the reverse version of [`advance_by`]. This method will diff --git a/library/coretests/tests/iter/adapters/filter.rs b/library/coretests/tests/iter/adapters/filter.rs index 167851e33336e..14cd3cbd58f10 100644 --- a/library/coretests/tests/iter/adapters/filter.rs +++ b/library/coretests/tests/iter/adapters/filter.rs @@ -63,3 +63,15 @@ fn test_next_chunk_does_not_leak() { assert_eq!(Rc::strong_count(w), 1); } } + +#[test] +fn test_next_chunk_back_does_not_leak() { + let drop_witness: [_; 5] = std::array::from_fn(|_| Rc::new(())); + + let v = (0..5).map(|i| drop_witness[i].clone()).collect::>(); + let _ = v.into_iter().filter(|_| false).next_chunk_back::<1>(); + + for ref w in drop_witness { + assert_eq!(Rc::strong_count(w), 1); + } +} diff --git a/library/coretests/tests/iter/traits/iterator.rs b/library/coretests/tests/iter/traits/iterator.rs index 386946461e9fe..0850b1e4edc26 100644 --- a/library/coretests/tests/iter/traits/iterator.rs +++ b/library/coretests/tests/iter/traits/iterator.rs @@ -619,6 +619,18 @@ fn test_next_chunk() { assert_eq!(it.next_chunk::<0>().unwrap(), []); } +#[test] +fn test_next_chunk_back() { + let mut it = 0..12; + assert_eq!(it.next_chunk_back().unwrap(), [8, 9, 10, 11]); + assert_eq!(it.next_chunk_back().unwrap(), []); + assert_eq!(it.next_chunk_back().unwrap(), [2, 3, 4, 5, 6, 7]); + assert_eq!(it.next_chunk_back::<4>().unwrap_err().as_slice(), &[0, 1]); + + let mut it = std::iter::once_with(|| panic!()); + assert_eq!(it.next_chunk_back::<0>().unwrap(), []); +} + #[test] fn test_next_chunk_untrusted() { struct Untrusted(I); @@ -636,6 +648,28 @@ fn test_next_chunk_untrusted() { assert_eq!(it.next_chunk::<4>().unwrap_err().as_slice(), &[10, 11]); } +#[test] +fn test_next_chunk_back_untrusted() { + struct Untrusted(I); + impl Iterator for Untrusted { + type Item = I::Item; + + fn next(&mut self) -> Option { + self.0.next() + } + } + impl DoubleEndedIterator for Untrusted { + fn next_back(&mut self) -> Option { + self.0.next_back() + } + } + let mut it = Untrusted(0..12); + assert_eq!(it.next_chunk_back().unwrap(), [8, 9, 10, 11]); + assert_eq!(it.next_chunk_back().unwrap(), []); + assert_eq!(it.next_chunk_back().unwrap(), [2, 3, 4, 5, 6, 7]); + assert_eq!(it.next_chunk::<4>().unwrap_err().as_slice(), &[0, 1]); +} + #[test] fn test_collect_into_tuples() { let a = vec![(1, 2, 3), (4, 5, 6), (7, 8, 9)];