Auto merge of rust-lang#91512 - scottmcm:array-intoiter-advance, r=Ma…

…rk-Simulacrum

Override `Iterator::advance(_back)_by` for `array::IntoIter`

Because I happened to notice that `nth` is currently getting codegen'd as a loop even for `Copy` types: <https://rust.godbolt.org/z/fPqv7Gvs7>

<details>
<summary>LLVM before and after</summary>

Rust:

```rust
#[no_mangle]
pub fn array_intoiter_nth(it: &mut std::array::IntoIter<i32, 100>, n: usize) -> Option<i32> {
    it.nth(n)
}
```

Current nightly:
```llvmir
define { i32, i32 } `@array_intoiter_nth(%"core::array::iter::IntoIter<i32,` 100_usize>"* noalias nocapture align 8 dereferenceable(416) %it, i64 %n) unnamed_addr #0 personality i32 (i32, i32, i64, %"unwind::libunwind::_Unwind_Exception"*, %"unwind::libunwind::_Unwind_Context"*)* `@rust_eh_personality` !dbg !6 {
start:
  %_3.i.i.i4.i.i = getelementptr inbounds %"core::array::iter::IntoIter<i32, 100_usize>", %"core::array::iter::IntoIter<i32, 100_usize>"* %it, i64 0, i32 0, i32 0
  %_4.i.i.i5.i.i = getelementptr inbounds %"core::array::iter::IntoIter<i32, 100_usize>", %"core::array::iter::IntoIter<i32, 100_usize>"* %it, i64 0, i32 0, i32 1
  %_4.i.i.i.i.i.i = load i64, i64* %_4.i.i.i5.i.i, align 8, !alias.scope !10
  %.not.i.i = icmp eq i64 %n, 0, !dbg !15
  %_3.i.i.i.i.pre.i = load i64, i64* %_3.i.i.i4.i.i, align 8, !dbg !40, !alias.scope !41
  br i1 %.not.i.i, label %bb4.i, label %bb4.preheader.i.i, !dbg !42

bb4.preheader.i.i:                                ; preds = %start
  %umax.i = tail call i64 `@llvm.umax.i64(i64` %_3.i.i.i.i.pre.i, i64 %_4.i.i.i.i.i.i) #3, !dbg !43
  %0 = sub i64 %umax.i, %_3.i.i.i.i.pre.i, !dbg !43
  br label %bb4.i.i, !dbg !43

bb4.i.i:                                          ; preds = %bb3.i.i.i.i, %bb4.preheader.i.i
  %_3.i.i.i.i.i.i = phi i64 [ %2, %bb3.i.i.i.i ], [ %_3.i.i.i.i.pre.i, %bb4.preheader.i.i ], !dbg !52
  %iter.sroa.0.016.i.i = phi i64 [ %1, %bb3.i.i.i.i ], [ 0, %bb4.preheader.i.i ]
  %1 = add nuw i64 %iter.sroa.0.016.i.i, 1, !dbg !54
  %exitcond.not.i = icmp eq i64 %iter.sroa.0.016.i.i, %0, !dbg !52
  br i1 %exitcond.not.i, label %core::iter::traits::iterator::Iterator::nth.exit, label %bb3.i.i.i.i, !dbg !43

bb3.i.i.i.i:                                      ; preds = %bb4.i.i
  %2 = add nuw i64 %_3.i.i.i.i.i.i, 1, !dbg !63
  store i64 %2, i64* %_3.i.i.i4.i.i, align 8, !dbg !66, !alias.scope !75
  %exitcond.not.i.i = icmp eq i64 %1, %n, !dbg !15
  br i1 %exitcond.not.i.i, label %bb4.i, label %bb4.i.i, !dbg !42

bb4.i:                                            ; preds = %bb3.i.i.i.i, %start
  %_3.i.i.i.i.i = phi i64 [ %_3.i.i.i.i.pre.i, %start ], [ %2, %bb3.i.i.i.i ], !dbg !84
  %3 = icmp ult i64 %_3.i.i.i.i.i, %_4.i.i.i.i.i.i, !dbg !84
  br i1 %3, label %bb3.i.i.i, label %core::iter::traits::iterator::Iterator::nth.exit, !dbg !89

bb3.i.i.i:                                        ; preds = %bb4.i
  %4 = add nuw i64 %_3.i.i.i.i.i, 1, !dbg !90
  store i64 %4, i64* %_3.i.i.i4.i.i, align 8, !dbg !93, !alias.scope !96
  %5 = getelementptr inbounds %"core::array::iter::IntoIter<i32, 100_usize>", %"core::array::iter::IntoIter<i32, 100_usize>"* %it, i64 0, i32 1, i64 %_3.i.i.i.i.i, !dbg !105
  %6 = load i32, i32* %5, align 4, !dbg !131, !alias.scope !141, !noalias !144
  br label %core::iter::traits::iterator::Iterator::nth.exit, !dbg !149

