Merge pull request #803 from rust-ndarray/generic-maybe-uninit

Add maybe_uninit constructor and .assume_init() method

Author: bluss

examples/sort-axis.rs

@@ -102,26 +102,31 @@ where
         assert_eq!(axis_len, perm.indices.len());
         debug_assert!(perm.correct());
 
-        let mut v = Vec::with_capacity(self.len());
-        let mut result;
+        let mut result = Array::maybe_uninit(self.dim());
 
         // panic-critical begin: we must not panic
         unsafe {
-            v.set_len(self.len());
-            result = Array::from_shape_vec_unchecked(self.dim(), v);
+            // logically move ownership of all elements from self into result
+            // the result realizes this ownership at .assume_init() further down
+            let mut moved_elements = 0;
             for i in 0..axis_len {
                 let perm_i = perm.indices[i];
                 Zip::from(result.index_axis_mut(axis, perm_i))
                     .and(self.index_axis(axis, i))
-                    .apply(|to, from| copy_nonoverlapping(from, to, 1));
+                    .apply(|to, from| {
+                        copy_nonoverlapping(from, to.as_mut_ptr(), 1);
+                        moved_elements += 1;
+                    });
             }
             // forget moved array elements but not its vec
+            // old_storage drops empty
             let mut old_storage = self.into_raw_vec();
             old_storage.set_len(0);
-            // old_storage drops empty
+
+            debug_assert_eq!(result.len(), moved_elements);
+            result.assume_init()
         }
         // panic-critical end
-        result
     }
 }
 

src/data_repr.rs

@@ -1,5 +1,6 @@
 
 use std::mem;
+use std::mem::ManuallyDrop;
 use std::ptr::NonNull;
 use std::slice;
 use crate::extension::nonnull;
