diff --git a/library/core/src/iter/sources/repeat_n.rs b/library/core/src/iter/sources/repeat_n.rs index 8390dab8e54..4c4ae39f836 100644 --- a/library/core/src/iter/sources/repeat_n.rs +++ b/library/core/src/iter/sources/repeat_n.rs @@ -114,19 +114,12 @@ impl Iterator for RepeatN { #[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) } + if self.count > 0 { + // SAFETY: Just checked it's not empty + unsafe { Some(self.next_unchecked()) } } else { - A::clone(&self.element) - }) + None + } } #[inline] @@ -194,4 +187,18 @@ impl FusedIterator for RepeatN {} #[unstable(feature = "trusted_len", issue = "37572")] unsafe impl TrustedLen for RepeatN {} #[unstable(feature = "trusted_len_next_unchecked", issue = "37572")] -impl UncheckedIterator for RepeatN {} +impl UncheckedIterator for RepeatN { + #[inline] + unsafe fn next_unchecked(&mut self) -> Self::Item { + // SAFETY: The caller promised the iterator isn't empty + self.count = unsafe { self.count.unchecked_sub(1) }; + 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(&self.element) + } + } +} diff --git a/tests/codegen/iter-repeat-n-trivial-drop.rs b/tests/codegen/iter-repeat-n-trivial-drop.rs index 31020b77984..7de224b92d8 100644 --- a/tests/codegen/iter-repeat-n-trivial-drop.rs +++ b/tests/codegen/iter-repeat-n-trivial-drop.rs @@ -1,8 +1,9 @@ -//@ compile-flags: -O +//@ compile-flags: -C opt-level=3 //@ only-x86_64 #![crate_type = "lib"] #![feature(iter_repeat_n)] +#![feature(array_repeat)] #[derive(Clone)] pub struct NotCopy(u16); @@ -54,3 +55,15 @@ pub fn vec_extend_via_iter_repeat_n() -> Vec { v.extend(std::iter::repeat_n(42_u8, n)); v } + +// Array repeat uses `RepeatN::next_unchecked` internally, +// so also check that the distinction disappears there. + +#[no_mangle] +// CHECK-LABEL: @array_repeat_not_copy +pub unsafe fn array_repeat_not_copy(item: NotCopy) -> [NotCopy; 8] { + // CHECK: insertelement {{.+}} i16 %item + // CHECK: shufflevector <8 x i16> {{.+}} zeroinitializer + // CHECK: store <8 x i16> + std::array::repeat(item) +}