core::iter::traits::iterator::Iterator::nth.exit: ; preds = %bb4.i.i, %bb4.i, %bb3.i.i.i
  %.sroa.3.0.i = phi i32 [ %6, %bb3.i.i.i ], [ undef, %bb4.i ], [ undef, %bb4.i.i ], !dbg !40
  %.sroa.0.0.i = phi i32 [ 1, %bb3.i.i.i ], [ 0, %bb4.i ], [ 0, %bb4.i.i ], !dbg !40
  %7 = insertvalue { i32, i32 } undef, i32 %.sroa.0.0.i, 0, !dbg !150
  %8 = insertvalue { i32, i32 } %7, i32 %.sroa.3.0.i, 1, !dbg !150
  ret { i32, i32 } %8, !dbg !151
}
```

With this PR:
```llvmir
define { i32, i32 } `@array_intoiter_nth(%"core::array::iter::IntoIter<i32,` 100_usize>"* noalias nocapture align 8 dereferenceable(416) %it, i64 %n) unnamed_addr #0 personality i32 (...)* `@__CxxFrameHandler3` {
start:
  %0 = getelementptr inbounds %"core::array::iter::IntoIter<i32, 100_usize>", %"core::array::iter::IntoIter<i32, 100_usize>"* %it, i64 0, i32 0, i32 1
  %_2.i.i.i.i = load i64, i64* %0, align 8, !alias.scope !6, !noalias !13
  %1 = getelementptr inbounds %"core::array::iter::IntoIter<i32, 100_usize>", %"core::array::iter::IntoIter<i32, 100_usize>"* %it, i64 0, i32 0, i32 0
  %_3.i.i.i.i = load i64, i64* %1, align 8, !alias.scope !16
  %2 = sub i64 %_2.i.i.i.i, %_3.i.i.i.i
  %3 = icmp ult i64 %2, %n
  %.0.sroa.speculated.i.i.i.i.i = select i1 %3, i64 %2, i64 %n
  %_10.i.i = add i64 %.0.sroa.speculated.i.i.i.i.i, %_3.i.i.i.i
  store i64 %_10.i.i, i64* %1, align 8, !alias.scope !16
  %.not.i = xor i1 %3, true
  %4 = icmp ult i64 %_10.i.i, %_2.i.i.i.i
  %or.cond.i = select i1 %.not.i, i1 %4, i1 false
  br i1 %or.cond.i, label %bb3.i.i.i, label %_ZN4core4iter6traits8iterator8Iterator3nth17hcbc727011e9e2a3bE.exit

bb3.i.i.i:                                        ; preds = %start
  %5 = add nuw i64 %_10.i.i, 1
  store i64 %5, i64* %1, align 8, !alias.scope !17
  %6 = getelementptr inbounds %"core::array::iter::IntoIter<i32, 100_usize>", %"core::array::iter::IntoIter<i32, 100_usize>"* %it, i64 0, i32 1, i64 %_10.i.i
  %7 = load i32, i32* %6, align 4, !alias.scope !26, !noalias !29
  br label %_ZN4core4iter6traits8iterator8Iterator3nth17hcbc727011e9e2a3bE.exit

