Skip to content

Commit

Permalink
Unrolled build for rust-lang#128530
Browse files Browse the repository at this point in the history
Rollup merge of rust-lang#128530 - scottmcm:repeat-n-unchecked, r=joboet

Implement `UncheckedIterator` directly for `RepeatN`

This just pulls the code out of `next` into `next_unchecked`, rather than making the `Some` and `unwrap_unchecked`ing it.

And while I was touching it, I added a codegen test that `array::repeat` for something that's just `Clone`, not `Copy`, still ends up optimizing to the same thing as `[x; n]`: <https://rust.godbolt.org/z/YY3a5ajMW>.
  • Loading branch information
rust-timer committed Aug 3, 2024
2 parents bbf60c8 + 77ca30f commit be145f3
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 14 deletions.
33 changes: 20 additions & 13 deletions library/core/src/iter/sources/repeat_n.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,19 +114,12 @@ impl<A: Clone> Iterator for RepeatN<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) }
if self.count > 0 {
// SAFETY: Just checked it's not empty
unsafe { Some(self.next_unchecked()) }
} else {
A::clone(&self.element)
})
None
}
}

#[inline]
Expand Down Expand Up @@ -194,4 +187,18 @@ impl<A: Clone> FusedIterator for RepeatN<A> {}
#[unstable(feature = "trusted_len", issue = "37572")]
unsafe impl<A: Clone> TrustedLen for RepeatN<A> {}
#[unstable(feature = "trusted_len_next_unchecked", issue = "37572")]
impl<A: Clone> UncheckedIterator for RepeatN<A> {}
impl<A: Clone> UncheckedIterator for RepeatN<A> {
#[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)
}
}
}
15 changes: 14 additions & 1 deletion tests/codegen/iter-repeat-n-trivial-drop.rs
Original file line number Diff line number Diff line change
@@ -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);
Expand Down Expand Up @@ -54,3 +55,15 @@ pub fn vec_extend_via_iter_repeat_n() -> Vec<u8> {
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)
}

0 comments on commit be145f3

Please sign in to comment.