Rollup merge of #104435 - scottmcm:iter-repeat-n, r=thomcc

`VecDeque::resize` should re-use the buffer in the passed-in element

Today it always copies it for *every* appended element, but one of those clones is avoidable.

This adds `iter::repeat_n` (#104434) as the primitive needed to do this.  If this PR is acceptable, I'll also use this in `Vec` rather than its custom `ExtendElement` type & infrastructure that is harder to share between multiple different containers:

https://github.com/rust-lang/rust/blob/101e1822c3e54e63996c8aaa014d55716f3937eb/library/alloc/src/vec/mod.rs#L2479-L2492

Author: Yuki Okushi

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<T: Clone, A: Allocator> VecDeque<T, A> {
     /// ```
     #[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);
+        }
     }
 }
 

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)]

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::<i32>::with_capacity(1234);
+    let mut d = VecDeque::new();
+    d.resize(1, v);
+    assert_eq!(d[0].capacity(), 1234);
+}

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")]

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};
 

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<i32> = 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<T: Clone>(element: T, count: usize) -> RepeatN<T> {
+    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<A> {
+    count: usize,
+    // Invariant: has been dropped iff count == 0.
+    element: ManuallyDrop<A>,
+}
+
+impl<A> RepeatN<A> {
+    /// 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<A> {
+        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<A> Drop for RepeatN<A> {
+    fn drop(&mut self) {
+        self.take_element();
+    }
+}
+
+#[unstable(feature = "iter_repeat_n", issue = "104434")]
+impl<A: Clone> Iterator for RepeatN<A> {
+    type Item = A;
+
+    #[inline]
+    fn next(&mut self) -> Option<A> {
+        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<usize>) {
+        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<A> {
+        self.take_element()
+    }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.len()
+    }
+}
+
+#[unstable(feature = "iter_repeat_n", issue = "104434")]
+impl<A: Clone> ExactSizeIterator for RepeatN<A> {
+    fn len(&self) -> usize {
+        self.count
+    }
+}
+
+#[unstable(feature = "iter_repeat_n", issue = "104434")]
+impl<A: Clone> DoubleEndedIterator for RepeatN<A> {
+    #[inline]
+    fn next_back(&mut self) -> Option<A> {
+        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<A> {
+        self.nth(n)
+    }
+}
+
+#[unstable(feature = "iter_repeat_n", issue = "104434")]
+impl<A: Clone> FusedIterator for RepeatN<A> {}
+
+#[unstable(feature = "trusted_len", issue = "37572")]
+unsafe impl<A: Clone> TrustedLen for RepeatN<A> {}

library/core/tests/iter/sources.rs

@@ -106,3 +106,52 @@ fn test_empty() {
     let mut it = empty::<i32>();
     assert_eq!(it.next(), None);
 }
+
+#[test]
+fn test_repeat_n_drop() {
+    #[derive(Clone, Debug)]
+    struct DropCounter<'a>(&'a Cell<usize>);
+    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);
+}

library/core/tests/lib.rs

@@ -75,6 +75,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)]

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<NotCopy>) -> Option<NotCopy> {
+    // 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<u8> {
+    // 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
+}