shape: Factor out common layout check

rust-ndarray · Apr 17, 2021 · 0cd0262 · 0cd0262
1 parent 62a450c
commit 0cd0262
Showing 1 changed file with 22 additions and 26 deletions.
diff --git a/src/impl_methods.rs b/src/impl_methods.rs
@@ -1642,14 +1642,8 @@ where
         A: Clone,
         S: Data,
     {
-        if size_of_shape_checked(&shape) != Ok(self.dim.size()) {
-            return Err(error::incompatible_shapes(&self.dim, &shape));
-        }
-        let layout = self.layout_impl();
-        if order == Order::Automatic {
-            order = preferred_order_for_layout(layout);
-        }
-
+        let layout = self.layout_and_order(&shape, &mut order)?;
+        // safe because: the number of elements is preserved and it's contiguous
         unsafe {
             if layout.is(Layout::CORDER) && order == Order::RowMajor {
                 let strides = shape.default_strides();
@@ -1672,6 +1666,24 @@ where
         }
     }
 
+    /// Check if `shape` is valid for reshaping the array (in terms of number of elements),
+    /// and also compute the Layout of self and resolve Order if it is Automatic.
+    ///
+    /// After returning successfully, `order` is either RowMajor or ColumnMajor.
+    fn layout_and_order<E>(&self, shape: &E, order: &mut Order) -> Result<Layout, ShapeError>
+    where
+        E: Dimension,
+    {
+        if size_of_shape_checked(shape) != Ok(self.dim.size()) {
+            return Err(error::incompatible_shapes(&self.dim, shape));
+        }
+        let layout = self.layout_impl();
+        if *order == Order::Automatic {
+            *order = preferred_order_for_layout(layout);
+        }
+        Ok(layout)
+    }
+
     /// Transform the array into `shape`; any shape with the same number of elements is accepted,
     /// but the source array or view must be in contiguous and stored in standard row-major (C) or
     /// column-major (Fortran) memory order.
@@ -1707,15 +1719,7 @@ where
     where
         E: Dimension,
     {
-        if size_of_shape_checked(&shape) != Ok(self.dim.size()) {
-            return Err(error::incompatible_shapes(&self.dim, &shape));
-        }
-
-        let layout = self.layout_impl();
-        if order == Order::Automatic {
-            order = preferred_order_for_layout(layout);
-        }
-
+        let layout = self.layout_and_order(&shape, &mut order)?;
         // safe because: the number of elements is preserved and it's contiguous
         unsafe {
             if layout.is(Layout::CORDER) && order == Order::RowMajor {
@@ -1769,15 +1773,7 @@ where
     where
         E: Dimension,
     {
-        if size_of_shape_checked(&shape) != Ok(self.dim.size()) {
-            return Err(error::incompatible_shapes(&self.dim, &shape));
-        }
-
-        let layout = self.layout_impl();
-        if order == Order::Automatic {
-            order = preferred_order_for_layout(layout);
-        }
-
+        let layout = self.layout_and_order(&shape, &mut order)?;
         // safe because: the number of elements is preserved and it's contiguous
         unsafe {
             if layout.is(Layout::CORDER) && order == Order::RowMajor {