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` (https://github.com/rust-lang/rust/issues/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:

101e1822c3/library/alloc/src/vec/mod.rs (L2479-L2492)
This commit is contained in:
Yuki Okushi 2022-11-20 13:15:59 +09:00 committed by GitHub
commit 785237d392
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 323 additions and 2 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

@ -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> {}

View File

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

View File

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

View File

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