Skip to content

Commit be145f3

Browse files
authored
Unrolled build for #128530
Rollup merge of #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 bbf60c8 + 77ca30f commit be145f3

File tree

2 files changed

+34
-14
lines changed

2 files changed

+34
-14
lines changed

Diff for: library/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+
}

Diff for: tests/codegen/iter-repeat-n-trivial-drop.rs

+14-1
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
1-
//@ compile-flags: -O
1+
//@ compile-flags: -C opt-level=3
22
//@ only-x86_64
33

44
#![crate_type = "lib"]
55
#![feature(iter_repeat_n)]
6+
#![feature(array_repeat)]
67

78
#[derive(Clone)]
89
pub struct NotCopy(u16);
@@ -54,3 +55,15 @@ pub fn vec_extend_via_iter_repeat_n() -> Vec<u8> {
5455
v.extend(std::iter::repeat_n(42_u8, n));
5556
v
5657
}
58+
59+
// Array repeat uses `RepeatN::next_unchecked` internally,
60+
// so also check that the distinction disappears there.
61+
62+
#[no_mangle]
63+
// CHECK-LABEL: @array_repeat_not_copy
64+
pub unsafe fn array_repeat_not_copy(item: NotCopy) -> [NotCopy; 8] {
65+
// CHECK: insertelement {{.+}} i16 %item
66+
// CHECK: shufflevector <8 x i16> {{.+}} zeroinitializer
67+
// CHECK: store <8 x i16>
68+
std::array::repeat(item)
69+
}

0 commit comments

Comments
 (0)