From b94faf724f3b6f91cb0d7c3f2f10988d762dc638 Mon Sep 17 00:00:00 2001
From: Barak Ugav <barakugav@gmail.com>
Date: Tue, 9 Jul 2024 13:42:06 +0300
Subject: [PATCH] Add flatten methods

---
 src/impl_methods.rs | 94 ++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 92 insertions(+), 2 deletions(-)

diff --git a/src/impl_methods.rs b/src/impl_methods.rs
index f21b407e1..eea017161 100644
--- a/src/impl_methods.rs
+++ b/src/impl_methods.rs
@@ -2017,8 +2017,6 @@ where
     /// possible, otherwise they are copied to create a new array.
     ///
     /// If an index ordering is not specified, the default is `RowMajor`.
-    /// The operation will only succeed if the array's memory layout is compatible with
-    /// the index ordering, so that the array elements can be rearranged in place.
     ///
     /// # `.to_shape` vs `.into_shape_clone`
     ///
@@ -2131,6 +2129,69 @@ where
         }
     }
 
+    /// Flatten the array to a one-dimensional array.
+    ///
+    /// The array is returned as a `CowArray`; a view if possible, otherwise an owned array.
+    ///
+    /// ```
+    /// use ndarray::{arr1, arr3};
+    ///
+    /// let array = arr3(&[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]);
+    /// let flattened = array.flatten();
+    /// assert_eq!(flattened, arr1(&[1, 2, 3, 4, 5, 6, 7, 8]));
+    /// ```
+    pub fn flatten(&self) -> CowArray<'_, A, Ix1>
+    where
+        A: Clone,
+        S: Data,
+    {
+        self.flatten_with_order(Order::RowMajor)
+    }
+
+    /// Flatten the array to a one-dimensional array.
+    ///
+    /// `order` specifies the *logical* order in which the array is to be read and reshaped.
+    /// The array is returned as a `CowArray`; a view if possible, otherwise an owned array.
+    ///
+    /// ```
+    /// use ndarray::{arr1, arr2};
+    /// use ndarray::Order;
+    ///
+    /// let array = arr2(&[[1, 2], [3, 4], [5, 6], [7, 8]]);
+    /// let flattened = array.flatten_with_order(Order::RowMajor);
+    /// assert_eq!(flattened, arr1(&[1, 2, 3, 4, 5, 6, 7, 8]));
+    /// let flattened = array.flatten_with_order(Order::ColumnMajor);
+    /// assert_eq!(flattened, arr1(&[1, 3, 5, 7, 2, 4, 6, 8]));
+    /// ```
+    pub fn flatten_with_order(&self, order: Order) -> CowArray<'_, A, Ix1>
+    where
+        A: Clone,
+        S: Data,
+    {
+        self.to_shape((self.len(), order)).unwrap()
+    }
+
+    /// Flatten the array to a one-dimensional array, consuming the array.
+    ///
+    /// If possible, no copy is made, and the new array use the same memory as the original array.
+    /// Otherwise, a new array is allocated and the elements are copied.
+    ///
+    /// ```
+    /// use ndarray::{arr1, arr3};
+    ///
+    /// let array = arr3(&[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]);
+    /// let flattened = array.into_flat();
+    /// assert_eq!(flattened, arr1(&[1, 2, 3, 4, 5, 6, 7, 8]));
+    /// ```
+    pub fn into_flat(self) -> ArrayBase<S, Ix1>
+    where
+        A: Clone,
+        S: DataOwned,
+    {
+        let len = self.len();
+        self.into_shape_clone(Ix1(len)).unwrap()
+    }
+
     /// Convert any array or array view to a dynamic dimensional array or
     /// array view (respectively).
     ///
@@ -3065,3 +3126,32 @@ unsafe fn unlimited_transmute<A, B>(data: A) -> B
 }
 
 type DimMaxOf<A, B> = <A as DimMax<B>>::Output;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::arr3;
+
+    #[test]
+    fn test_flatten() {
+        let array = arr3(&[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]);
+        let flattened = array.flatten();
+        assert_eq!(flattened, arr1(&[1, 2, 3, 4, 5, 6, 7, 8]));
+    }
+
+    #[test]
+    fn test_flatten_with_order() {
+        let array = arr2(&[[1, 2], [3, 4], [5, 6], [7, 8]]);
+        let flattened = array.flatten_with_order(Order::RowMajor);
+        assert_eq!(flattened, arr1(&[1, 2, 3, 4, 5, 6, 7, 8]));
+        let flattened = array.flatten_with_order(Order::ColumnMajor);
+        assert_eq!(flattened, arr1(&[1, 3, 5, 7, 2, 4, 6, 8]));
+    }
+
+    #[test]
+    fn test_into_flat() {
+        let array = arr3(&[[[1, 2], [3, 4]], [[5, 6], [7, 8]]]);
+        let flattened = array.into_flat();
+        assert_eq!(flattened, arr1(&[1, 2, 3, 4, 5, 6, 7, 8]));
+    }
+}