From f6a6d2fef6c090806e961acb89817be8329f8208 Mon Sep 17 00:00:00 2001 From: Josh Stone Date: Sat, 1 Aug 2020 11:07:12 -0700 Subject: [PATCH 1/5] Add slice::array_chunks_mut --- library/core/src/slice/mod.rs | 151 ++++++++++++++++++++++++++++++++-- 1 file changed, 146 insertions(+), 5 deletions(-) diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs index 68977a983aa37..18ff50b8dbd41 100644 --- a/library/core/src/slice/mod.rs +++ b/library/core/src/slice/mod.rs @@ -997,9 +997,9 @@ impl [T] { /// Returns an iterator over `N` elements of the slice at a time, starting at the /// beginning of the slice. /// - /// The chunks are slices and do not overlap. If `N` does not divide the length of the - /// slice, then the last up to `N-1` elements will be omitted and can be retrieved - /// from the `remainder` function of the iterator. + /// The chunks are array references and do not overlap. If `N` does not divide the + /// length of the slice, then the last up to `N-1` elements will be omitted and can be + /// retrieved from the `remainder` function of the iterator. /// /// This method is the const generic equivalent of [`chunks_exact`]. /// @@ -1033,6 +1033,51 @@ impl [T] { ArrayChunks { iter: array_slice.iter(), rem: snd } } + /// Returns an iterator over `N` elements of the slice at a time, starting at the + /// beginning of the slice. + /// + /// The chunks are mutable array references and do not overlap. If `N` does not divide + /// the length of the slice, then the last up to `N-1` elements will be omitted and + /// can be retrieved from the `into_remainder` function of the iterator. + /// + /// This method is the const generic equivalent of [`chunks_exact_mut`]. + /// + /// # Panics + /// + /// Panics if `N` is 0. This check will most probably get changed to a compile time + /// error before this method gets stabilized. + /// + /// # Examples + /// + /// ``` + /// #![feature(array_chunks)] + /// let v = &mut [0, 0, 0, 0, 0]; + /// let mut count = 1; + /// + /// for chunk in v.array_chunks_mut() { + /// *chunk = [count; 2]; + /// count += 1; + /// } + /// assert_eq!(v, &[1, 1, 2, 2, 0]); + /// ``` + /// + /// [`chunks_exact_mut`]: #method.chunks_exact_mut + #[unstable(feature = "array_chunks", issue = "74985")] + #[inline] + pub fn array_chunks_mut(&mut self) -> ArrayChunksMut<'_, T, N> { + assert_ne!(N, 0); + let len = self.len() / N; + let (fst_ptr, snd) = { + // Scope the first slice into a pointer to avoid aliasing the new slice below. + let (fst, snd) = self.split_at_mut(len * N); + (fst.as_mut_ptr(), snd) + }; + // SAFETY: We cast a slice of `len * N` elements into + // a slice of `len` many `N` elements chunks. + let array_slice: &mut [[T; N]] = unsafe { from_raw_parts_mut(fst_ptr.cast(), len) }; + ArrayChunksMut { iter: array_slice.iter_mut(), rem: snd } + } + /// Returns an iterator over `chunk_size` elements of the slice at a time, starting at the end /// of the slice. /// @@ -5826,7 +5871,7 @@ unsafe impl<'a, T> TrustedRandomAccess for ChunksExactMut<'a, T> { /// time), starting at the beginning of the slice. /// /// When the slice len is not evenly divided by the chunk size, the last -/// up to `chunk_size-1` elements will be omitted but can be retrieved from +/// up to `N-1` elements will be omitted but can be retrieved from /// the [`remainder`] function from the iterator. /// /// This struct is created by the [`array_chunks`] method on [slices]. @@ -5843,7 +5888,7 @@ pub struct ArrayChunks<'a, T: 'a, const N: usize> { impl<'a, T, const N: usize> ArrayChunks<'a, T, N> { /// Returns the remainder of the original slice that is not going to be - /// returned by the iterator. The returned slice has at most `chunk_size-1` + /// returned by the iterator. The returned slice has at most `N-1` /// elements. #[unstable(feature = "array_chunks", issue = "74985")] pub fn remainder(&self) -> &'a [T] { @@ -5929,6 +5974,102 @@ unsafe impl<'a, T, const N: usize> TrustedRandomAccess for ArrayChunks<'a, T, N> } } +/// An iterator over a slice in (non-overlapping) mutable chunks (`N` elements +/// at a time), starting at the beginning of the slice. +/// +/// When the slice len is not evenly divided by the chunk size, the last +/// up to `N-1` elements will be omitted but can be retrieved from +/// the [`into_remainder`] function from the iterator. +/// +/// This struct is created by the [`array_chunks_mut`] method on [slices]. +/// +/// [`array_chunks_mut`]: ../../std/primitive.slice.html#method.array_chunks_mut +/// [`into_remainder`]: ../../std/slice/struct.ArrayChunksMut.html#method.into_remainder +/// [slices]: ../../std/primitive.slice.html +#[derive(Debug)] +#[unstable(feature = "array_chunks", issue = "74985")] +pub struct ArrayChunksMut<'a, T: 'a, const N: usize> { + iter: IterMut<'a, [T; N]>, + rem: &'a mut [T], +} + +impl<'a, T, const N: usize> ArrayChunksMut<'a, T, N> { + /// Returns the remainder of the original slice that is not going to be + /// returned by the iterator. The returned slice has at most `N-1` + /// elements. + #[unstable(feature = "array_chunks", issue = "74985")] + pub fn into_remainder(self) -> &'a mut [T] { + self.rem + } +} + +#[unstable(feature = "array_chunks", issue = "74985")] +impl<'a, T, const N: usize> Iterator for ArrayChunksMut<'a, T, N> { + type Item = &'a mut [T; N]; + + #[inline] + fn next(&mut self) -> Option<&'a mut [T; N]> { + self.iter.next() + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + self.iter.size_hint() + } + + #[inline] + fn count(self) -> usize { + self.iter.count() + } + + #[inline] + fn nth(&mut self, n: usize) -> Option { + self.iter.nth(n) + } + + #[inline] + fn last(self) -> Option { + self.iter.last() + } +} + +#[unstable(feature = "array_chunks", issue = "74985")] +impl<'a, T, const N: usize> DoubleEndedIterator for ArrayChunksMut<'a, T, N> { + #[inline] + fn next_back(&mut self) -> Option<&'a mut [T; N]> { + self.iter.next_back() + } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option { + self.iter.nth_back(n) + } +} + +#[unstable(feature = "array_chunks", issue = "74985")] +impl ExactSizeIterator for ArrayChunksMut<'_, T, N> { + fn is_empty(&self) -> bool { + self.iter.is_empty() + } +} + +#[unstable(feature = "trusted_len", issue = "37572")] +unsafe impl TrustedLen for ArrayChunksMut<'_, T, N> {} + +#[unstable(feature = "array_chunks", issue = "74985")] +impl FusedIterator for ArrayChunksMut<'_, T, N> {} + +#[doc(hidden)] +#[unstable(feature = "array_chunks", issue = "74985")] +unsafe impl<'a, T, const N: usize> TrustedRandomAccess for ArrayChunksMut<'a, T, N> { + unsafe fn get_unchecked(&mut self, i: usize) -> &'a mut [T; N] { + unsafe { self.iter.get_unchecked(i) } + } + fn may_have_side_effect() -> bool { + false + } +} + /// An iterator over a slice in (non-overlapping) chunks (`chunk_size` elements at a /// time), starting at the end of the slice. /// From b9fd6734e8ce315f0825157c0bf3aad67521e664 Mon Sep 17 00:00:00 2001 From: Josh Stone Date: Sat, 1 Aug 2020 11:07:28 -0700 Subject: [PATCH 2/5] Add tests for array_chunks_mut --- library/core/tests/slice.rs | 93 +++++++++++++++++++++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/library/core/tests/slice.rs b/library/core/tests/slice.rs index 743df68699c28..9b31e532a6a9f 100644 --- a/library/core/tests/slice.rs +++ b/library/core/tests/slice.rs @@ -564,6 +564,99 @@ fn test_array_chunks_zip() { assert_eq!(res, vec![14, 22]); } +#[test] +fn test_array_chunks_mut_infer() { + let v: &mut [i32] = &mut [0, 1, 2, 3, 4, 5, 6]; + for a in v.array_chunks_mut() { + let sum = a.iter().sum::(); + *a = [sum; 3]; + } + assert_eq!(v, &[3, 3, 3, 12, 12, 12, 6]); + + let v2: &mut [i32] = &mut [0, 1, 2, 3, 4, 5, 6]; + v2.array_chunks_mut().for_each(|[a, b]| core::mem::swap(a, b)); + assert_eq!(v2, &[1, 0, 3, 2, 5, 4, 6]); +} + +#[test] +fn test_array_chunks_mut_count() { + let v: &mut [i32] = &mut [0, 1, 2, 3, 4, 5]; + let c = v.array_chunks_mut::<3>(); + assert_eq!(c.count(), 2); + + let v2: &mut [i32] = &mut [0, 1, 2, 3, 4]; + let c2 = v2.array_chunks_mut::<2>(); + assert_eq!(c2.count(), 2); + + let v3: &mut [i32] = &mut []; + let c3 = v3.array_chunks_mut::<2>(); + assert_eq!(c3.count(), 0); +} + +#[test] +fn test_array_chunks_mut_nth() { + let v: &mut [i32] = &mut [0, 1, 2, 3, 4, 5]; + let mut c = v.array_chunks_mut::<2>(); + assert_eq!(c.nth(1).unwrap(), &[2, 3]); + assert_eq!(c.next().unwrap(), &[4, 5]); + + let v2: &mut [i32] = &mut [0, 1, 2, 3, 4, 5, 6]; + let mut c2 = v2.array_chunks_mut::<3>(); + assert_eq!(c2.nth(1).unwrap(), &[3, 4, 5]); + assert_eq!(c2.next(), None); +} + +#[test] +fn test_array_chunks_mut_nth_back() { + let v: &mut [i32] = &mut [0, 1, 2, 3, 4, 5]; + let mut c = v.array_chunks_mut::<2>(); + assert_eq!(c.nth_back(1).unwrap(), &[2, 3]); + assert_eq!(c.next().unwrap(), &[0, 1]); + assert_eq!(c.next(), None); + + let v2: &mut [i32] = &mut [0, 1, 2, 3, 4]; + let mut c2 = v2.array_chunks_mut::<3>(); + assert_eq!(c2.nth_back(0).unwrap(), &[0, 1, 2]); + assert_eq!(c2.next(), None); + assert_eq!(c2.next_back(), None); + + let v3: &mut [i32] = &mut [0, 1, 2, 3, 4]; + let mut c3 = v3.array_chunks_mut::<10>(); + assert_eq!(c3.nth_back(0), None); +} + +#[test] +fn test_array_chunks_mut_last() { + let v: &mut [i32] = &mut [0, 1, 2, 3, 4, 5]; + let c = v.array_chunks_mut::<2>(); + assert_eq!(c.last().unwrap(), &[4, 5]); + + let v2: &mut [i32] = &mut [0, 1, 2, 3, 4]; + let c2 = v2.array_chunks_mut::<2>(); + assert_eq!(c2.last().unwrap(), &[2, 3]); +} + +#[test] +fn test_array_chunks_mut_remainder() { + let v: &mut [i32] = &mut [0, 1, 2, 3, 4]; + let c = v.array_chunks_mut::<2>(); + assert_eq!(c.into_remainder(), &[4]); +} + +#[test] +fn test_array_chunks_mut_zip() { + let v1: &mut [i32] = &mut [0, 1, 2, 3, 4]; + let v2: &[i32] = &[6, 7, 8, 9, 10]; + + for (a, b) in v1.array_chunks_mut::<2>().zip(v2.array_chunks::<2>()) { + let sum = b.iter().sum::(); + for v in a { + *v += sum; + } + } + assert_eq!(v1, [13, 14, 19, 20, 4]); +} + #[test] fn test_rchunks_count() { let v: &[i32] = &[0, 1, 2, 3, 4, 5]; From 864a28e01d70cd73ef7f0b64f6793697f638b8b0 Mon Sep 17 00:00:00 2001 From: Josh Stone Date: Sat, 1 Aug 2020 11:48:34 -0700 Subject: [PATCH 3/5] Re-export ArrayChunksMut in alloc --- library/alloc/src/slice.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/library/alloc/src/slice.rs b/library/alloc/src/slice.rs index 8ea2c6dc859b2..677bfdd2349ec 100644 --- a/library/alloc/src/slice.rs +++ b/library/alloc/src/slice.rs @@ -93,6 +93,8 @@ use crate::vec::Vec; #[unstable(feature = "array_chunks", issue = "74985")] pub use core::slice::ArrayChunks; +#[unstable(feature = "array_chunks", issue = "74985")] +pub use core::slice::ArrayChunksMut; #[stable(feature = "slice_get_slice", since = "1.28.0")] pub use core::slice::SliceIndex; #[stable(feature = "from_ref", since = "1.28.0")] From 21903532eee96a5311d08b8a9e9cc9f9231bc478 Mon Sep 17 00:00:00 2001 From: Josh Stone Date: Sun, 2 Aug 2020 10:35:57 -0700 Subject: [PATCH 4/5] Build the slice directly in array_chunks_mut Review discussion found that the concern about aliasing was overblown, so we can simplify this to cast from one slice to another directly. --- library/core/src/slice/mod.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs index 18ff50b8dbd41..a32ea062de494 100644 --- a/library/core/src/slice/mod.rs +++ b/library/core/src/slice/mod.rs @@ -1067,14 +1067,11 @@ impl [T] { pub fn array_chunks_mut(&mut self) -> ArrayChunksMut<'_, T, N> { assert_ne!(N, 0); let len = self.len() / N; - let (fst_ptr, snd) = { - // Scope the first slice into a pointer to avoid aliasing the new slice below. - let (fst, snd) = self.split_at_mut(len * N); - (fst.as_mut_ptr(), snd) - }; + let (fst, snd) = self.split_at_mut(len * N); // SAFETY: We cast a slice of `len * N` elements into // a slice of `len` many `N` elements chunks. - let array_slice: &mut [[T; N]] = unsafe { from_raw_parts_mut(fst_ptr.cast(), len) }; + let array_slice: &mut [[T; N]] = + unsafe { from_raw_parts_mut(fst.as_mut_ptr().cast(), len) }; ArrayChunksMut { iter: array_slice.iter_mut(), rem: snd } } From 86b9f710d0c90068866e736bbbb7b89ac93ff2e6 Mon Sep 17 00:00:00 2001 From: Josh Stone Date: Fri, 4 Sep 2020 20:08:12 -0700 Subject: [PATCH 5/5] Move ArrayChunksMut::get_unchecked per #73565 --- library/core/src/slice/mod.rs | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/library/core/src/slice/mod.rs b/library/core/src/slice/mod.rs index a32ea062de494..609fdbd599274 100644 --- a/library/core/src/slice/mod.rs +++ b/library/core/src/slice/mod.rs @@ -1070,9 +1070,10 @@ impl [T] { let (fst, snd) = self.split_at_mut(len * N); // SAFETY: We cast a slice of `len * N` elements into // a slice of `len` many `N` elements chunks. - let array_slice: &mut [[T; N]] = - unsafe { from_raw_parts_mut(fst.as_mut_ptr().cast(), len) }; - ArrayChunksMut { iter: array_slice.iter_mut(), rem: snd } + unsafe { + let array_slice: &mut [[T; N]] = from_raw_parts_mut(fst.as_mut_ptr().cast(), len); + ArrayChunksMut { iter: array_slice.iter_mut(), rem: snd } + } } /// Returns an iterator over `chunk_size` elements of the slice at a time, starting at the end @@ -6028,6 +6029,12 @@ impl<'a, T, const N: usize> Iterator for ArrayChunksMut<'a, T, N> { fn last(self) -> Option { self.iter.last() } + + unsafe fn get_unchecked(&mut self, i: usize) -> &'a mut [T; N] { + // SAFETY: The safety guarantees of `get_unchecked` are transferred to + // the caller. + unsafe { self.iter.get_unchecked(i) } + } } #[unstable(feature = "array_chunks", issue = "74985")] @@ -6059,9 +6066,6 @@ impl FusedIterator for ArrayChunksMut<'_, T, N> {} #[doc(hidden)] #[unstable(feature = "array_chunks", issue = "74985")] unsafe impl<'a, T, const N: usize> TrustedRandomAccess for ArrayChunksMut<'a, T, N> { - unsafe fn get_unchecked(&mut self, i: usize) -> &'a mut [T; N] { - unsafe { self.iter.get_unchecked(i) } - } fn may_have_side_effect() -> bool { false }