From 026d7fae030b5c461ebcfd5784ccdcc3452b201a Mon Sep 17 00:00:00 2001
From: Ruslan Piasetskyi <ruslan.piasetskyi0@gmail.com>
Date: Fri, 13 Oct 2023 17:38:11 +0200
Subject: [PATCH] hybrid-array: implement Concat and Split

These traits add convenient methods for concatenation and splitting.
The size of the result is calculated in the compile time.

This implementation was inspired by GenericArray implementation of
relevant traits.
---
 hybrid-array/src/lib.rs | 133 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 129 insertions(+), 4 deletions(-)

diff --git a/hybrid-array/src/lib.rs b/hybrid-array/src/lib.rs
index e4817adc..13ef4865 100644
--- a/hybrid-array/src/lib.rs
+++ b/hybrid-array/src/lib.rs
@@ -32,10 +32,12 @@ use core::{
     cmp::Ordering,
     fmt::{self, Debug},
     hash::{Hash, Hasher},
-    ops::{Deref, DerefMut, Index, IndexMut, Range},
+    mem::{ManuallyDrop, MaybeUninit},
+    ops::{Add, Deref, DerefMut, Index, IndexMut, Range, Sub},
+    ptr,
     slice::{Iter, IterMut},
 };
-use typenum::Unsigned;
+use typenum::{Diff, Sum, Unsigned};
 
 /// Hybrid typenum-based and const generic array type.
 ///
@@ -613,6 +615,110 @@ macro_rules! impl_array_size {
      };
 }
 
+/// Defines trait for concatenation.
+pub trait Concat<T, M: ArraySize> {
+    /// Part to be appended to `self`
+    type Suffix;
+
+    /// Result of the concatenation.
+    type Output;
+
+    /// Concatenates `self` with `suffix`.
+    fn concat(self, suffix: Self::Suffix) -> Self::Output;
+}
+
+impl<T, M, N> Concat<T, M> for Array<T, N>
+where
+    N: ArraySize + Add<M>,
+    M: ArraySize,
+    Sum<N, M>: ArraySize,
+{
+    type Suffix = Array<T, M>;
+    type Output = Array<T, Sum<N, M>>;
+
+    fn concat(self, suffix: Self::Suffix) -> Self::Output {
+        let mut output: MaybeUninit<Self::Output> = MaybeUninit::uninit();
+        let out_ptr = output.as_mut_ptr() as *mut Self;
+
+        unsafe {
+            ptr::write(out_ptr, self);
+            ptr::write(out_ptr.add(1) as *mut _, suffix);
+            output.assume_init()
+        }
+    }
+}
+
+/// Defines trait for splitting.
+pub trait Split<T, K: ArraySize> {
+    /// First part of the split.
+    type Head;
+    /// Second part of the split.
+    type Tail;
+
+    /// Splits `self` in two parts.
+    fn split(self) -> (Self::Head, Self::Tail);
+}
+
+impl<T, N, K> Split<T, K> for Array<T, N>
+where
+    N: ArraySize,
+    K: ArraySize,
+    N: Sub<K>,
+    Diff<N, K>: ArraySize,
+{
+    type Head = Array<T, K>;
+    type Tail = Array<T, Diff<N, K>>;
+
+    fn split(self) -> (Self::Head, Self::Tail) {
+        unsafe {
+            let array = ManuallyDrop::new(self);
+            let head = ptr::read(array.as_ptr() as *const _);
+            let tail = ptr::read(array.as_ptr().add(K::USIZE) as *const _);
+            (head, tail)
+        }
+    }
+}
+
+impl<'a, T, N, K> Split<T, K> for &'a Array<T, N>
+where
+    N: ArraySize,
+    K: ArraySize,
+    N: Sub<K>,
+    Diff<N, K>: ArraySize,
+{
+    type Head = &'a Array<T, K>;
+    type Tail = &'a Array<T, Diff<N, K>>;
+
+    fn split(self) -> (Self::Head, Self::Tail) {
+        unsafe {
+            let array_ptr: *const T = self.as_ptr();
+            let head = &*(array_ptr as *const _);
+            let tail = &*(array_ptr.add(K::USIZE) as *const _);
+            (head, tail)
+        }
+    }
+}
+
+impl<'a, T, N, K> Split<T, K> for &'a mut Array<T, N>
+where
+    N: ArraySize,
+    K: ArraySize,
+    N: Sub<K>,
+    Diff<N, K>: ArraySize,
+{
+    type Head = &'a mut Array<T, K>;
+    type Tail = &'a mut Array<T, Diff<N, K>>;
+
+    fn split(self) -> (Self::Head, Self::Tail) {
+        unsafe {
+            let array_ptr: *mut T = self.as_mut_ptr();
+            let head = &mut *(array_ptr as *mut _);
+            let tail = &mut *(array_ptr.add(K::USIZE) as *mut _);
+            (head, tail)
+        }
+    }
+}
+
 impl_array_size! {
     0 => U0,
     1 => U1,
@@ -696,9 +802,9 @@ impl_array_size! {
 
 #[cfg(test)]
 mod tests {
-    use super::ByteArray;
+    use super::{ByteArray, Concat, Split};
     use crate::Array;
-    use typenum::{U0, U3, U6, U7};
+    use typenum::{U0, U2, U3, U4, U6, U7};
 
     const EXAMPLE_SLICE: &[u8] = &[1, 2, 3, 4, 5, 6];
 
@@ -729,4 +835,23 @@ mod tests {
 
         assert!(<&ByteArray::<U7>>::try_from(EXAMPLE_SLICE).is_err());
     }
+
+    #[test]
+    fn concat() {
+        let prefix = ByteArray::<U2>::clone_from_slice(&EXAMPLE_SLICE[..2]);
+        let suffix = ByteArray::<U4>::clone_from_slice(&EXAMPLE_SLICE[2..]);
+
+        let array = prefix.concat(suffix);
+        assert_eq!(array.as_slice(), EXAMPLE_SLICE);
+    }
+
+    #[test]
+    fn split() {
+        let array = ByteArray::<U6>::clone_from_slice(EXAMPLE_SLICE);
+
+        let (prefix, suffix) = Split::<_, U2>::split(array);
+
+        assert_eq!(prefix.as_slice(), &EXAMPLE_SLICE[..2]);
+        assert_eq!(suffix.as_slice(), &EXAMPLE_SLICE[2..]);
+    }
 }