FIX: Use Partial instead of PartialArray

Partial is just a contiguous slice, and much simpler than PartialArray;
Partial is all that's needed, because the area written will always be
contiguous.

Author: bluss

src/lib.rs

@@ -176,6 +176,7 @@ mod linalg_traits;
 mod linspace;
 mod logspace;
 mod numeric_util;
+mod partial;
 mod shape_builder;
 #[macro_use]
 mod slice;

src/parallel/impl_par_methods.rs

@@ -5,6 +5,8 @@ use crate::parallel::prelude::*;
 use crate::parallel::par::ParallelSplits;
 use super::send_producer::SendProducer;
 
+use crate::partial::Partial;
+
 /// # Parallel methods
 ///
 /// These methods require crate feature `rayon`.
@@ -163,76 +165,3 @@ zip_impl! {
     [true P1 P2 P3 P4 P5],
     [false P1 P2 P3 P4 P5 P6],
 }
-
-/// Partial is a partially written contiguous slice of data;
-/// it is the owner of the elements, but not the allocation,
-/// and will drop the elements on drop.
-#[must_use]
-pub(crate) struct Partial<T> {
-    /// Data pointer
-    ptr: *mut T,
-    /// Current length
-    len: usize,
-}
-
-impl<T> Partial<T> {
-    /// Create an empty partial for this data pointer
-    ///
-    /// Safety: Unless ownership is released, the 
-    /// Partial acts as an owner of the slice of data (not the allocation);
-    /// and will free the elements on drop; the pointer must be dereferenceable
-    /// and the `len` elements following it valid.
-    pub(crate) unsafe fn new(ptr: *mut T) -> Self {
-        Self {
-            ptr,
-            len: 0,
-        }
-    }
-
-    pub(crate) fn stub() -> Self {
-        Self { len: 0, ptr: 0 as *mut _ }
-    }
-
-    pub(crate) fn is_stub(&self) -> bool {
-        self.ptr.is_null()
-    }
-
-    /// Release Partial's ownership of the written elements, and return the current length
-    pub(crate) fn release_ownership(mut self) -> usize {
-        let ret = self.len;
-        self.len = 0;
-        ret
-    }
-
-    /// Merge if they are in order (left to right) and contiguous.
-    /// Skips merge if T does not need drop.
-    pub(crate) fn try_merge(mut left: Self, right: Self) -> Self {
-        if !std::mem::needs_drop::<T>() {
-            return left;
-        }
-        // Merge the partial collect results; the final result will be a slice that
-        // covers the whole output.
-        if left.is_stub() {
-            right
-        } else if left.ptr.wrapping_add(left.len) == right.ptr {
-            left.len += right.release_ownership();
-            left
-        } else {
-            // failure to merge; this is a bug in collect, so we will never reach this
-            debug_assert!(false, "Partial: failure to merge left and right parts");
-            left
-        }
-    }
-}
-
-unsafe impl<T> Send for Partial<T> where T: Send { }
-
-impl<T> Drop for Partial<T> {
-    fn drop(&mut self) {
-        if !self.ptr.is_null() {
-            unsafe {
-                std::ptr::drop_in_place(std::slice::from_raw_parts_mut(self.ptr, self.len));
-            }
-        }
-    }
-}

src/partial.rs

