shape: Add internal constructor from_shape_trusted_iter_unchecked

Author: bluss

src/impl_constructors.rs

@@ -29,6 +29,7 @@ use crate::indices;
 #[cfg(feature = "std")]
 use crate::iterators::to_vec;
 use crate::iterators::to_vec_mapped;
+use crate::iterators::TrustedIterator;
 use crate::StrideShape;
 #[cfg(feature = "std")]
 use crate::{geomspace, linspace, logspace};
@@ -495,6 +496,27 @@ where
         ArrayBase::from_data_ptr(DataOwned::new(v), ptr).with_strides_dim(strides, dim)
     }
 
+    /// Creates an array from an iterator, mapped by `map` and interpret it according to the
+    /// provided shape and strides.
+    ///
+    /// # Safety
+    ///
+    /// See from_shape_vec_unchecked
+    pub(crate) unsafe fn from_shape_trusted_iter_unchecked<Sh, I, F>(shape: Sh, iter: I, map: F)
+        -> Self
+    where
+        Sh: Into<StrideShape<D>>,
+        I: TrustedIterator + ExactSizeIterator,
+        F: FnMut(I::Item) -> A,
+    {
+        let shape = shape.into();
+        let dim = shape.dim;
+        let strides = shape.strides.strides_for_dim(&dim);
+        let v = to_vec_mapped(iter, map);
+        Self::from_vec_dim_stride_unchecked(dim, strides, v)
+    }
+
+
     /// Create an array with uninitalized elements, shape `shape`.
     ///
     /// The uninitialized elements of type `A` are represented by the type `MaybeUninit<A>`,

src/impl_methods.rs

@@ -2226,17 +2226,14 @@ where
         A: 'a,
         S: Data,
     {
-        if let Some(slc) = self.as_slice_memory_order() {
-            let v = crate::iterators::to_vec_mapped(slc.iter(), f);
-            unsafe {
-                ArrayBase::from_shape_vec_unchecked(
+        unsafe {
+            if let Some(slc) = self.as_slice_memory_order() {
+                ArrayBase::from_shape_trusted_iter_unchecked(
                     self.dim.clone().strides(self.strides.clone()),
-                    v,
-                )
+                    slc.iter(), f)
+            } else {
+                ArrayBase::from_shape_trusted_iter_unchecked(self.dim.clone(), self.iter(), f)
             }
-        } else {
-            let v = crate::iterators::to_vec_mapped(self.iter(), f);
-            unsafe { ArrayBase::from_shape_vec_unchecked(self.dim.clone(), v) }
         }
     }
 
@@ -2256,11 +2253,10 @@ where
         if self.is_contiguous() {
             let strides = self.strides.clone();
             let slc = self.as_slice_memory_order_mut().unwrap();
-            let v = crate::iterators::to_vec_mapped(slc.iter_mut(), f);
-            unsafe { ArrayBase::from_shape_vec_unchecked(dim.strides(strides), v) }
+            unsafe { ArrayBase::from_shape_trusted_iter_unchecked(dim.strides(strides),
+                        slc.iter_mut(), f) }
         } else {
-            let v = crate::iterators::to_vec_mapped(self.iter_mut(), f);
-            unsafe { ArrayBase::from_shape_vec_unchecked(dim, v) }
+            unsafe { ArrayBase::from_shape_trusted_iter_unchecked(dim, self.iter_mut(), f) }
         }
     }