diff --git a/library/alloc/src/vec/drain.rs b/library/alloc/src/vec/drain.rs index fb32d144f872c..08cdf213ce810 100644 --- a/library/alloc/src/vec/drain.rs +++ b/library/alloc/src/vec/drain.rs @@ -1,7 +1,7 @@ use crate::alloc::{Allocator, Global}; use core::fmt; use core::iter::{FusedIterator, TrustedLen}; -use core::mem::{self}; +use core::mem::{self, MaybeUninit}; use core::ptr::{self, NonNull}; use core::slice::{self}; @@ -102,16 +102,11 @@ impl DoubleEndedIterator for Drain<'_, T, A> { #[stable(feature = "drain", since = "1.6.0")] impl Drop for Drain<'_, T, A> { fn drop(&mut self) { - /// Continues dropping the remaining elements in the `Drain`, then moves back the - /// un-`Drain`ed elements to restore the original `Vec`. + /// Moves back the un-`Drain`ed elements to restore the original `Vec`. struct DropGuard<'r, 'a, T, A: Allocator>(&'r mut Drain<'a, T, A>); impl<'r, 'a, T, A: Allocator> Drop for DropGuard<'r, 'a, T, A> { fn drop(&mut self) { - // Continue the same loop we have below. If the loop already finished, this does - // nothing. - self.0.for_each(drop); - if self.0.tail_len > 0 { unsafe { let source_vec = self.0.vec.as_mut(); @@ -129,15 +124,43 @@ impl Drop for Drain<'_, T, A> { } } - // exhaust self first - while let Some(item) = self.next() { - let guard = DropGuard(self); - drop(item); - mem::forget(guard); + let iter = mem::replace(&mut self.iter, (&mut []).iter()); + let drop_len = iter.len(); + let drop_ptr = iter.as_slice().as_ptr(); + + // forget iter so there's no aliasing reference + drop(iter); + + let mut vec = self.vec; + + if mem::size_of::() == 0 { + // ZSTs have no identity, so we don't need to move them around, we only need to drop the correct amount. + // this can be achieved by manipulating the Vec length instead of moving values out from `iter`. + unsafe { + let vec = vec.as_mut(); + let old_len = vec.len(); + vec.set_len(old_len + drop_len + self.tail_len); + vec.truncate(old_len + self.tail_len); + } + + return; + } + + // ensure elements are moved back into their appropriate places, even when drop_in_place panics + let _guard = DropGuard(self); + + if drop_len == 0 { + return; } - // Drop a `DropGuard` to move back the non-drained tail of `self`. - DropGuard(self); + unsafe { + let vec = vec.as_mut(); + let spare_capacity = vec.spare_capacity_mut(); + let drop_offset = drop_ptr.offset_from(spare_capacity.as_ptr() as *const _) as usize; + let drop_range = drop_offset..(drop_offset + drop_len); + let to_drop = &mut spare_capacity[drop_range]; + ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(to_drop)); + } } } diff --git a/library/core/src/iter/traits/double_ended.rs b/library/core/src/iter/traits/double_ended.rs index 6f8cb6b5a65b6..ac2501928f82d 100644 --- a/library/core/src/iter/traits/double_ended.rs +++ b/library/core/src/iter/traits/double_ended.rs @@ -214,17 +214,13 @@ pub trait DoubleEndedIterator: Iterator { /// ``` #[inline] #[stable(feature = "iterator_try_fold", since = "1.27.0")] - fn try_rfold(&mut self, init: B, mut f: F) -> R + fn try_rfold(&mut self, init: B, f: F) -> R where Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try, { - let mut accum = init; - while let Some(x) = self.next_back() { - accum = f(accum, x)?; - } - try { accum } + default_try_rfold(self, init, f) } /// An iterator method that reduces the iterator's elements to a single, @@ -278,16 +274,12 @@ pub trait DoubleEndedIterator: Iterator { /// ``` #[inline] #[stable(feature = "iter_rfold", since = "1.27.0")] - fn rfold(mut self, init: B, mut f: F) -> B + fn rfold(self, init: B, f: F) -> B where Self: Sized, F: FnMut(B, Self::Item) -> B, { - let mut accum = init; - while let Some(x) = self.next_back() { - accum = f(accum, x); - } - accum + default_rfold(self, init, f) } /// Searches for an element of an iterator from the back that satisfies a predicate. @@ -349,6 +341,33 @@ pub trait DoubleEndedIterator: Iterator { } } +#[inline] +fn default_try_rfold(iter: &mut I, init: B, mut f: F) -> R +where + I: DoubleEndedIterator + ?Sized, + F: FnMut(B, I::Item) -> R, + R: Try, +{ + let mut accum = init; + while let Some(x) = iter.next_back() { + accum = f(accum, x)?; + } + try { accum } +} + +#[inline] +fn default_rfold(mut iter: I, init: B, mut f: F) -> B +where + I: DoubleEndedIterator, + F: FnMut(B, I::Item) -> B, +{ + let mut accum = init; + while let Some(x) = iter.next_back() { + accum = f(accum, x); + } + accum +} + #[stable(feature = "rust1", since = "1.0.0")] impl<'a, I: DoubleEndedIterator + ?Sized> DoubleEndedIterator for &'a mut I { fn next_back(&mut self) -> Option { @@ -360,4 +379,69 @@ impl<'a, I: DoubleEndedIterator + ?Sized> DoubleEndedIterator for &'a mut I { fn nth_back(&mut self, n: usize) -> Option { (**self).nth_back(n) } + #[inline] + fn try_rfold(&mut self, init: B, f: F) -> R + where + F: FnMut(B, Self::Item) -> R, + R: Try, + { + SpecSizedDoubleEndedIterator::try_rfold(self, init, f) + } + #[inline] + fn rfold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + SpecSizedDoubleEndedIterator::rfold(self, init, f) + } +} + +trait SpecSizedDoubleEndedIterator: DoubleEndedIterator { + fn try_rfold(&mut self, init: B, f: F) -> R + where + F: FnMut(B, Self::Item) -> R, + R: Try; + fn rfold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B; +} + +impl SpecSizedDoubleEndedIterator for &mut I { + #[inline] + default fn try_rfold(&mut self, init: B, f: F) -> R + where + F: FnMut(B, Self::Item) -> R, + R: Try, + { + default_try_rfold(self, init, f) + } + #[inline] + default fn rfold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + default_rfold(self, init, f) + } +} + +impl SpecSizedDoubleEndedIterator for &mut I { + #[inline] + fn try_rfold(&mut self, init: B, f: F) -> R + where + F: FnMut(B, Self::Item) -> R, + R: Try, + { + (**self).try_rfold(init, f) + } + #[inline] + fn rfold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + #[inline] + fn ok(mut f: impl FnMut(T, U) -> T) -> impl FnMut(T, U) -> Result { + move |acc, x| Ok(f(acc, x)) + } + self.try_rfold(init, ok(f)).unwrap() + } } diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs index da9e5fde7ccec..2f9ca3b152e8f 100644 --- a/library/core/src/iter/traits/iterator.rs +++ b/library/core/src/iter/traits/iterator.rs @@ -1995,17 +1995,13 @@ pub trait Iterator { /// ``` #[inline] #[stable(feature = "iterator_try_fold", since = "1.27.0")] - fn try_fold(&mut self, init: B, mut f: F) -> R + fn try_fold(&mut self, init: B, f: F) -> R where Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try, { - let mut accum = init; - while let Some(x) = self.next() { - accum = f(accum, x)?; - } - try { accum } + default_try_fold(self, init, f) } /// An iterator method that applies a fallible function to each item in the @@ -2136,16 +2132,12 @@ pub trait Iterator { #[doc(alias = "inject")] #[inline] #[stable(feature = "rust1", since = "1.0.0")] - fn fold(mut self, init: B, mut f: F) -> B + fn fold(self, init: B, f: F) -> B where Self: Sized, F: FnMut(B, Self::Item) -> B, { - let mut accum = init; - while let Some(x) = self.next() { - accum = f(accum, x); - } - accum + default_fold(self, init, f) } /// Reduces the elements to a single one, by repeatedly applying a reducing @@ -3415,6 +3407,33 @@ pub trait Iterator { } } +#[inline] +fn default_try_fold(iter: &mut I, init: B, mut f: F) -> R +where + I: Iterator + ?Sized, + F: FnMut(B, I::Item) -> R, + R: Try, +{ + let mut accum = init; + while let Some(x) = iter.next() { + accum = f(accum, x)?; + } + try { accum } +} + +#[inline] +fn default_fold(mut iter: I, init: B, mut f: F) -> B +where + I: Iterator, + F: FnMut(B, I::Item) -> B, +{ + let mut accum = init; + while let Some(x) = iter.next() { + accum = f(accum, x); + } + accum +} + #[stable(feature = "rust1", since = "1.0.0")] impl Iterator for &mut I { type Item = I::Item; @@ -3430,4 +3449,69 @@ impl Iterator for &mut I { fn nth(&mut self, n: usize) -> Option { (**self).nth(n) } + #[inline] + fn try_fold(&mut self, init: B, f: F) -> R + where + F: FnMut(B, Self::Item) -> R, + R: Try, + { + SpecSizedIterator::try_fold(self, init, f) + } + #[inline] + fn fold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + SpecSizedIterator::fold(self, init, f) + } +} + +trait SpecSizedIterator: Iterator { + fn try_fold(&mut self, init: B, f: F) -> R + where + F: FnMut(B, Self::Item) -> R, + R: Try; + fn fold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B; +} + +impl SpecSizedIterator for &mut I { + #[inline] + default fn try_fold(&mut self, init: B, f: F) -> R + where + F: FnMut(B, Self::Item) -> R, + R: Try, + { + default_try_fold(self, init, f) + } + #[inline] + default fn fold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + default_fold(self, init, f) + } +} + +impl SpecSizedIterator for &mut I { + #[inline] + fn try_fold(&mut self, init: B, f: F) -> R + where + F: FnMut(B, Self::Item) -> R, + R: Try, + { + (**self).try_fold(init, f) + } + #[inline] + fn fold(self, init: B, f: F) -> B + where + F: FnMut(B, Self::Item) -> B, + { + #[inline] + fn ok(mut f: impl FnMut(T, U) -> T) -> impl FnMut(T, U) -> Result { + move |acc, x| Ok(f(acc, x)) + } + self.try_fold(init, ok(f)).unwrap() + } } diff --git a/src/test/ui/issues/issue-3044.stderr b/src/test/ui/issues/issue-3044.stderr index b93aeade95e42..b8300d5071ccd 100644 --- a/src/test/ui/issues/issue-3044.stderr +++ b/src/test/ui/issues/issue-3044.stderr @@ -11,7 +11,7 @@ LL | | }); note: associated function defined here --> $SRC_DIR/core/src/iter/traits/iterator.rs:LL:COL | -LL | fn fold(mut self, init: B, mut f: F) -> B +LL | fn fold(self, init: B, f: F) -> B | ^^^^ error: aborting due to previous error