Update into_raw_vec

Without the offset from start of allocation to the logically first
element of the array, correctly reinterpreting the results of
.into_raw_vec() as an n-D array is tricky.

Author: jturner314

examples/sort-axis.rs

@@ -157,7 +157,7 @@ where D: Dimension
                 });
             debug_assert_eq!(result.len(), moved_elements);
             // forget the old elements but not the allocation
-            let mut old_storage = self.into_raw_vec();
+            let mut old_storage = self.into_raw_vec().0;
             old_storage.set_len(0);
 
             // transfer ownership of the elements into the result

src/dimension/mod.rs

@@ -409,8 +409,8 @@ fn to_abs_slice(axis_len: usize, slice: Slice) -> (usize, usize, isize)
     (start, end, step)
 }
 
-/// Returns the offset from the lowest-address element to the logically first
-/// element.
+/// This function computes the offset from the lowest address element to the
+/// logically first element. The result is always >= 0.
 pub fn offset_from_low_addr_ptr_to_logical_ptr<D: Dimension>(dim: &D, strides: &D) -> usize
 {
     let offset = izip!(dim.slice(), strides.slice()).fold(0, |_offset, (&d, &s)| {

src/impl_owned_array.rs

@@ -59,14 +59,89 @@ impl<A> Array<A, Ix0>
 impl<A, D> Array<A, D>
 where D: Dimension
 {
+    /// Returns the offset (in units of `A`) from the start of the allocation
+    /// to the first element, or `None` if the array is empty.
+    fn offset_from_alloc_to_logical_ptr(&self) -> Option<usize>
+    {
+        if self.is_empty() {
+            return None;
+        }
+        if std::mem::size_of::<A>() == 0 {
+            Some(dimension::offset_from_low_addr_ptr_to_logical_ptr(&self.dim, &self.strides))
+        } else {
+            let offset = unsafe { self.as_ptr().offset_from(self.data.as_ptr()) };
+            debug_assert!(offset >= 0);
+            Some(offset as usize)
+        }
+    }
+
     /// Return a vector of the elements in the array, in the way they are
-    /// stored internally.
+    /// stored internally, and the index in the vector corresponding to the
+    /// logically first element of the array (or `None` if the array is empty).
     ///
     /// If the array is in standard memory layout, the logical element order
     /// of the array (`.iter()` order) and of the returned vector will be the same.
-    pub fn into_raw_vec(self) -> Vec<A>
+    ///
+    /// ```
+    /// use ndarray::{array, Array2, Axis};
+    ///
+    /// let mut arr: Array2<f64> = array![[1., 2.], [3., 4.], [5., 6.]];
+    /// arr.slice_axis_inplace(Axis(0), (1..).into());
+    /// assert_eq!(arr[[0, 0]], 3.);
+    /// let copy = arr.clone();
+    ///
+    /// let shape = arr.shape().to_owned();
+    /// let strides = arr.strides().to_owned();
+    /// let (v, offset) = arr.into_raw_vec();
+    ///
+    /// assert_eq!(v, &[1., 2., 3., 4., 5., 6.]);
+    /// assert_eq!(offset, Some(2));
+    /// assert_eq!(v[offset.unwrap()], 3.);
+    /// for row in 0..shape[0] {
+    ///     for col in 0..shape[1] {
+    ///         let index = (
+    ///             offset.unwrap() as isize
+    ///             + row as isize * strides[0]
+    ///             + col as isize * strides[1]
+    ///         ) as usize;
+    ///         assert_eq!(v[index], copy[[row, col]]);
+    ///     }
+    /// }
+    /// ```
+    ///
+    /// In the case of zero-sized elements, the offset to the logically first
+    /// element is somewhat meaningless. For convenience, an offset will be
+    /// returned such that all indices computed using the offset, shape, and
+    /// strides will be in-bounds for the `Vec<A>`. Note that this offset won't
+    /// necessarily be the same as the offset for an array of nonzero-sized
+    /// elements sliced in the same way.
+    ///
+    /// ```
+    /// use ndarray::{array, Array2, Axis};
+    ///
+    /// let mut arr: Array2<()> = array![[(), ()], [(), ()], [(), ()]];
+    /// arr.slice_axis_inplace(Axis(0), (1..).into());
+    ///
+    /// let shape = arr.shape().to_owned();
+    /// let strides = arr.strides().to_owned();
+    /// let (v, offset) = arr.into_raw_vec();
+    ///
+    /// assert_eq!(v, &[(), (), (), (), (), ()]);
+    /// for row in 0..shape[0] {
+    ///     for col in 0..shape[1] {
+    ///         let index = (
+    ///             offset.unwrap() as isize
+    ///             + row as isize * strides[0]
+    ///             + col as isize * strides[1]
+    ///         ) as usize;
+    ///         assert_eq!(v[index], ());
+    ///     }
+    /// }
+    /// ```
+    pub fn into_raw_vec(self) -> (Vec<A>, Option<usize>)
     {
-        self.data.into_vec()
+        let offset = self.offset_from_alloc_to_logical_ptr();
+        (self.data.into_vec(), offset)
     }
 }
 
@@ -575,16 +650,11 @@ where D: Dimension
 
         unsafe {
             // grow backing storage and update head ptr
-            let data_to_array_offset = if std::mem::size_of::<A>() != 0 {
-                self.as_ptr().offset_from(self.data.as_ptr())
-            } else {
-                0
-            };
-            debug_assert!(data_to_array_offset >= 0);
+            let offset_from_alloc_to_logical = self.offset_from_alloc_to_logical_ptr().unwrap_or(0);
             self.ptr = self
                 .data
                 .reserve(len_to_append)
-                .offset(data_to_array_offset);
+                .add(offset_from_alloc_to_logical);
 
             // clone elements from view to the array now
             //

tests/array-construct.rs

@@ -19,9 +19,9 @@ fn test_from_shape_fn()
 fn test_dimension_zero()
 {
     let a: Array2<f32> = Array2::from(vec![[], [], []]);
-    assert_eq!(vec![0.; 0], a.into_raw_vec());
+    assert_eq!(vec![0.; 0], a.into_raw_vec().0);
     let a: Array3<f32> = Array3::from(vec![[[]], [[]], [[]]]);
-    assert_eq!(vec![0.; 0], a.into_raw_vec());
+    assert_eq!(vec![0.; 0], a.into_raw_vec().0);
 }
 
 #[test]

tests/array.rs

@@ -1157,7 +1157,7 @@ fn array0_into_scalar()
     // With this kind of setup, the `Array`'s pointer is not the same as the
     // underlying `Vec`'s pointer.
     let a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
-    assert_ne!(a.as_ptr(), a.into_raw_vec().as_ptr());
+    assert_ne!(a.as_ptr(), a.into_raw_vec().0.as_ptr());
     // `.into_scalar()` should still work correctly.
     let a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
     assert_eq!(a.into_scalar(), 6);
@@ -1173,7 +1173,7 @@ fn array_view0_into_scalar()
     // With this kind of setup, the `Array`'s pointer is not the same as the
     // underlying `Vec`'s pointer.
     let a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
-    assert_ne!(a.as_ptr(), a.into_raw_vec().as_ptr());
+    assert_ne!(a.as_ptr(), a.into_raw_vec().0.as_ptr());
     // `.into_scalar()` should still work correctly.
     let a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
     assert_eq!(a.view().into_scalar(), &6);
@@ -1189,7 +1189,7 @@ fn array_view_mut0_into_scalar()
     // With this kind of setup, the `Array`'s pointer is not the same as the
     // underlying `Vec`'s pointer.
     let a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
-    assert_ne!(a.as_ptr(), a.into_raw_vec().as_ptr());
+    assert_ne!(a.as_ptr(), a.into_raw_vec().0.as_ptr());
     // `.into_scalar()` should still work correctly.
     let mut a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
     assert_eq!(a.view_mut().into_scalar(), &6);