Skip to content

Commit

Permalink
Auto merge of rust-lang#78818 - scottmcm:as_rchunks, r=KodrAus
Browse files Browse the repository at this point in the history 
Add `as_rchunks` (and friends) to slices

`@est31` mentioned (rust-lang#76354 (comment)) that, for completeness, there needed to be an `as_chunks`-like method that chunks from the end (with the remainder at the beginning) like `rchunks` does.

So here's a PR for `as_rchunks: &[T] -> (&[T], &[[T; N]])` and `as_rchunks_mut: &mut [T] -> (&mut [T], &mut [[T; N]])`.

But as I was doing this and copy-pasting `from_raw_parts` calls, I thought that I should extract that into an unsafe method.  It started out a private helper, but it seemed like `as_chunks_unchecked` could be reasonable as a "real" method, so I added docs and made it public.  Let me know if you think it doesn't pull its weight.
  • Loading branch information
@bors
bors committed Jan 17, 2021
2 parents 95cbcad + 6bcaba9 commit 49d7889
Show file tree
Hide file tree
Showing 2 changed files with 187 additions and 7 deletions.
161 changes: 154 additions & 7 deletions library/core/src/slice/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -888,6 +888,46 @@ impl<T> [T] {
ChunksExactMut::new(self, chunk_size)
}

/// Splits the slice into a slice of `N`-element arrays,
/// assuming that there's no remainder.
///
/// # Safety
///
/// This may only be called when
/// - The slice splits exactly into `N`-element chunks (aka `self.len() % N == 0`).
/// - `N != 0`.
///
/// # Examples
///
/// ```
/// #![feature(slice_as_chunks)]
/// let slice: &[char] = &['l', 'o', 'r', 'e', 'm', '!'];
/// let chunks: &[[char; 1]] =
/// // SAFETY: 1-element chunks never have remainder
/// unsafe { slice.as_chunks_unchecked() };
/// assert_eq!(chunks, &[['l'], ['o'], ['r'], ['e'], ['m'], ['!']]);
/// let chunks: &[[char; 3]] =
/// // SAFETY: The slice length (6) is a multiple of 3
/// unsafe { slice.as_chunks_unchecked() };
/// assert_eq!(chunks, &[['l', 'o', 'r'], ['e', 'm', '!']]);
///
/// // These would be unsound:
/// // let chunks: &[[_; 5]] = slice.as_chunks_unchecked() // The slice length is not a multiple of 5
/// // let chunks: &[[_; 0]] = slice.as_chunks_unchecked() // Zero-length chunks are never allowed
/// ```
#[unstable(feature = "slice_as_chunks", issue = "74985")]
#[inline]
pub unsafe fn as_chunks_unchecked<const N: usize>(&self) -> &[[T; N]] {
debug_assert_ne!(N, 0);
debug_assert_eq!(self.len() % N, 0);
let new_len =
// SAFETY: Our precondition is exactly what's needed to call this
unsafe { crate::intrinsics::exact_div(self.len(), N) };
// SAFETY: We cast a slice of `new_len * N` elements into
// a slice of `new_len` many `N` elements chunks.
unsafe { from_raw_parts(self.as_ptr().cast(), new_len) }
}

/// Splits the slice into a slice of `N`-element arrays,
/// starting at the beginning of the slice,
/// and a remainder slice with length strictly less than `N`.
Expand All @@ -912,12 +952,42 @@ impl<T> [T] {
assert_ne!(N, 0);
let len = self.len() / N;
let (multiple_of_n, remainder) = self.split_at(len * N);
// SAFETY: We cast a slice of `len * N` elements into
// a slice of `len` many `N` elements chunks.
let array_slice: &[[T; N]] = unsafe { from_raw_parts(multiple_of_n.as_ptr().cast(), len) };
// SAFETY: We already panicked for zero, and ensured by construction
// that the length of the subslice is a multiple of N.
let array_slice = unsafe { multiple_of_n.as_chunks_unchecked() };
(array_slice, remainder)
}

/// Splits the slice into a slice of `N`-element arrays,
/// starting at the end of the slice,
/// and a remainder slice with length strictly less than `N`.
///
/// # 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(slice_as_chunks)]
/// let slice = ['l', 'o', 'r', 'e', 'm'];
/// let (remainder, chunks) = slice.as_rchunks();
/// assert_eq!(remainder, &['l']);
/// assert_eq!(chunks, &[['o', 'r'], ['e', 'm']]);
/// ```
#[unstable(feature = "slice_as_chunks", issue = "74985")]
#[inline]
pub fn as_rchunks<const N: usize>(&self) -> (&[T], &[[T; N]]) {
assert_ne!(N, 0);
let len = self.len() / N;
let (remainder, multiple_of_n) = self.split_at(self.len() - len * N);
// SAFETY: We already panicked for zero, and ensured by construction
// that the length of the subslice is a multiple of N.
let array_slice = unsafe { multiple_of_n.as_chunks_unchecked() };
(remainder, array_slice)
}

/// Returns an iterator over `N` elements of the slice at a time, starting at the
/// beginning of the slice.
///
Expand Down Expand Up @@ -952,6 +1022,48 @@ impl<T> [T] {
ArrayChunks::new(self)
}

/// Splits the slice into a slice of `N`-element arrays,
/// assuming that there's no remainder.
///
/// # Safety
///
/// This may only be called when
/// - The slice splits exactly into `N`-element chunks (aka `self.len() % N == 0`).
/// - `N != 0`.
///
/// # Examples
///
/// ```
/// #![feature(slice_as_chunks)]
/// let slice: &mut [char] = &mut ['l', 'o', 'r', 'e', 'm', '!'];
/// let chunks: &mut [[char; 1]] =
/// // SAFETY: 1-element chunks never have remainder
/// unsafe { slice.as_chunks_unchecked_mut() };
/// chunks[0] = ['L'];
/// assert_eq!(chunks, &[['L'], ['o'], ['r'], ['e'], ['m'], ['!']]);
/// let chunks: &mut [[char; 3]] =
/// // SAFETY: The slice length (6) is a multiple of 3
/// unsafe { slice.as_chunks_unchecked_mut() };
/// chunks[1] = ['a', 'x', '?'];
/// assert_eq!(slice, &['L', 'o', 'r', 'a', 'x', '?']);
///
/// // These would be unsound:
/// // let chunks: &[[_; 5]] = slice.as_chunks_unchecked_mut() // The slice length is not a multiple of 5
/// // let chunks: &[[_; 0]] = slice.as_chunks_unchecked_mut() // Zero-length chunks are never allowed
/// ```
#[unstable(feature = "slice_as_chunks", issue = "74985")]
#[inline]
pub unsafe fn as_chunks_unchecked_mut<const N: usize>(&mut self) -> &mut [[T; N]] {
debug_assert_ne!(N, 0);
debug_assert_eq!(self.len() % N, 0);
let new_len =
// SAFETY: Our precondition is exactly what's needed to call this
unsafe { crate::intrinsics::exact_div(self.len(), N) };
// SAFETY: We cast a slice of `new_len * N` elements into
// a slice of `new_len` many `N` elements chunks.
unsafe { from_raw_parts_mut(self.as_mut_ptr().cast(), new_len) }
}

/// Splits the slice into a slice of `N`-element arrays,
/// starting at the beginning of the slice,
/// and a remainder slice with length strictly less than `N`.
Expand Down Expand Up @@ -982,13 +1094,48 @@ impl<T> [T] {
assert_ne!(N, 0);
let len = self.len() / N;
let (multiple_of_n, remainder) = self.split_at_mut(len * N);
let array_slice: &mut [[T; N]] =
// SAFETY: We cast a slice of `len * N` elements into
// a slice of `len` many `N` elements chunks.
unsafe { from_raw_parts_mut(multiple_of_n.as_mut_ptr().cast(), len) };
// SAFETY: We already panicked for zero, and ensured by construction
// that the length of the subslice is a multiple of N.
let array_slice = unsafe { multiple_of_n.as_chunks_unchecked_mut() };
(array_slice, remainder)
}

/// Splits the slice into a slice of `N`-element arrays,
/// starting at the end of the slice,
/// and a remainder slice with length strictly less than `N`.
///
/// # 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(slice_as_chunks)]
/// let v = &mut [0, 0, 0, 0, 0];
/// let mut count = 1;
///
/// let (remainder, chunks) = v.as_rchunks_mut();
/// remainder[0] = 9;
/// for chunk in chunks {
/// *chunk = [count; 2];
/// count += 1;
/// }
/// assert_eq!(v, &[9, 1, 1, 2, 2]);
/// ```
#[unstable(feature = "slice_as_chunks", issue = "74985")]
#[inline]
pub fn as_rchunks_mut<const N: usize>(&mut self) -> (&mut [T], &mut [[T; N]]) {
assert_ne!(N, 0);
let len = self.len() / N;
let (remainder, multiple_of_n) = self.split_at_mut(self.len() - len * N);
// SAFETY: We already panicked for zero, and ensured by construction
// that the length of the subslice is a multiple of N.
let array_slice = unsafe { multiple_of_n.as_chunks_unchecked_mut() };
(remainder, array_slice)
}

/// Returns an iterator over `N` elements of the slice at a time, starting at the
/// beginning of the slice.
///
Expand Down
33 changes: 33 additions & 0 deletions src/test/codegen/slice-as_chunks.rs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
// no-system-llvm
// compile-flags: -O
// only-64bit (because the LLVM type of i64 for usize shows up)
// ignore-debug: the debug assertions get in the way

#![crate_type = "lib"]
#![feature(slice_as_chunks)]

// CHECK-LABEL: @chunks4
#[no_mangle]
pub fn chunks4(x: &[u8]) -> &[[u8; 4]] {
// CHECK-NEXT: start:
// CHECK-NEXT: lshr i64 %x.1, 2
// CHECK-NOT: shl
// CHECK-NOT: mul
// CHECK-NOT: udiv
// CHECK-NOT: urem
// CHECK: ret
x.as_chunks().0
}

// CHECK-LABEL: @chunks4_with_remainder
#[no_mangle]
pub fn chunks4_with_remainder(x: &[u8]) -> (&[[u8; 4]], &[u8]) {
// CHECK: and i64 %x.1, -4
// CHECK: and i64 %x.1, 3
// CHECK: lshr exact
// CHECK-NOT: mul
// CHECK-NOT: udiv
// CHECK-NOT: urem
// CHECK: ret
x.as_chunks()
}

0 comments on commit 49d7889

Please sign in to comment.