From 77ca30f1950a062e9eb1caeff790618488058057 Mon Sep 17 00:00:00 2001
From: Scott McMurray <scottmcm@users.noreply.github.com>
Date: Thu, 1 Aug 2024 21:58:34 -0700
Subject: [PATCH] Implement `UncheckedIterator` directly for `RepeatN`

---
 library/core/src/iter/sources/repeat_n.rs   | 33 +++++++++++++--------
 tests/codegen/iter-repeat-n-trivial-drop.rs | 15 +++++++++-
 2 files changed, 34 insertions(+), 14 deletions(-)

diff --git a/library/core/src/iter/sources/repeat_n.rs b/library/core/src/iter/sources/repeat_n.rs
index 8390dab8e543e..4c4ae39f836ca 100644
--- a/library/core/src/iter/sources/repeat_n.rs
+++ b/library/core/src/iter/sources/repeat_n.rs
@@ -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]
@@ -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)
+        }
+    }
+}
diff --git a/tests/codegen/iter-repeat-n-trivial-drop.rs b/tests/codegen/iter-repeat-n-trivial-drop.rs
index 31020b77984ba..7de224b92d853 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<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)
+}