diff --git a/library/core/src/array/iter.rs b/library/core/src/array/iter.rs index 5d63cf03fcbc4..a146966c674d4 100644 --- a/library/core/src/array/iter.rs +++ b/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, @@ -150,6 +150,27 @@ impl Iterator for IntoIter { fn last(mut self) -> Option { 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")] @@ -170,6 +191,27 @@ impl DoubleEndedIterator for IntoIter { 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")] diff --git a/library/core/tests/array.rs b/library/core/tests/array.rs index 43d30d12355eb..d212a3a3a05da 100644 --- a/library/core/tests/array.rs +++ b/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); + 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); + 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); +}