diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs index 2a57dad89a770..537fe22a4be72 100644 --- a/library/alloc/src/collections/vec_deque/mod.rs +++ b/library/alloc/src/collections/vec_deque/mod.rs @@ -10,7 +10,7 @@ use core::cmp::{self, Ordering}; use core::fmt; use core::hash::{Hash, Hasher}; -use core::iter::{repeat_with, FromIterator}; +use core::iter::{repeat_n, repeat_with, FromIterator}; use core::marker::PhantomData; use core::mem::{ManuallyDrop, MaybeUninit, SizedTypeProperties}; use core::ops::{Index, IndexMut, Range, RangeBounds}; @@ -2833,7 +2833,12 @@ impl VecDeque { /// ``` #[stable(feature = "deque_extras", since = "1.16.0")] pub fn resize(&mut self, new_len: usize, value: T) { - self.resize_with(new_len, || value.clone()); + if new_len > self.len() { + let extra = new_len - self.len(); + self.extend(repeat_n(value, extra)) + } else { + self.truncate(new_len); + } } } diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index 008926666c136..5e13547abcb70 100644 --- a/library/alloc/src/lib.rs +++ b/library/alloc/src/lib.rs @@ -124,6 +124,7 @@ #![feature(inplace_iteration)] #![feature(iter_advance_by)] #![feature(iter_next_chunk)] +#![feature(iter_repeat_n)] #![feature(layout_for_ptr)] #![feature(maybe_uninit_slice)] #![feature(maybe_uninit_uninit_array)] diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 766006939fa48..69a107440c20a 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -2163,7 +2163,7 @@ impl Vec { { let len = self.len(); if new_len > len { - self.extend_with(new_len - len, ExtendFunc(f)); + self.extend(iter::repeat_with(f).take(new_len - len)); } else { self.truncate(new_len); } @@ -2361,7 +2361,7 @@ impl Vec { let len = self.len(); if new_len > len { - self.extend_with(new_len - len, ExtendElement(value)) + self.extend(iter::repeat_n(value, new_len - len)) } else { self.truncate(new_len); } @@ -2475,64 +2475,6 @@ impl Vec<[T; N], A> { } } -// This code generalizes `extend_with_{element,default}`. -trait ExtendWith { - fn next(&mut self) -> T; - fn last(self) -> T; -} - -struct ExtendElement(T); -impl ExtendWith for ExtendElement { - fn next(&mut self) -> T { - self.0.clone() - } - fn last(self) -> T { - self.0 - } -} - -struct ExtendFunc(F); -impl T> ExtendWith for ExtendFunc { - fn next(&mut self) -> T { - (self.0)() - } - fn last(mut self) -> T { - (self.0)() - } -} - -impl Vec { - #[cfg(not(no_global_oom_handling))] - /// Extend the vector by `n` values, using the given generator. - fn extend_with>(&mut self, n: usize, mut value: E) { - self.reserve(n); - - unsafe { - let mut ptr = self.as_mut_ptr().add(self.len()); - // Use SetLenOnDrop to work around bug where compiler - // might not realize the store through `ptr` through self.set_len() - // don't alias. - let mut local_len = SetLenOnDrop::new(&mut self.len); - - // Write all elements except the last one - for _ in 1..n { - ptr::write(ptr, value.next()); - ptr = ptr.add(1); - // Increment the length in every step in case next() panics - local_len.increment_len(1); - } - - if n > 0 { - // We can write the last element directly without cloning needlessly - ptr::write(ptr, value.last()); - local_len.increment_len(1); - } - - // len set by scope guard - } - } -} - impl Vec { /// Removes consecutive repeated elements in the vector according to the /// [`PartialEq`] trait implementation. diff --git a/library/alloc/src/vec/spec_from_elem.rs b/library/alloc/src/vec/spec_from_elem.rs index ff364c033ee98..f71cd7c35c5b6 100644 --- a/library/alloc/src/vec/spec_from_elem.rs +++ b/library/alloc/src/vec/spec_from_elem.rs @@ -1,61 +1,33 @@ -use core::ptr; +use core::iter; use crate::alloc::Allocator; use crate::raw_vec::RawVec; -use super::{ExtendElement, IsZero, Vec}; +use super::{IsZero, Vec}; // Specialization trait used for Vec::from_elem pub(super) trait SpecFromElem: Sized { fn from_elem(elem: Self, n: usize, alloc: A) -> Vec; } +fn basic_from_elem(elem: T, n: usize, alloc: A) -> Vec { + let mut v = Vec::with_capacity_in(n, alloc); + v.extend(iter::repeat_n(elem, n)); + v +} + impl SpecFromElem for T { default fn from_elem(elem: Self, n: usize, alloc: A) -> Vec { - let mut v = Vec::with_capacity_in(n, alloc); - v.extend_with(n, ExtendElement(elem)); - v + basic_from_elem(elem, n, alloc) } } impl SpecFromElem for T { #[inline] - default fn from_elem(elem: T, n: usize, alloc: A) -> Vec { + fn from_elem(elem: T, n: usize, alloc: A) -> Vec { if elem.is_zero() { return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n }; } - let mut v = Vec::with_capacity_in(n, alloc); - v.extend_with(n, ExtendElement(elem)); - v - } -} - -impl SpecFromElem for i8 { - #[inline] - fn from_elem(elem: i8, n: usize, alloc: A) -> Vec { - if elem == 0 { - return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n }; - } - unsafe { - let mut v = Vec::with_capacity_in(n, alloc); - ptr::write_bytes(v.as_mut_ptr(), elem as u8, n); - v.set_len(n); - v - } - } -} - -impl SpecFromElem for u8 { - #[inline] - fn from_elem(elem: u8, n: usize, alloc: A) -> Vec { - if elem == 0 { - return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n }; - } - unsafe { - let mut v = Vec::with_capacity_in(n, alloc); - ptr::write_bytes(v.as_mut_ptr(), elem, n); - v.set_len(n); - v - } + basic_from_elem(elem, n, alloc) } } diff --git a/library/alloc/tests/vec_deque.rs b/library/alloc/tests/vec_deque.rs index 019d73c0b16bb..c1b9e7af9d857 100644 --- a/library/alloc/tests/vec_deque.rs +++ b/library/alloc/tests/vec_deque.rs @@ -1727,3 +1727,11 @@ fn test_from_zero_sized_vec() { let queue = VecDeque::from(v); assert_eq!(queue.len(), 100); } + +#[test] +fn test_resize_keeps_reserved_space_from_item() { + let v = Vec::::with_capacity(1234); + let mut d = VecDeque::new(); + d.resize(1, v); + assert_eq!(d[0].capacity(), 1234); +} diff --git a/library/core/src/iter/mod.rs b/library/core/src/iter/mod.rs index ef0f397825b15..bb35d50b4bfda 100644 --- a/library/core/src/iter/mod.rs +++ b/library/core/src/iter/mod.rs @@ -401,6 +401,8 @@ pub use self::sources::{once, Once}; pub use self::sources::{once_with, OnceWith}; #[stable(feature = "rust1", since = "1.0.0")] pub use self::sources::{repeat, Repeat}; +#[unstable(feature = "iter_repeat_n", issue = "104434")] +pub use self::sources::{repeat_n, RepeatN}; #[stable(feature = "iterator_repeat_with", since = "1.28.0")] pub use self::sources::{repeat_with, RepeatWith}; #[stable(feature = "iter_successors", since = "1.34.0")] diff --git a/library/core/src/iter/sources.rs b/library/core/src/iter/sources.rs index d34772cd304ff..3ec426a3ad9a1 100644 --- a/library/core/src/iter/sources.rs +++ b/library/core/src/iter/sources.rs @@ -4,6 +4,7 @@ mod from_generator; mod once; mod once_with; mod repeat; +mod repeat_n; mod repeat_with; mod successors; @@ -16,6 +17,9 @@ pub use self::empty::{empty, Empty}; #[stable(feature = "iter_once", since = "1.2.0")] pub use self::once::{once, Once}; +#[unstable(feature = "iter_repeat_n", issue = "104434")] +pub use self::repeat_n::{repeat_n, RepeatN}; + #[stable(feature = "iterator_repeat_with", since = "1.28.0")] pub use self::repeat_with::{repeat_with, RepeatWith}; diff --git a/library/core/src/iter/sources/repeat_n.rs b/library/core/src/iter/sources/repeat_n.rs new file mode 100644 index 0000000000000..dc69bf4df595e --- /dev/null +++ b/library/core/src/iter/sources/repeat_n.rs @@ -0,0 +1,195 @@ +use crate::iter::{FusedIterator, TrustedLen}; +use crate::mem::ManuallyDrop; + +/// Creates a new iterator that repeats a single element a given number of times. +/// +/// The `repeat_n()` function repeats a single value exactly `n` times. +/// +/// This is very similar to using [`repeat()`] with [`Iterator::take()`], +/// but there are two differences: +/// - `repeat_n()` can return the original value, rather than always cloning. +/// - `repeat_n()` produces an [`ExactSizeIterator`]. +/// +/// [`repeat()`]: crate::iter::repeat +/// +/// # Examples +/// +/// Basic usage: +/// +/// ``` +/// #![feature(iter_repeat_n)] +/// use std::iter; +/// +/// // four of the the number four: +/// let mut four_fours = iter::repeat_n(4, 4); +/// +/// assert_eq!(Some(4), four_fours.next()); +/// assert_eq!(Some(4), four_fours.next()); +/// assert_eq!(Some(4), four_fours.next()); +/// assert_eq!(Some(4), four_fours.next()); +/// +/// // no more fours +/// assert_eq!(None, four_fours.next()); +/// ``` +/// +/// For non-`Copy` types, +/// +/// ``` +/// #![feature(iter_repeat_n)] +/// use std::iter; +/// +/// let v: Vec = Vec::with_capacity(123); +/// let mut it = iter::repeat_n(v, 5); +/// +/// for i in 0..4 { +/// // It starts by cloning things +/// let cloned = it.next().unwrap(); +/// assert_eq!(cloned.len(), 0); +/// assert_eq!(cloned.capacity(), 0); +/// } +/// +/// // ... but the last item is the original one +/// let last = it.next().unwrap(); +/// assert_eq!(last.len(), 0); +/// assert_eq!(last.capacity(), 123); +/// +/// // ... and now we're done +/// assert_eq!(None, it.next()); +/// ``` +#[inline] +#[unstable(feature = "iter_repeat_n", issue = "104434")] +#[doc(hidden)] // waiting on ACP#120 to decide whether to expose publicly +pub fn repeat_n(element: T, count: usize) -> RepeatN { + let mut element = ManuallyDrop::new(element); + + if count == 0 { + // SAFETY: we definitely haven't dropped it yet, since we only just got + // passed it in, and because the count is zero the instance we're about + // to create won't drop it, so to avoid leaking we need to now. + unsafe { ManuallyDrop::drop(&mut element) }; + } + + RepeatN { element, count } +} + +/// An iterator that repeats an element an exact number of times. +/// +/// This `struct` is created by the [`repeat_n()`] function. +/// See its documentation for more. +#[derive(Clone, Debug)] +#[unstable(feature = "iter_repeat_n", issue = "104434")] +#[doc(hidden)] // waiting on ACP#120 to decide whether to expose publicly +pub struct RepeatN { + count: usize, + // Invariant: has been dropped iff count == 0. + element: ManuallyDrop, +} + +impl RepeatN { + /// If we haven't already dropped the element, return it in an option. + /// + /// Clears the count so it won't be dropped again later. + #[inline] + fn take_element(&mut self) -> Option { + if self.count > 0 { + self.count = 0; + // SAFETY: We just set count to zero so it won't be dropped again, + // and it used to be non-zero so it hasn't already been dropped. + unsafe { Some(ManuallyDrop::take(&mut self.element)) } + } else { + None + } + } +} + +#[unstable(feature = "iter_repeat_n", issue = "104434")] +impl Drop for RepeatN { + fn drop(&mut self) { + self.take_element(); + } +} + +#[unstable(feature = "iter_repeat_n", issue = "104434")] +impl Iterator for RepeatN { + type Item = A; + + #[inline] + fn next(&mut self) -> Option { + if self.count == 0 { + return None; + } + + self.count -= 1; + Some(if self.count == 0 { + // SAFETY: the check above ensured that the count used to be non-zero, + // so element hasn't been dropped yet, and we just lowered the count to + // zero so it won't be dropped later, and thus it's okay to take it here. + unsafe { ManuallyDrop::take(&mut self.element) } + } else { + A::clone(&mut self.element) + }) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + let len = self.len(); + (len, Some(len)) + } + + #[inline] + fn advance_by(&mut self, skip: usize) -> Result<(), usize> { + let len = self.count; + + if skip >= len { + self.take_element(); + } + + if skip > len { + Err(len) + } else { + self.count = len - skip; + Ok(()) + } + } + + #[inline] + fn last(mut self) -> Option { + self.take_element() + } + + #[inline] + fn count(self) -> usize { + self.len() + } +} + +#[unstable(feature = "iter_repeat_n", issue = "104434")] +impl ExactSizeIterator for RepeatN { + fn len(&self) -> usize { + self.count + } +} + +#[unstable(feature = "iter_repeat_n", issue = "104434")] +impl DoubleEndedIterator for RepeatN { + #[inline] + fn next_back(&mut self) -> Option { + self.next() + } + + #[inline] + fn advance_back_by(&mut self, n: usize) -> Result<(), usize> { + self.advance_by(n) + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option { + self.nth(n) + } +} + +#[unstable(feature = "iter_repeat_n", issue = "104434")] +impl FusedIterator for RepeatN {} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl TrustedLen for RepeatN {} diff --git a/library/core/tests/iter/sources.rs b/library/core/tests/iter/sources.rs index d0114ade6e463..a15f3a5148f0a 100644 --- a/library/core/tests/iter/sources.rs +++ b/library/core/tests/iter/sources.rs @@ -106,3 +106,52 @@ fn test_empty() { let mut it = empty::(); assert_eq!(it.next(), None); } + +#[test] +fn test_repeat_n_drop() { + #[derive(Clone, Debug)] + struct DropCounter<'a>(&'a Cell); + impl Drop for DropCounter<'_> { + fn drop(&mut self) { + self.0.set(self.0.get() + 1); + } + } + + // `repeat_n(x, 0)` drops `x` immediately + let count = Cell::new(0); + let item = DropCounter(&count); + let mut it = repeat_n(item, 0); + assert_eq!(count.get(), 1); + assert!(it.next().is_none()); + assert_eq!(count.get(), 1); + drop(it); + assert_eq!(count.get(), 1); + + // Dropping the iterator needs to drop the item if it's non-empty + let count = Cell::new(0); + let item = DropCounter(&count); + let it = repeat_n(item, 3); + assert_eq!(count.get(), 0); + drop(it); + assert_eq!(count.get(), 1); + + // Dropping the iterator doesn't drop the item if it was exhausted + let count = Cell::new(0); + let item = DropCounter(&count); + let mut it = repeat_n(item, 3); + assert_eq!(count.get(), 0); + let x0 = it.next().unwrap(); + assert_eq!(count.get(), 0); + let x1 = it.next().unwrap(); + assert_eq!(count.get(), 0); + let x2 = it.next().unwrap(); + assert_eq!(count.get(), 0); + assert!(it.next().is_none()); + assert_eq!(count.get(), 0); + assert!(it.next().is_none()); + assert_eq!(count.get(), 0); + drop(it); + assert_eq!(count.get(), 0); + drop((x0, x1, x2)); + assert_eq!(count.get(), 3); +} diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index 61d31b3448734..58e9045c9b72d 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -73,6 +73,7 @@ #![feature(iter_is_partitioned)] #![feature(iter_next_chunk)] #![feature(iter_order_by)] +#![feature(iter_repeat_n)] #![feature(iterator_try_collect)] #![feature(iterator_try_reduce)] #![feature(const_mut_refs)] diff --git a/src/test/codegen/iter-repeat-n-trivial-drop.rs b/src/test/codegen/iter-repeat-n-trivial-drop.rs new file mode 100644 index 0000000000000..20e1d9b4d5988 --- /dev/null +++ b/src/test/codegen/iter-repeat-n-trivial-drop.rs @@ -0,0 +1,56 @@ +// compile-flags: -O +// only-x86_64 +// ignore-debug: the debug assertions get in the way + +#![crate_type = "lib"] +#![feature(iter_repeat_n)] + +#[derive(Clone)] +pub struct NotCopy(u16); + +impl Drop for NotCopy { + fn drop(&mut self) {} +} + +// For a type where `Drop::drop` doesn't do anything observable and a clone is the +// same as a move, make sure that the extra case for the last item disappears. + +#[no_mangle] +// CHECK-LABEL: @iter_repeat_n_next +pub fn iter_repeat_n_next(it: &mut std::iter::RepeatN) -> Option { + // CHECK-NEXT: start: + // CHECK-NOT: br + // CHECK: %[[COUNT:.+]] = load i64 + // CHECK-NEXT: %[[COUNT_ZERO:.+]] = icmp eq i64 %[[COUNT]], 0 + // CHECK-NEXT: br i1 %[[COUNT_ZERO]], label %[[EMPTY:.+]], label %[[NOT_EMPTY:.+]] + + // CHECK: [[NOT_EMPTY]]: + // CHECK-NEXT: %[[DEC:.+]] = add i64 %[[COUNT]], -1 + // CHECK-NEXT: store i64 %[[DEC]] + // CHECK-NOT: br + // CHECK: %[[VAL:.+]] = load i16 + // CHECK-NEXT: br label %[[EMPTY]] + + // CHECK: [[EMPTY]]: + // CHECK-NOT: br + // CHECK: phi i16 [ undef, %start ], [ %[[VAL]], %[[NOT_EMPTY]] ] + // CHECK-NOT: br + // CHECK: ret + + it.next() +} + +// And as a result, using the iterator can optimize without special cases for +// the last iteration, like `memset`ing all the items in one call. + +#[no_mangle] +// CHECK-LABEL: @vec_extend_via_iter_repeat_n +pub fn vec_extend_via_iter_repeat_n() -> Vec { + // CHECK: %[[ADDR:.+]] = tail call dereferenceable_or_null(1234) ptr @__rust_alloc(i64 1234, i64 1) + // CHECK: tail call void @llvm.memset.p0.i64(ptr noundef nonnull align 1 dereferenceable(1234) %[[ADDR]], i8 42, i64 1234, + + let n = 1234_usize; + let mut v = Vec::with_capacity(n); + v.extend(std::iter::repeat_n(42_u8, n)); + v +} diff --git a/src/test/codegen/vec-calloc.rs b/src/test/codegen/vec-calloc.rs index ae6e448f172f7..7fd350dc0f45d 100644 --- a/src/test/codegen/vec-calloc.rs +++ b/src/test/codegen/vec-calloc.rs @@ -6,6 +6,7 @@ #![crate_type = "lib"] // CHECK-LABEL: @vec_zero_bytes +// CHECK-SAME: i64 %n #[no_mangle] pub fn vec_zero_bytes(n: usize) -> Vec { // CHECK-NOT: call {{.*}}alloc::vec::from_elem @@ -13,7 +14,7 @@ pub fn vec_zero_bytes(n: usize) -> Vec { // CHECK-NOT: call {{.*}}__rust_alloc( // CHECK-NOT: call {{.*}}llvm.memset - // CHECK: call {{.*}}__rust_alloc_zeroed( + // CHECK: call {{.*}}__rust_alloc_zeroed(i64 %n // CHECK-NOT: call {{.*}}alloc::vec::from_elem // CHECK-NOT: call {{.*}}reserve @@ -22,17 +23,19 @@ pub fn vec_zero_bytes(n: usize) -> Vec { // CHECK: ret void vec![0; n] -} +} // CHECK-LABEL: @vec_one_bytes +// CHECK-SAME: i64 %n #[no_mangle] pub fn vec_one_bytes(n: usize) -> Vec { // CHECK-NOT: call {{.*}}alloc::vec::from_elem // CHECK-NOT: call {{.*}}reserve // CHECK-NOT: call {{.*}}__rust_alloc_zeroed( - // CHECK: call {{.*}}__rust_alloc( + // CHECK: call {{.*}}__rust_alloc(i64 %n // CHECK: call {{.*}}llvm.memset + // CHECK-SAME: i8 1, i64 %n // CHECK-NOT: call {{.*}}alloc::vec::from_elem // CHECK-NOT: call {{.*}}reserve @@ -43,13 +46,20 @@ pub fn vec_one_bytes(n: usize) -> Vec { } // CHECK-LABEL: @vec_zero_scalar +// CHECK-SAME: i64 %n #[no_mangle] pub fn vec_zero_scalar(n: usize) -> Vec { // CHECK-NOT: call {{.*}}alloc::vec::from_elem // CHECK-NOT: call {{.*}}reserve // CHECK-NOT: call {{.*}}__rust_alloc( - // CHECK: call {{.*}}__rust_alloc_zeroed( + // CHECK: %[[BYTES:.+]] = shl i64 %n, 2 + + // CHECK-NOT: call {{.*}}alloc::vec::from_elem + // CHECK-NOT: call {{.*}}reserve + // CHECK-NOT: call {{.*}}__rust_alloc( + + // CHECK: call {{.*}}__rust_alloc_zeroed(i64 %[[BYTES]] // CHECK-NOT: call {{.*}}alloc::vec::from_elem // CHECK-NOT: call {{.*}}reserve @@ -60,13 +70,20 @@ pub fn vec_zero_scalar(n: usize) -> Vec { } // CHECK-LABEL: @vec_one_scalar +// CHECK-SAME: i64 %n #[no_mangle] pub fn vec_one_scalar(n: usize) -> Vec { // CHECK-NOT: call {{.*}}alloc::vec::from_elem // CHECK-NOT: call {{.*}}reserve // CHECK-NOT: call {{.*}}__rust_alloc_zeroed( - // CHECK: call {{.*}}__rust_alloc( + // CHECK: %[[BYTES:.+]] = shl i64 %n, 2 + + // CHECK-NOT: call {{.*}}alloc::vec::from_elem + // CHECK-NOT: call {{.*}}reserve + // CHECK-NOT: call {{.*}}__rust_alloc_zeroed( + + // CHECK: call {{.*}}__rust_alloc(i64 %[[BYTES]] // CHECK-NOT: call {{.*}}alloc::vec::from_elem // CHECK-NOT: call {{.*}}reserve