Skip to content

Commit

Permalink
Rollup merge of rust-lang#128530 - scottmcm:repeat-n-unchecked, r=joboet
Browse files Browse the repository at this point in the history
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
matthiaskrgr authored Aug 3, 2024
2 parents 0b5f1b8 + 05d8d7c commit 43a1e93
Showing 1 changed file with 20 additions and 13 deletions.
33 changes: 20 additions & 13 deletions 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)
}
}
}

0 comments on commit 43a1e93

Please sign in to comment.