Stop using `into_iter` in `array::map`

2 files changed, 74 insertions, 11 deletions

diff --git a/library/core/src/array/drain.rs b/library/core/src/array/drain.rs
new file mode 100644
index 00000000000..5ca93d54f87
--- /dev/null
+++ b/library/core/src/array/drain.rs
@@ -0,0 +1,51 @@
+use crate::iter::TrustedLen;
+use crate::mem::ManuallyDrop;
+use crate::ptr::drop_in_place;
+use crate::slice;
+
+// INVARIANT: It's ok to drop the remainder of the inner iterator.
+pub(crate) struct Drain<'a, T>(slice::IterMut<'a, T>);
+
+pub(crate) fn drain_array_with<T, R, const N: usize>(
+    array: [T; N],
+    func: impl for<'a> FnOnce(Drain<'a, T>) -> R,
+) -> R {
+    let mut array = ManuallyDrop::new(array);
+    // SAFETY: Now that the local won't drop it, it's ok to construct the `Drain` which will.
+    let drain = Drain(array.iter_mut());
+    func(drain)
+}
+
+impl<T> Drop for Drain<'_, T> {
+    fn drop(&mut self) {
+        // SAFETY: By the type invariant, we're allowed to drop all these.
+        unsafe { drop_in_place(self.0.as_mut_slice()) }
+    }
+}
+
+impl<T> Iterator for Drain<'_, T> {
+    type Item = T;
+
+    #[inline]
+    fn next(&mut self) -> Option<T> {
+        let p: *const T = self.0.next()?;
+        // SAFETY: The iterator was already advanced, so we won't drop this later.
+        Some(unsafe { p.read() })
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let n = self.len();
+        (n, Some(n))
+    }
+}
+
+impl<T> ExactSizeIterator for Drain<'_, T> {
+    #[inline]
+    fn len(&self) -> usize {
+        self.0.len()
+    }
+}
+
+// SAFETY: This is a 1:1 wrapper for a slice iterator, which is also `TrustedLen`.
+unsafe impl<T> TrustedLen for Drain<'_, T> {}
diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs
index 2825e0bbb43..ee340f38543 100644
--- a/library/core/src/array/mod.rs
+++ b/library/core/src/array/mod.rs
@@ -17,9 +17,12 @@ use crate::ops::{
 };
 use crate::slice::{Iter, IterMut};
 
+mod drain;
 mod equality;
 mod iter;
 
+pub(crate) use drain::drain_array_with;
+
 #[stable(feature = "array_value_iter", since = "1.51.0")]
 pub use iter::IntoIter;
 
@@ -513,9 +516,12 @@ impl<T, const N: usize> [T; N] {
     where
         F: FnMut(T) -> U,
     {
-        // SAFETY: we know for certain that this iterator will yield exactly `N`
-        // items.
-        unsafe { collect_into_array_unchecked(&mut IntoIterator::into_iter(self).map(f)) }
+        drain_array_with(self, |iter| {
+            let mut iter = iter.map(f);
+            // SAFETY: we know for certain that this iterator will yield exactly `N`
+            // items.
+            unsafe { collect_into_array_unchecked(&mut iter) }
+        })
     }
 
     /// A fallible function `f` applied to each element on array `self` in order to
@@ -552,9 +558,12 @@ impl<T, const N: usize> [T; N] {
         R: Try,
         R::Residual: Residual<[R::Output; N]>,
     {
-        // SAFETY: we know for certain that this iterator will yield exactly `N`
-        // items.
-        unsafe { try_collect_into_array_unchecked(&mut IntoIterator::into_iter(self).map(f)) }
+        drain_array_with(self, |iter| {
+            let mut iter = iter.map(f);
+            // SAFETY: we know for certain that this iterator will yield exactly `N`
+            // items.
+            unsafe { try_collect_into_array_unchecked(&mut iter) }
+        })
     }
 
     /// 'Zips up' two arrays into a single array of pairs.
@@ -575,11 +584,14 @@ impl<T, const N: usize> [T; N] {
     /// ```
     #[unstable(feature = "array_zip", issue = "80094")]
     pub fn zip<U>(self, rhs: [U; N]) -> [(T, U); N] {
-        let mut iter = IntoIterator::into_iter(self).zip(rhs);
-
-        // SAFETY: we know for certain that this iterator will yield exactly `N`
-        // items.
-        unsafe { collect_into_array_unchecked(&mut iter) }
+        drain_array_with(self, |lhs| {
+            drain_array_with(rhs, |rhs| {
+                let mut iter = crate::iter::zip(lhs, rhs);
+                // SAFETY: we know for certain that this iterator will yield exactly `N`
+                // items.
+                unsafe { collect_into_array_unchecked(&mut iter) }
+            })
+        })
     }
 
     /// Returns a slice containing the entire array. Equivalent to `&s[..]`.