Skip to content

Commit 43a1e93

Browse files
authored
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>.
2 parents 0b5f1b8 + 05d8d7c commit 43a1e93

File tree

1 file changed

+20
-13
lines changed

1 file changed

+20
-13
lines changed

core/src/iter/sources/repeat_n.rs

+20-13
Original file line numberDiff line numberDiff line change
@@ -114,19 +114,12 @@ impl<A: Clone> Iterator for RepeatN<A> {
114114

115115
#[inline]
116116
fn next(&mut self) -> Option<A> {
117-
if self.count == 0 {
118-
return None;
119-
}
120-
121-
self.count -= 1;
122-
Some(if self.count == 0 {
123-
// SAFETY: the check above ensured that the count used to be non-zero,
124-
// so element hasn't been dropped yet, and we just lowered the count to
125-
// zero so it won't be dropped later, and thus it's okay to take it here.
126-
unsafe { ManuallyDrop::take(&mut self.element) }
117+
if self.count > 0 {
118+
// SAFETY: Just checked it's not empty
119+
unsafe { Some(self.next_unchecked()) }
127120
} else {
128-
A::clone(&self.element)
129-
})
121+
None
122+
}
130123
}
131124

132125
#[inline]
@@ -194,4 +187,18 @@ impl<A: Clone> FusedIterator for RepeatN<A> {}
194187
#[unstable(feature = "trusted_len", issue = "37572")]
195188
unsafe impl<A: Clone> TrustedLen for RepeatN<A> {}
196189
#[unstable(feature = "trusted_len_next_unchecked", issue = "37572")]
197-
impl<A: Clone> UncheckedIterator for RepeatN<A> {}
190+
impl<A: Clone> UncheckedIterator for RepeatN<A> {
191+
#[inline]
192+
unsafe fn next_unchecked(&mut self) -> Self::Item {
193+
// SAFETY: The caller promised the iterator isn't empty
194+
self.count = unsafe { self.count.unchecked_sub(1) };
195+
if self.count == 0 {
196+
// SAFETY: the check above ensured that the count used to be non-zero,
197+
// so element hasn't been dropped yet, and we just lowered the count to
198+
// zero so it won't be dropped later, and thus it's okay to take it here.
199+
unsafe { ManuallyDrop::take(&mut self.element) }
200+
} else {
201+
A::clone(&self.element)
202+
}
203+
}
204+
}

0 commit comments

Comments
 (0)