Skip to content

Commit 49d7889

Browse files
borsbors
committed
Auto merge of #78818 - scottmcm:as_rchunks, r=KodrAus
Add `as_rchunks` (and friends) to slices `@est31` mentioned (#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.
2 parents 95cbcad + 6bcaba9 commit 49d7889

File tree

2 files changed

+187
-7
lines changed

2 files changed

+187
-7
lines changed

‎library/core/src/slice/mod.rs‎

Lines changed: 154 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -888,6 +888,46 @@ impl<T> [T] {
888888
ChunksExactMut::new(self, chunk_size)
889889
}
890890

891+
/// Splits the slice into a slice of `N`-element arrays,
892+
/// assuming that there's no remainder.
893+
///
894+
/// # Safety
895+
///
896+
/// This may only be called when
897+
/// - The slice splits exactly into `N`-element chunks (aka `self.len() % N == 0`).
898+
/// - `N != 0`.
899+
///
900+
/// # Examples
901+
///
902+
/// ```
903+
/// #![feature(slice_as_chunks)]
904+
/// let slice: &[char] = &['l', 'o', 'r', 'e', 'm', '!'];
905+
/// let chunks: &[[char; 1]] =
906+
/// // SAFETY: 1-element chunks never have remainder
907+
/// unsafe { slice.as_chunks_unchecked() };
908+
/// assert_eq!(chunks, &[['l'], ['o'], ['r'], ['e'], ['m'], ['!']]);
909+
/// let chunks: &[[char; 3]] =
910+
/// // SAFETY: The slice length (6) is a multiple of 3
911+
/// unsafe { slice.as_chunks_unchecked() };
912+
/// assert_eq!(chunks, &[['l', 'o', 'r'], ['e', 'm', '!']]);
913+
///
914+
/// // These would be unsound:
915+
/// // let chunks: &[[_; 5]] = slice.as_chunks_unchecked() // The slice length is not a multiple of 5
916+
/// // let chunks: &[[_; 0]] = slice.as_chunks_unchecked() // Zero-length chunks are never allowed
917+
/// ```
918+
#[unstable(feature = "slice_as_chunks", issue = "74985")]
919+
#[inline]
920+
pub unsafe fn as_chunks_unchecked<const N: usize>(&self) -> &[[T; N]] {
921+
debug_assert_ne!(N, 0);
922+
debug_assert_eq!(self.len() % N, 0);
923+
let new_len =
924+
// SAFETY: Our precondition is exactly what's needed to call this
925+
unsafe { crate::intrinsics::exact_div(self.len(), N) };
926+
// SAFETY: We cast a slice of `new_len * N` elements into
927+
// a slice of `new_len` many `N` elements chunks.
928+
unsafe { from_raw_parts(self.as_ptr().cast(), new_len) }
929+
}
930+
891931
/// Splits the slice into a slice of `N`-element arrays,
892932
/// starting at the beginning of the slice,
893933
/// and a remainder slice with length strictly less than `N`.
@@ -912,12 +952,42 @@ impl<T> [T] {
912952
assert_ne!(N, 0);
913953
let len = self.len() / N;
914954
let (multiple_of_n, remainder) = self.split_at(len * N);
915-
// SAFETY: We cast a slice of `len * N` elements into
916-
// a slice of `len` many `N` elements chunks.
917-
let array_slice: &[[T; N]] = unsafe { from_raw_parts(multiple_of_n.as_ptr().cast(), len) };
955+
// SAFETY: We already panicked for zero, and ensured by construction
956+
// that the length of the subslice is a multiple of N.
957+
let array_slice = unsafe { multiple_of_n.as_chunks_unchecked() };
918958
(array_slice, remainder)
919959
}
920960

961+
/// Splits the slice into a slice of `N`-element arrays,
962+
/// starting at the end of the slice,
963+
/// and a remainder slice with length strictly less than `N`.
964+
///
965+
/// # Panics
966+
///
967+
/// Panics if `N` is 0. This check will most probably get changed to a compile time
968+
/// error before this method gets stabilized.
969+
///
970+
/// # Examples
971+
///
972+
/// ```
973+
/// #![feature(slice_as_chunks)]
974+
/// let slice = ['l', 'o', 'r', 'e', 'm'];
975+
/// let (remainder, chunks) = slice.as_rchunks();
976+
/// assert_eq!(remainder, &['l']);
977+
/// assert_eq!(chunks, &[['o', 'r'], ['e', 'm']]);
978+
/// ```
979+
#[unstable(feature = "slice_as_chunks", issue = "74985")]
980+
#[inline]
981+
pub fn as_rchunks<const N: usize>(&self) -> (&[T], &[[T; N]]) {
982+
assert_ne!(N, 0);
983+
let len = self.len() / N;
984+
let (remainder, multiple_of_n) = self.split_at(self.len() - len * N);
985+
// SAFETY: We already panicked for zero, and ensured by construction
986+
// that the length of the subslice is a multiple of N.
987+
let array_slice = unsafe { multiple_of_n.as_chunks_unchecked() };
988+
(remainder, array_slice)
989+
}
990+
921991
/// Returns an iterator over `N` elements of the slice at a time, starting at the
922992
/// beginning of the slice.
923993
///
@@ -952,6 +1022,48 @@ impl<T> [T] {
9521022
ArrayChunks::new(self)
9531023
}
9541024

1025+
/// Splits the slice into a slice of `N`-element arrays,
1026+
/// assuming that there's no remainder.
1027+
///
1028+
/// # Safety
1029+
///
1030+
/// This may only be called when
1031+
/// - The slice splits exactly into `N`-element chunks (aka `self.len() % N == 0`).
1032+
/// - `N != 0`.
1033+
///
1034+
/// # Examples
1035+
///
1036+
/// ```
1037+
/// #![feature(slice_as_chunks)]
1038+
/// let slice: &mut [char] = &mut ['l', 'o', 'r', 'e', 'm', '!'];
1039+
/// let chunks: &mut [[char; 1]] =
1040+
/// // SAFETY: 1-element chunks never have remainder
1041+
/// unsafe { slice.as_chunks_unchecked_mut() };
1042+
/// chunks[0] = ['L'];
1043+
/// assert_eq!(chunks, &[['L'], ['o'], ['r'], ['e'], ['m'], ['!']]);
1044+
/// let chunks: &mut [[char; 3]] =
1045+
/// // SAFETY: The slice length (6) is a multiple of 3
1046+
/// unsafe { slice.as_chunks_unchecked_mut() };
1047+
/// chunks[1] = ['a', 'x', '?'];
1048+
/// assert_eq!(slice, &['L', 'o', 'r', 'a', 'x', '?']);
1049+
///
1050+
/// // These would be unsound:
1051+
/// // let chunks: &[[_; 5]] = slice.as_chunks_unchecked_mut() // The slice length is not a multiple of 5
1052+
/// // let chunks: &[[_; 0]] = slice.as_chunks_unchecked_mut() // Zero-length chunks are never allowed
1053+
/// ```
1054+
#[unstable(feature = "slice_as_chunks", issue = "74985")]
1055+
#[inline]
1056+
pub unsafe fn as_chunks_unchecked_mut<const N: usize>(&mut self) -> &mut [[T; N]] {
1057+
debug_assert_ne!(N, 0);
1058+
debug_assert_eq!(self.len() % N, 0);
1059+
let new_len =
1060+
// SAFETY: Our precondition is exactly what's needed to call this
1061+
unsafe { crate::intrinsics::exact_div(self.len(), N) };
1062+
// SAFETY: We cast a slice of `new_len * N` elements into
1063+
// a slice of `new_len` many `N` elements chunks.
1064+
unsafe { from_raw_parts_mut(self.as_mut_ptr().cast(), new_len) }
1065+
}
1066+
9551067
/// Splits the slice into a slice of `N`-element arrays,
9561068
/// starting at the beginning of the slice,
9571069
/// and a remainder slice with length strictly less than `N`.
@@ -982,13 +1094,48 @@ impl<T> [T] {
9821094
assert_ne!(N, 0);
9831095
let len = self.len() / N;
9841096
let (multiple_of_n, remainder) = self.split_at_mut(len * N);
985-
let array_slice: &mut [[T; N]] =
986-
// SAFETY: We cast a slice of `len * N` elements into
987-
// a slice of `len` many `N` elements chunks.
988-
unsafe { from_raw_parts_mut(multiple_of_n.as_mut_ptr().cast(), len) };
1097+
// SAFETY: We already panicked for zero, and ensured by construction
1098+
// that the length of the subslice is a multiple of N.
1099+
let array_slice = unsafe { multiple_of_n.as_chunks_unchecked_mut() };
9891100
(array_slice, remainder)
9901101
}
9911102

1103+
/// Splits the slice into a slice of `N`-element arrays,
1104+
/// starting at the end of the slice,
1105+
/// and a remainder slice with length strictly less than `N`.
1106+
///
1107+
/// # Panics
1108+
///
1109+
/// Panics if `N` is 0. This check will most probably get changed to a compile time
1110+
/// error before this method gets stabilized.
1111+
///
1112+
/// # Examples
1113+
///
1114+
/// ```
1115+
/// #![feature(slice_as_chunks)]
1116+
/// let v = &mut [0, 0, 0, 0, 0];
1117+
/// let mut count = 1;
1118+
///
1119+
/// let (remainder, chunks) = v.as_rchunks_mut();
1120+
/// remainder[0] = 9;
1121+
/// for chunk in chunks {
1122+
/// *chunk = [count; 2];
1123+
/// count += 1;
1124+
/// }
1125+
/// assert_eq!(v, &[9, 1, 1, 2, 2]);
1126+
/// ```
1127+
#[unstable(feature = "slice_as_chunks", issue = "74985")]
1128+
#[inline]
1129+
pub fn as_rchunks_mut<const N: usize>(&mut self) -> (&mut [T], &mut [[T; N]]) {
1130+
assert_ne!(N, 0);
1131+
let len = self.len() / N;
1132+
let (remainder, multiple_of_n) = self.split_at_mut(self.len() - len * N);
1133+
// SAFETY: We already panicked for zero, and ensured by construction
1134+
// that the length of the subslice is a multiple of N.
1135+
let array_slice = unsafe { multiple_of_n.as_chunks_unchecked_mut() };
1136+
(remainder, array_slice)
1137+
}
1138+
9921139
/// Returns an iterator over `N` elements of the slice at a time, starting at the
9931140
/// beginning of the slice.
9941141
///

‎src/test/codegen/slice-as_chunks.rs‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
// no-system-llvm
2+
// compile-flags: -O
3+
// only-64bit (because the LLVM type of i64 for usize shows up)
4+
// ignore-debug: the debug assertions get in the way
5+
6+
#![crate_type = "lib"]
7+
#![feature(slice_as_chunks)]
8+
9+
// CHECK-LABEL: @chunks4
10+
#[no_mangle]
11+
pub fn chunks4(x: &[u8]) -> &[[u8; 4]] {
12+
// CHECK-NEXT: start:
13+
// CHECK-NEXT: lshr i64 %x.1, 2
14+
// CHECK-NOT: shl
15+
// CHECK-NOT: mul
16+
// CHECK-NOT: udiv
17+
// CHECK-NOT: urem
18+
// CHECK: ret
19+
x.as_chunks().0
20+
}
21+
22+
// CHECK-LABEL: @chunks4_with_remainder
23+
#[no_mangle]
24+
pub fn chunks4_with_remainder(x: &[u8]) -> (&[[u8; 4]], &[u8]) {
25+
// CHECK: and i64 %x.1, -4
26+
// CHECK: and i64 %x.1, 3
27+
// CHECK: lshr exact
28+
// CHECK-NOT: mul
29+
// CHECK-NOT: udiv
30+
// CHECK-NOT: urem
31+
// CHECK: ret
32+
x.as_chunks()
33+
}

0 commit comments

Comments
 (0)