_ZN4core4iter6traits8iterator8Iterator3nth17hcbc727011e9e2a3bE.exit: ; preds = %start, %bb3.i.i.i
  %.sroa.3.0.i = phi i32 [ undef, %start ], [ %7, %bb3.i.i.i ]
  %.sroa.0.0.i = phi i32 [ 0, %start ], [ 1, %bb3.i.i.i ]
  %8 = insertvalue { i32, i32 } undef, i32 %.sroa.0.0.i, 0
  %9 = insertvalue { i32, i32 } %8, i32 %.sroa.3.0.i, 1
  ret { i32, i32 } %9
}
```
</details>

library/core/src/array/iter.rs

@@ -1,7 +1,7 @@
 //! Defines the `IntoIter` owned iterator for arrays.
 
 use crate::{
-    fmt,
+    cmp, fmt,
     iter::{self, ExactSizeIterator, FusedIterator, TrustedLen},
     mem::{self, MaybeUninit},
     ops::Range,
@@ -281,6 +281,27 @@ impl<T, const N: usize> Iterator for IntoIter<T, N> {
     fn last(mut self) -> Option<Self::Item> {
         self.next_back()
     }
+
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        let len = self.len();
+
+        // The number of elements to drop.  Always in-bounds by construction.
+        let delta = cmp::min(n, len);
+
+        let range_to_drop = self.alive.start..(self.alive.start + delta);
+
+        // Moving the start marks them as conceptually "dropped", so if anything
+        // goes bad then our drop impl won't double-free them.
+        self.alive.start += delta;
+
+        // SAFETY: These elements are currently initialized, so it's fine to drop them.
+        unsafe {
+            let slice = self.data.get_unchecked_mut(range_to_drop);
+            ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(slice));
+        }
+
+        if n > len { Err(len) } else { Ok(()) }
+    }
 }
 
 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]
@@ -301,6 +322,27 @@ impl<T, const N: usize> DoubleEndedIterator for IntoIter<T, N> {
             unsafe { self.data.get_unchecked(idx).assume_init_read() }
         })
     }
+
+    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+        let len = self.len();
+
+        // The number of elements to drop.  Always in-bounds by construction.
+        let delta = cmp::min(n, len);
+
+        let range_to_drop = (self.alive.end - delta)..self.alive.end;
+
+        // Moving the end marks them as conceptually "dropped", so if anything
+        // goes bad then our drop impl won't double-free them.
+        self.alive.end -= delta;
+
+        // SAFETY: These elements are currently initialized, so it's fine to drop them.
+        unsafe {
+            let slice = self.data.get_unchecked_mut(range_to_drop);
+            ptr::drop_in_place(MaybeUninit::slice_assume_init_mut(slice));
+        }
+
+        if n > len { Err(len) } else { Ok(()) }
+    }
 }
 
 #[stable(feature = "array_value_iter_impls", since = "1.40.0")]

library/core/tests/array.rs

@@ -474,3 +474,109 @@ fn array_split_array_mut_out_of_bounds() {
 
     v.split_array_mut::<7>();
 }
+
+#[test]
+fn array_intoiter_advance_by() {
+    use std::cell::Cell;
+    struct DropCounter<'a>(usize, &'a Cell<usize>);
+    impl Drop for DropCounter<'_> {
+        fn drop(&mut self) {
+            let x = self.1.get();
+            self.1.set(x + 1);
+        }
+    }
+
+    let counter = Cell::new(0);
+    let a: [_; 100] = std::array::from_fn(|i| DropCounter(i, &counter));
+    let mut it = IntoIterator::into_iter(a);
+
+    let r = it.advance_by(1);
+    assert_eq!(r, Ok(()));
+    assert_eq!(it.len(), 99);
+    assert_eq!(counter.get(), 1);
+
+    let r = it.advance_by(0);
+    assert_eq!(r, Ok(()));
+    assert_eq!(it.len(), 99);
+    assert_eq!(counter.get(), 1);
+
+    let r = it.advance_by(11);
+    assert_eq!(r, Ok(()));
+    assert_eq!(it.len(), 88);
+    assert_eq!(counter.get(), 12);
+
+    let x = it.next();
+    assert_eq!(x.as_ref().map(|x| x.0), Some(12));
+    assert_eq!(it.len(), 87);
+    assert_eq!(counter.get(), 12);
+    drop(x);
+    assert_eq!(counter.get(), 13);
+
+    let r = it.advance_by(123456);
+    assert_eq!(r, Err(87));
+    assert_eq!(it.len(), 0);
+    assert_eq!(counter.get(), 100);
+
+    let r = it.advance_by(0);
+    assert_eq!(r, Ok(()));
+    assert_eq!(it.len(), 0);
+    assert_eq!(counter.get(), 100);
+
+    let r = it.advance_by(10);
+    assert_eq!(r, Err(0));
+    assert_eq!(it.len(), 0);
+    assert_eq!(counter.get(), 100);
+}
+
+#[test]
+fn array_intoiter_advance_back_by() {
+    use std::cell::Cell;
+    struct DropCounter<'a>(usize, &'a Cell<usize>);
+    impl Drop for DropCounter<'_> {
+        fn drop(&mut self) {
+            let x = self.1.get();
+            self.1.set(x + 1);
+        }
+    }
+
+    let counter = Cell::new(0);
+    let a: [_; 100] = std::array::from_fn(|i| DropCounter(i, &counter));
+    let mut it = IntoIterator::into_iter(a);
+
+    let r = it.advance_back_by(1);
+    assert_eq!(r, Ok(()));
+    assert_eq!(it.len(), 99);
+    assert_eq!(counter.get(), 1);
+
+    let r = it.advance_back_by(0);
+    assert_eq!(r, Ok(()));
+    assert_eq!(it.len(), 99);
+    assert_eq!(counter.get(), 1);
+
+    let r = it.advance_back_by(11);
+    assert_eq!(r, Ok(()));
+    assert_eq!(it.len(), 88);
+    assert_eq!(counter.get(), 12);
+
+    let x = it.next_back();
+    assert_eq!(x.as_ref().map(|x| x.0), Some(87));
+    assert_eq!(it.len(), 87);
+    assert_eq!(counter.get(), 12);
+    drop(x);
+    assert_eq!(counter.get(), 13);
+
+    let r = it.advance_back_by(123456);
+    assert_eq!(r, Err(87));
+    assert_eq!(it.len(), 0);
+    assert_eq!(counter.get(), 100);
+
+    let r = it.advance_back_by(0);
+    assert_eq!(r, Ok(()));
+    assert_eq!(it.len(), 0);
+    assert_eq!(counter.get(), 100);
+
+    let r = it.advance_back_by(10);
+    assert_eq!(r, Err(0));
+    assert_eq!(it.len(), 0);
+    assert_eq!(counter.get(), 100);
+}