FEAT: Factor out traits SplitAt and SplitPreference

To be used by Zip and parallel Zip

Author: bluss

src/indexes.rs

@@ -7,7 +7,8 @@
 // except according to those terms.
 use super::Dimension;
 use crate::dimension::IntoDimension;
-use crate::zip::{Offset, Splittable};
+use crate::zip::Offset;
+use crate::split_at::SplitAt;
 use crate::Axis;
 use crate::Layout;
 use crate::NdProducer;

src/lib.rs

@@ -179,6 +179,7 @@ mod numeric_util;
 mod shape_builder;
 #[macro_use]
 mod slice;
+mod split_at;
 mod stacking;
 #[macro_use]
 mod zip;

src/split_at.rs

@@ -0,0 +1,49 @@
+
+use crate::imp_prelude::*;
+
+/// Arrays and similar that can be split along an axis
+pub(crate) trait SplitAt  {
+    fn split_at(self, axis: Axis, index: usize) -> (Self, Self) where Self: Sized;
+}
+
+pub(crate) trait SplitPreference : SplitAt {
+    fn can_split(&self) -> bool;
+    fn size(&self) -> usize;
+    fn split_preference(&self) -> (Axis, usize);
+    fn split(self) -> (Self, Self) where Self: Sized {
+        let (axis, index) = self.split_preference();
+        self.split_at(axis, index)
+    }
+}
+
+impl<D> SplitAt for D
+where
+    D: Dimension,
+{
+    fn split_at(self, axis: Axis, index: Ix) -> (Self, Self) {
+        let mut d1 = self;
+        let mut d2 = d1.clone();
+        let i = axis.index();
+        let len = d1[i];
+        d1[i] = index;
+        d2[i] = len - index;
+        (d1, d2)
+    }
+}
+
+impl<'a, A, D> SplitAt for ArrayViewMut<'a, A, D>
+    where D: Dimension
+{
+    fn split_at(self, axis: Axis, index: usize) -> (Self, Self) {
+        self.split_at(axis, index)
+    }
+}
+
+
+impl<A, D> SplitAt for RawArrayViewMut<A, D>
+    where D: Dimension
+{
+    fn split_at(self, axis: Axis, index: usize) -> (Self, Self) {
+        self.split_at(axis, index)
+    }
+}

src/zip/mod.rs

@@ -20,6 +20,7 @@ use crate::NdIndex;
 
 use crate::indexes::{indices, Indices};
 use crate::layout::{CORDER, FORDER};
+use crate::split_at::{SplitPreference, SplitAt};
 
 use partial_array::PartialArray;
 
@@ -92,25 +93,6 @@ where
     private_impl! {}
 }
 
-pub trait Splittable: Sized {
-    fn split_at(self, axis: Axis, index: Ix) -> (Self, Self);
-}
-
-impl<D> Splittable for D
-where
-    D: Dimension,
-{
-    fn split_at(self, axis: Axis, index: Ix) -> (Self, Self) {
-        let mut d1 = self;
-        let mut d2 = d1.clone();
-        let i = axis.index();
-        let len = d1[i];
-        d1[i] = index;
-        d2[i] = len - index;
-        (d1, d2)
-    }
-}
-
 /// Argument conversion into a producer.
 ///
 /// Slices and vectors can be used (equivalent to 1-dimensional array views).
@@ -1121,9 +1103,31 @@ macro_rules! map_impl {
             pub fn split(self) -> (Self, Self) {
                 debug_assert_ne!(self.size(), 0, "Attempt to split empty zip");
                 debug_assert_ne!(self.size(), 1, "Attempt to split zip with 1 elem");
+                SplitPreference::split(self)
+            }
+        }
+
+        impl<D, $($p),*> SplitPreference for Zip<($($p,)*), D>
+            where D: Dimension,
+                  $($p: NdProducer<Dim=D> ,)*
+        {
+            fn can_split(&self) -> bool { self.size() > 1 }
+
+            fn size(&self) -> usize { self.size() }
+
+            fn split_preference(&self) -> (Axis, usize) {
                 // Always split in a way that preserves layout (if any)
                 let axis = self.max_stride_axis();
                 let index = self.len_of(axis) / 2;
+                (axis, index)
+            }
+        }
+
+        impl<D, $($p),*> SplitAt for Zip<($($p,)*), D>
+            where D: Dimension,
+                  $($p: NdProducer<Dim=D> ,)*
+        {
+            fn split_at(self, axis: Axis, index: usize) -> (Self, Self) {
                 let (p1, p2) = self.parts.split_at(axis, index);
                 let (d1, d2) = self.dimension.split_at(axis, index);
                 (Zip {
@@ -1139,7 +1143,9 @@ macro_rules! map_impl {
                     layout_tendency: self.layout_tendency,
                 })
             }
+
         }
+
         )+
     }
 }