From 4718c94b6affa7fc82dc4d0fe46170564f100c2a Mon Sep 17 00:00:00 2001
From: jmwample <8297368+jmwample@users.noreply.github.com>
Date: Wed, 11 Sep 2024 12:42:30 -0600
Subject: [PATCH] implement as_flattened{_mut} for Array<Array<>>

---
 src/iter.rs | 90 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 90 insertions(+)

diff --git a/src/iter.rs b/src/iter.rs
index 281bab1..01b1b74 100644
--- a/src/iter.rs
+++ b/src/iter.rs
@@ -4,6 +4,7 @@
 use crate::{Array, ArraySize};
 use core::{
     fmt,
+    mem::size_of,
     slice::{Iter, IterMut},
 };
 
@@ -105,3 +106,92 @@ where
         self.iter_mut()
     }
 }
+
+impl<T, U, V> Array<Array<T, U>, V>
+where
+    U: ArraySize,
+    V: ArraySize,
+{
+    /// Takes a `&mut Array<Array<T, N>,M>`, and flattens it to a `&mut [T]`.
+    ///
+    /// # Panics
+    ///
+    /// This panics if the length of the resulting slice would overflow a `usize`.
+    ///
+    /// This is only possible when flattening a slice of arrays of zero-sized
+    /// types, and thus tends to be irrelevant in practice. If
+    /// `size_of::<T>() > 0`, this will never panic.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use hybrid_array::{Array, typenum::U3};
+    ///
+    /// fn add_5_to_all(slice: &mut [i32]) {
+    ///     for i in slice {
+    ///         *i += 5;
+    ///     }
+    /// }
+    ///
+    /// let mut array: Array<Array<i32, U3>, U3> = Array([Array([1_i32, 2, 3]), Array([4, 5, 6]), Array([7, 8, 9])]);
+    /// add_5_to_all(array.as_flattened_mut());
+    /// assert_eq!(array, Array([Array([6, 7, 8]), Array([9, 10, 11]), Array([12, 13, 14])]));
+    /// ```
+    pub fn as_flattened_mut(&mut self) -> &mut [T] {
+        let len = if size_of::<T>() == 0 {
+            self.len()
+                .checked_mul(U::USIZE)
+                .expect("slice len overflow")
+        } else {
+            // SAFETY: `self.len() * N` cannot overflow because `self` is
+            // already in the address space.
+            unsafe { self.len().unchecked_mul(U::USIZE) }
+        };
+        // SAFETY: `[T]` is layout-identical to `[T; U]`
+        unsafe { core::slice::from_raw_parts_mut(self.as_mut_ptr().cast(), len) }
+    }
+
+    /// Takes a `&Array<Array<T, N>, >>`, and flattens it to a `&[T]`.
+    ///
+    /// # Panics
+    ///
+    /// This panics if the length of the resulting slice would overflow a `usize`.
+    ///
+    /// This is only possible when flattening a slice of arrays of zero-sized
+    /// types, and thus tends to be irrelevant in practice. If
+    /// `size_of::<T>() > 0`, this will never panic.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use hybrid_array::{Array, typenum::{U0, U2, U3, U5, U10}};
+    ///
+    /// let a: Array<Array<usize, U3>, U2> = Array([Array([1, 2, 3]), Array([4, 5, 6])]);
+    /// assert_eq!(a.as_flattened(), &[1, 2, 3, 4, 5, 6]);
+    ///
+    /// let b: Array<Array<usize, U2>, U3> = Array([Array([1, 2]), Array([3, 4]), Array([5, 6])]);
+    /// assert_eq!(a.as_flattened(), b.as_flattened());
+    ///
+    /// let c: Array<[usize; 2], U3> = Array([[1, 2], [3, 4], [5, 6]]);
+    /// assert_eq!(a.as_flattened(), c.as_flattened());
+    ///
+    /// let slice_of_empty_arrays: &Array<Array<i32, U5>, U0> = &Array::from_fn(|_| Array([1, 2, 3, 4, 5]));
+    /// assert!(slice_of_empty_arrays.as_flattened().is_empty());
+    ///
+    /// let empty_slice_of_arrays: &Array<Array<u32, U10>, U0>  = &Array([]);
+    /// assert!(empty_slice_of_arrays.as_flattened().is_empty());
+    /// ```
+    pub fn as_flattened(&self) -> &[T] {
+        let len = if size_of::<T>() == 0 {
+            self.len()
+                .checked_mul(U::USIZE)
+                .expect("slice len overflow")
+        } else {
+            // SAFETY: `self.len() * N` cannot overflow because `self` is
+            // already in the address space.
+            unsafe { self.len().unchecked_mul(U::USIZE) }
+        };
+        // SAFETY: `[T]` is layout-identical to `[T; U]`
+        unsafe { core::slice::from_raw_parts(self.as_ptr().cast(), len) }
+    }
+}