ring/polyfill/slice/
as_chunks_mut.rs

1// Copyright 2025 Brian Smith.
2//
3// Permission to use, copy, modify, and/or distribute this software for any
4// purpose with or without fee is hereby granted, provided that the above
5// copyright notice and this permission notice appear in all copies.
6//
7// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
8// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
9// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
10// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
11// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
12// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
13// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
14
15use super::AsChunks;
16
17#[inline(always)]
18pub fn as_chunks_mut<T, const N: usize>(slice: &mut [T]) -> (AsChunksMut<T, N>, &mut [T]) {
19    assert!(N != 0, "chunk size must be non-zero");
20    let len = slice.len() / N;
21    let (multiple_of_n, remainder) = slice.split_at_mut(len * N);
22    (AsChunksMut(multiple_of_n), remainder)
23}
24
25pub struct AsChunksMut<'a, T, const N: usize>(&'a mut [T]);
26
27impl<T, const N: usize> AsChunksMut<'_, T, N> {
28    #[inline(always)]
29    pub fn as_flattened(&self) -> &[T] {
30        self.0
31    }
32
33    #[inline(always)]
34    pub fn as_flattened_mut(&mut self) -> &mut [T] {
35        self.0
36    }
37
38    #[cfg(target_arch = "aarch64")]
39    pub fn as_ptr(&self) -> *const [T; N] {
40        self.0.as_ptr().cast()
41    }
42
43    #[cfg(target_arch = "x86_64")]
44    pub fn as_ptr(&self) -> *const [T; N] {
45        self.0.as_ptr().cast()
46    }
47
48    #[cfg(target_arch = "aarch64")]
49    pub fn as_mut_ptr(&mut self) -> *mut [T; N] {
50        self.0.as_mut_ptr().cast()
51    }
52
53    #[cfg(target_arch = "x86_64")]
54    #[inline(always)]
55    pub fn as_mut(&mut self) -> AsChunksMut<T, N> {
56        AsChunksMut(self.0)
57    }
58
59    #[inline(always)]
60    pub fn as_ref(&self) -> AsChunks<T, N> {
61        AsChunks::<T, N>::from(self)
62    }
63
64    // Argument moved from runtime argument to `const` argument so that
65    // `CHUNK_LEN * N` is checked at compile time for overflow.
66    #[inline(always)]
67    pub fn chunks_mut<const CHUNK_LEN: usize>(&mut self) -> AsChunksMutChunksMutIter<T, N> {
68        AsChunksMutChunksMutIter(self.0.chunks_mut(CHUNK_LEN * N))
69    }
70
71    #[cfg(target_arch = "x86_64")]
72    #[inline(always)]
73    pub fn split_at_mut(&mut self, mid: usize) -> (AsChunksMut<T, N>, AsChunksMut<T, N>) {
74        let (before, after) = self.0.split_at_mut(mid * N);
75        (AsChunksMut(before), AsChunksMut(after))
76    }
77}
78
79pub struct AsChunksMutChunksMutIter<'a, T, const N: usize>(core::slice::ChunksMut<'a, T>);
80
81impl<'a, T, const N: usize> Iterator for AsChunksMutChunksMutIter<'a, T, N> {
82    type Item = AsChunksMut<'a, T, N>;
83
84    #[inline(always)]
85    fn next(&mut self) -> Option<Self::Item> {
86        self.0.next().map(AsChunksMut)
87    }
88}