Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add as_rchunks (and friends) to slices #78818

Merged
merged 3 commits into from
Jan 17, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
161 changes: 154 additions & 7 deletions library/core/src/slice/mod.rs
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
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
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wanted to ensure I didn't accidentally introduce extra instructions with the refactor, so codegen test. It did change, but only from lshr/and/and to and/and/lshr exact, which is a thoroughly uninteresting difference.

// CHECK-NOT: mul
// CHECK-NOT: udiv
// CHECK-NOT: urem
// CHECK: ret
x.as_chunks()
}