@@ -17,22 +18,20 @@ pub struct OwnedRepr<A> {
 }
 
 impl<A> OwnedRepr<A> {
-    pub(crate) fn from(mut v: Vec<A>) -> Self {
+    pub(crate) fn from(v: Vec<A>) -> Self {
+        let mut v = ManuallyDrop::new(v);
         let len = v.len();
         let capacity = v.capacity();
         let ptr = nonnull::nonnull_from_vec_data(&mut v);
-        mem::forget(v);
         Self {
             ptr,
             len,
             capacity,
         }
     }
 
-    pub(crate) fn into_vec(mut self) -> Vec<A> {
-        let v = self.take_as_vec();
-        mem::forget(self);
-        v
+    pub(crate) fn into_vec(self) -> Vec<A> {
+        ManuallyDrop::new(self).take_as_vec()
     }
 
     pub(crate) fn as_slice(&self) -> &[A] {

src/data_traits.rs

@@ -528,3 +528,38 @@ unsafe impl<'a, A> Data for CowRepr<'a, A> {
 }
 
 unsafe impl<'a, A> DataMut for CowRepr<'a, A> where A: Clone {}
+
+/// Array representation trait.
+///
+/// The RawDataSubst trait maps the element type of array storage, while
+/// keeping the same kind of storage.
+///
+/// For example, `RawDataSubst<B>` can map the type `OwnedRepr<A>` to `OwnedRepr<B>`.
+pub trait RawDataSubst<A>: RawData {
+    /// The resulting array storage of the same kind but substituted element type
+    type Output: RawData<Elem = A>;
+}
+
+impl<A, B> RawDataSubst<B> for OwnedRepr<A> {
+    type Output = OwnedRepr<B>;
+}
+
+impl<A, B> RawDataSubst<B> for OwnedArcRepr<A> {
+    type Output = OwnedArcRepr<B>;
+}
+
+impl<A, B> RawDataSubst<B> for RawViewRepr<*const A> {
+    type Output = RawViewRepr<*const B>;
+}
+
+impl<A, B> RawDataSubst<B> for RawViewRepr<*mut A> {
+    type Output = RawViewRepr<*mut B>;
+}
+
+impl<'a, A: 'a, B: 'a> RawDataSubst<B> for ViewRepr<&'a A> {
+    type Output = ViewRepr<&'a B>;
+}
+
+impl<'a, A: 'a, B: 'a> RawDataSubst<B> for ViewRepr<&'a mut A> {
+    type Output = ViewRepr<&'a mut B>;
+}

src/impl_constructors.rs

@@ -467,6 +467,9 @@ where
 
     /// Create an array with uninitalized elements, shape `shape`.
     ///
+    /// Prefer to use [`maybe_uninit()`](ArrayBase::maybe_uninit) if possible, because it is
+    /// easier to use correctly.
+    ///
     /// **Panics** if the number of elements in `shape` would overflow isize.
     ///
     /// ### Safety
@@ -524,7 +527,66 @@ where
     S: DataOwned<Elem = MaybeUninit<A>>,
     D: Dimension,
 {
-    pub(crate) fn maybe_uninit<Sh>(shape: Sh) -> Self
+    /// Create an array with uninitalized elements, shape `shape`.
+    ///
+    /// The uninitialized elements of type `A` are represented by the type `MaybeUninit<A>`,
+    /// an easier way to handle uninit values correctly.
+    ///
+    /// Only *when* the array is completely initialized with valid elements, can it be
+    /// converted to an array of `A` elements using [`.assume_init()`].
+    ///
+    /// **Panics** if the number of elements in `shape` would overflow isize.
+    ///
+    /// ### Safety
+    ///
+    /// The whole of the array must be initialized before it is converted
+    /// using [`.assume_init()`] or otherwise traversed.
+    ///
+    /// ### Examples
+    ///
+    /// It is possible to assign individual values through `*elt = MaybeUninit::new(value)`
+    /// and so on.
+    ///
+    /// [`.assume_init()`]: ArrayBase::assume_init
+    ///
+    /// ```
+    /// use ndarray::{s, Array2};
+    /// use ndarray::Zip;
+    /// use ndarray::Axis;
+    ///
+    /// // Example Task: Let's create a transposed copy of the input
+    ///
+    /// fn shift_by_two(a: &Array2<f32>) -> Array2<f32> {
+    ///     // create an uninitialized array
+    ///     let mut b = Array2::maybe_uninit(a.dim());
+    ///
+    ///     // two first columns in b are two last in a
+    ///     // rest of columns in b are the initial columns in a
+    ///
+    ///     assign_to(a.slice(s![.., -2..]), b.slice_mut(s![.., ..2]));
+    ///     assign_to(a.slice(s![.., 2..]), b.slice_mut(s![.., ..-2]));
+    ///
+    ///     // Now we can promise that `b` is safe to use with all operations
+    ///     unsafe {
+    ///         b.assume_init()
+    ///     }
+    /// }
+    ///
+    /// use ndarray::{IntoNdProducer, AssignElem};
+    ///
+    /// fn assign_to<'a, P1, P2, A>(from: P1, to: P2)
+    ///     where P1: IntoNdProducer<Item = &'a A>,
+    ///           P2: IntoNdProducer<Dim = P1::Dim>,
+    ///           P2::Item: AssignElem<A>,
+    ///           A: Clone + 'a
+    /// {
+    ///     Zip::from(from)
+    ///         .apply_assign_into(to, A::clone);
+    /// }
+    ///
+    /// # shift_by_two(&Array2::zeros((8, 8)));
+    /// ```
+    pub fn maybe_uninit<Sh>(shape: Sh) -> Self
     where
         Sh: ShapeBuilder<Dim = D>,
     {

src/impl_owned_array.rs

@@ -1,8 +1,5 @@
-use std::mem::MaybeUninit;
-use std::mem::transmute;
 
 use crate::imp_prelude::*;
-use crate::OwnedRepr;
 
 /// Methods specific to `Array0`.
 ///
@@ -61,33 +58,3 @@ where
         self.data.into_vec()
     }
 }
-
-/// Methods specific to `Array` of `MaybeUninit`.
-///
-/// ***See also all methods for [`ArrayBase`]***
-///
-/// [`ArrayBase`]: struct.ArrayBase.html
-impl<A, D> Array<MaybeUninit<A>, D>
-where
-    D: Dimension,
-{
-    /// Assert that the array's storage's elements are all fully initialized, and conver
-    /// the array from element type `MaybeUninit<A>` to `A`.
-    pub(crate) unsafe fn assume_init(self) -> Array<A, D> {
-        // NOTE: Fully initialized includes elements not reachable in current slicing/view.
-        //
-        // Should this method be generalized to all array types?
-        // (Will need a way to map the RawData<Elem=X> to RawData<Elem=Y> of same kind)
-
-        let Array { data, ptr, dim, strides } = self;
-        let data = transmute::<OwnedRepr<MaybeUninit<A>>, OwnedRepr<A>>(data);
-        let ptr = ptr.cast::<A>();
-
-        Array {
-            data,
-            ptr,
-            dim,
-            strides,
-        }
-    }
-}

src/impl_special_element_types.rs

@@ -0,0 +1,67 @@
+// 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::mem::size_of;
+use std::mem::ManuallyDrop;
+use std::mem::MaybeUninit;
+
+use crate::imp_prelude::*;
+use crate::RawDataSubst;
+
+
+/// Methods specific to arrays with `MaybeUninit` elements.
+///
+/// ***See also all methods for [`ArrayBase`]***
+///
+/// [`ArrayBase`]: struct.ArrayBase.html
+impl<A, S, D> ArrayBase<S, D>
+where
+    S: RawDataSubst<A, Elem=MaybeUninit<A>>,
+    D: Dimension,
+{
+    /// **Promise** that the array's elements are all fully initialized, and convert
+    /// the array from element type `MaybeUninit<A>` to `A`.
+    ///
+    /// For example, it can convert an `Array<MaybeUninit<f64>, D>` to `Array<f64, D>`.
+    ///
+    /// ## Safety
+    ///
+    /// Safe to use if all the array's elements have been initialized.
+    ///
+    /// Note that for owned and shared ownership arrays, the promise must include all of the
+    /// array's storage; it is for example possible to slice these in place, but that must
+    /// only be done after all elements have been initialized.
+    pub unsafe fn assume_init(self) -> ArrayBase<<S as RawDataSubst<A>>::Output, D> {
+        // NOTE: Fully initialized includes elements not reachable in current slicing/view.
+
+        let ArrayBase { data, ptr, dim, strides } = self;
+
+        // transmute from storage of MaybeUninit<A> to storage of A
+        let data = unlimited_transmute::<S, S::Output>(data);
+        let ptr = ptr.cast::<A>();
+
+        ArrayBase {
+            data,
+            ptr,
+            dim,
+            strides,
+        }
+    }
+}
+
+/// Transmute from A to B.
+///
+/// Like transmute, but does not have the compile-time size check which blocks
+/// using regular transmute for "S to S::Output".
+///
+/// **Panics** if the size of A and B are different.
+unsafe fn unlimited_transmute<A, B>(data: A) -> B {
+    assert_eq!(size_of::<A>(), size_of::<B>());
+    let old_data = ManuallyDrop::new(data);
+    (&*old_data as *const A as *const B).read()
+}

src/lib.rs

@@ -159,6 +159,7 @@ pub use crate::aliases::*;
 #[allow(deprecated)]
 pub use crate::data_traits::{
     Data, DataClone, DataMut, DataOwned, DataShared, RawData, RawDataClone, RawDataMut,
+    RawDataSubst,
 };
 
 mod free_functions;
@@ -1483,6 +1484,7 @@ mod impl_constructors;
 
 mod impl_methods;
 mod impl_owned_array;
+mod impl_special_element_types;
 
 /// Private Methods
 impl<A, S, D> ArrayBase<S, D>

tests/array-construct.rs

@@ -7,6 +7,7 @@
 
 use defmac::defmac;
 use ndarray::prelude::*;
+use ndarray::Zip;
 
 #[test]
 fn test_from_shape_fn() {
@@ -194,3 +195,46 @@ fn deny_wraparound_uninit() {
         let _five_large = Array::<f32, _>::uninitialized((3, 7, 29, 36760123, 823996703));
     }
 }
+
+
+#[test]
+fn maybe_uninit_1() {
+    use std::mem::MaybeUninit;
+
+    unsafe {
+        // Array
+        type Mat<D> = Array<MaybeUninit<f32>, D>;
+
+        let mut a = Mat::maybe_uninit((10, 10));
+        a.mapv_inplace(|_| MaybeUninit::new(1.));
+
+        let a_init = a.assume_init();
+        assert_eq!(a_init, Array2::from_elem(a_init.dim(), 1.));
+
+        // ArcArray
+        type ArcMat<D> = ArcArray<MaybeUninit<f32>, D>;
+
+        let mut a = ArcMat::maybe_uninit((10, 10));
+        a.mapv_inplace(|_| MaybeUninit::new(1.));
+        let a2 = a.clone();
+
+        let a_init = a.assume_init();
+        assert_eq!(a_init, Array2::from_elem(a_init.dim(), 1.));
+
+        // ArrayView
+        let av_init = a2.view().assume_init();
+        assert_eq!(av_init, Array2::from_elem(a_init.dim(), 1.));
+
+        // RawArrayViewMut
+        let mut a = Mat::maybe_uninit((10, 10));
+        let v = a.raw_view_mut();
+        Zip::from(v)
+            .apply(|ptr| *(*ptr).as_mut_ptr() = 1.);
+
+        let u = a.raw_view_mut().assume_init();
+
+        Zip::from(u)
+            .apply(|ptr| assert_eq!(*ptr, 1.));
+
+    }
+}