inout: split a buffer in equal chunks

Author: baloo

inout/src/inout.rs

@@ -1,6 +1,6 @@
 use crate::InOutBuf;
-use core::{marker::PhantomData, ptr};
-use hybrid_array::{Array, ArraySize};
+use core::{marker::PhantomData, mem::MaybeUninit, ops::Mul, ptr};
+use hybrid_array::{Array, ArraySize, typenum::Prod};
 
 /// Custom pointer type which contains one immutable (input) and one mutable
 /// (output) pointer, which are either equal or non-overlapping.
@@ -151,6 +151,35 @@ impl<'inp, 'out, T, N: ArraySize> InOut<'inp, 'out, Array<T, N>> {
     }
 }
 
+impl<'inp, 'out, T, N, M> From<InOut<'inp, 'out, Array<T, Prod<N, M>>>>
+    for Array<InOut<'inp, 'out, Array<T, N>>, M>
+where
+    N: ArraySize,
+    M: ArraySize,
+    N: Mul<M>,
+    Prod<N, M>: ArraySize,
+{
+    fn from(buf: InOut<'inp, 'out, Array<T, Prod<N, M>>>) -> Self {
+        let split_point = N::USIZE;
+        let mut out = Array::<MaybeUninit<InOut<Array<T, N>>>, M>::uninit();
+        let (mut tail_in_ptr, mut tail_out_ptr) = (buf.in_ptr as *const T, buf.out_ptr as *mut T);
+
+        for i in 0..M::USIZE {
+            let el = InOut {
+                in_ptr: tail_in_ptr as *const Array<T, N>,
+                out_ptr: tail_out_ptr as *mut Array<T, N>,
+                _pd: PhantomData,
+            };
+            out[i].write(el);
+
+            (tail_in_ptr, tail_out_ptr) =
+                unsafe { (tail_in_ptr.add(split_point), tail_out_ptr.add(split_point)) };
+        }
+
+        unsafe { out.assume_init() }
+    }
+}
+
 impl<N: ArraySize> InOut<'_, '_, Array<u8, N>> {
     /// XOR `data` with values behind the input slice and write
     /// result to the output slice.

inout/tests/split-inout.rs

@@ -0,0 +1,34 @@
+use hybrid_array::{
+    Array,
+    sizes::{U2, U4, U8},
+};
+use inout::{InOut, InOutBuf};
+
+#[test]
+fn test_split() {
+    let mut buf = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
+    let inout: InOutBuf<'_, '_, u8> = buf.as_mut_slice().into();
+
+    let expected: Vec<&[u8]> = vec![
+        &[1, 2],
+        &[3, 4],
+        &[5, 6],
+        &[7, 8],
+        &[9, 10],
+        &[11, 12],
+        &[13, 14],
+        &[15, 16],
+    ];
+    let mut expected = expected.into_iter();
+
+    let (blocks, _tail) = inout.into_chunks::<U8>();
+    for block in blocks.into_iter() {
+        type SubBlock = Array<u8, U2>;
+
+        let subblocks = Array::<InOut<'_, '_, SubBlock>, U4>::from(block);
+
+        for subblock in subblocks {
+            assert_eq!(Some(subblock.get_in().as_slice()), expected.next());
+        }
+    }
+}