@@ -0,0 +1,88 @@
+// Copyright 2020 bluss and ndarray developers.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use std::ptr;
+
+/// Partial is a partially written contiguous slice of data;
+/// it is the owner of the elements, but not the allocation,
+/// and will drop the elements on drop.
+#[must_use]
+pub(crate) struct Partial<T> {
+    /// Data pointer
+    ptr: *mut T,
+    /// Current length
+    pub(crate) len: usize,
+}
+
+impl<T> Partial<T> {
+    /// Create an empty partial for this data pointer
+    ///
+    /// ## Safety
+    ///
+    /// Unless ownership is released, the Partial acts as an owner of the slice of data (not the
+    /// allocation); and will free the elements on drop; the pointer must be dereferenceable and
+    /// the `len` elements following it valid.
+    ///
+    /// The Partial has an accessible length field which must only be modified in trusted code.
+    pub(crate) unsafe fn new(ptr: *mut T) -> Self {
+        Self {
+            ptr,
+            len: 0,
+        }
+    }
+
+    #[cfg(feature = "rayon")]
+    pub(crate) fn stub() -> Self {
+        Self { len: 0, ptr: 0 as *mut _ }
+    }
+
+    #[cfg(feature = "rayon")]
+    pub(crate) fn is_stub(&self) -> bool {
+        self.ptr.is_null()
+    }
+
+    /// Release Partial's ownership of the written elements, and return the current length
+    pub(crate) fn release_ownership(mut self) -> usize {
+        let ret = self.len;
+        self.len = 0;
+        ret
+    }
+
+    #[cfg(feature = "rayon")]
+    /// Merge if they are in order (left to right) and contiguous.
+    /// Skips merge if T does not need drop.
+    pub(crate) fn try_merge(mut left: Self, right: Self) -> Self {
+        if !std::mem::needs_drop::<T>() {
+            return left;
+        }
+        // Merge the partial collect results; the final result will be a slice that
+        // covers the whole output.
+        if left.is_stub() {
+            right
+        } else if left.ptr.wrapping_add(left.len) == right.ptr {
+            left.len += right.release_ownership();
+            left
+        } else {
+            // failure to merge; this is a bug in collect, so we will never reach this
+            debug_assert!(false, "Partial: failure to merge left and right parts");
+            left
+        }
+    }
+}
+
+unsafe impl<T> Send for Partial<T> where T: Send { }
+
+impl<T> Drop for Partial<T> {
+    fn drop(&mut self) {
+        if !self.ptr.is_null() {
+            unsafe {
+                ptr::drop_in_place(std::slice::from_raw_parts_mut(self.ptr, self.len));
+            }
+        }
+    }
+}

src/zip/mod.rs

@@ -8,7 +8,6 @@
 
 #[macro_use]
 mod zipmacro;
-mod partial_array;
 
 use std::mem::MaybeUninit;
 
@@ -17,13 +16,12 @@ use crate::AssignElem;
 use crate::IntoDimension;
 use crate::Layout;
 use crate::NdIndex;
+use crate::partial::Partial;
 
 use crate::indexes::{indices, Indices};
 use crate::layout::{CORDER, FORDER};
 use crate::split_at::{SplitPreference, SplitAt};
 
-use partial_array::PartialArray;
-
 /// Return if the expression is a break value.
 macro_rules! fold_while {
     ($e:expr) => {
@@ -1070,21 +1068,28 @@ macro_rules! map_impl {
             /// inputs.
             ///
             /// If all inputs are c- or f-order respectively, that is preserved in the output.
-            pub fn apply_collect<R>(self, f: impl FnMut($($p::Item,)* ) -> R) -> Array<R, D>
+            pub fn apply_collect<R>(self, mut f: impl FnMut($($p::Item,)* ) -> R) -> Array<R, D>
             {
                 // Make uninit result
                 let mut output = self.uninitalized_for_current_layout::<R>();
                 if !std::mem::needs_drop::<R>() {
                     // For elements with no drop glue, just overwrite into the array
                     self.apply_assign_into(&mut output, f);
                 } else {
-                    // For generic elements, use a proxy that counts the number of filled elements,
+                    // For generic elements, use a Partial to counts the number of filled elements,
                     // and can drop the right number of elements on unwinding
                     unsafe {
-                        PartialArray::scope(output.view_mut(), move |partial| {
-                            debug_assert_eq!(partial.layout().tendency() >= 0, self.layout_tendency >= 0);
-                            self.apply_assign_into(partial, f);
+                        let mut output = output.raw_view_mut().cast::<R>();
+                        let mut partial = Partial::new(output.as_mut_ptr());
+                        let partial_ref = &mut partial;
+                        debug_assert!(output.is_contiguous());
+                        debug_assert_eq!(output.layout().tendency() >= 0, self.layout_tendency >= 0);
+                        self.and(output)
+                            .apply(move |$($p, )* output_: *mut R| {
+                                output_.write(f($($p ),*));
+                                partial_ref.len += 1;
                         });
+                        partial.release_ownership();
                     }
                 }