Add reserve method for owned arrays

Author: ssande7

src/impl_owned_array.rs

@@ -660,14 +660,10 @@ where D: Dimension
             self.strides.clone()
         };
 
-        unsafe {
-            // grow backing storage and update head ptr
-            let offset_from_alloc_to_logical = self.offset_from_alloc_to_logical_ptr().unwrap_or(0);
-            self.ptr = self
-                .data
-                .reserve(len_to_append)
-                .add(offset_from_alloc_to_logical);
+        // grow backing storage and update head ptr
+        self.reserve(axis, array_dim[axis.index()]);
 
+        unsafe {
             // clone elements from view to the array now
             //
             // To be robust for panics and drop the right elements, we want
@@ -751,6 +747,44 @@ where D: Dimension
 
         Ok(())
     }
+
+    /// Reserve capacity to grow array along `axis` by at least `additional` elements.
+    ///
+    /// Existing elements of `array` are untouched and the backing storage is grown by
+    /// calling the underlying `reserve` method of the `OwnedRepr`.
+    ///
+    /// This is useful when pushing or appending repeatedly to an array to avoid multiple
+    /// allocations.
+    ///
+    /// ```rust
+    /// use ndarray::{Array3, Axis};
+    /// let mut a = Array3::<i32>::zeros((0,2,4));
+    /// a.reserve(Axis(0), 1000);
+    /// assert!(a.into_raw_vec().capacity() >= 2*4*1000);
+    pub fn reserve(&mut self, axis: Axis, additional: usize)
+    where D: RemoveAxis
+    {
+        if additional == 0 {
+            return;
+        }
+        let self_dim = self.raw_dim();
+        let remaining_shape = self_dim.remove_axis(axis);
+        let len_to_append = remaining_shape.size() * additional;
+
+        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);
+            self.ptr = self
+                .data
+                .reserve(len_to_append)
+                .offset(data_to_array_offset);
+        }
+    }
 }
 
 /// This drops all "unreachable" elements in `self_` given the data pointer and data length.

tests/reserve.rs

@@ -0,0 +1,52 @@
+use ndarray::prelude::*;
+
+#[test]
+fn reserve_1d()
+{
+    let mut a = Array1::<i32>::zeros((4,));
+    a.reserve(Axis(0), 1000);
+    assert_eq!(a.shape(), &[4]);
+    assert!(a.into_raw_vec().capacity() >= 1004);
+}
+
+#[test]
+fn reserve_3d()
+{
+    let mut a = Array3::<i32>::zeros((0, 4, 8));
+    a.reserve(Axis(0), 10);
+    assert_eq!(a.shape(), &[0, 4, 8]);
+    assert!(a.into_raw_vec().capacity() >= 4 * 8 * 10);
+}
+
+#[test]
+fn reserve_empty_3d()
+{
+    let mut a = Array3::<i32>::zeros((0, 0, 0));
+    a.reserve(Axis(0), 10);
+}
+
+#[test]
+fn reserve_3d_axis1()
+{
+    let mut a = Array3::<i32>::zeros((2, 4, 8));
+    a.reserve(Axis(1), 10);
+    assert!(a.into_raw_vec().capacity() >= 2 * 8 * 10);
+}
+
+#[test]
+fn reserve_3d_repeat()
+{
+    let mut a = Array3::<i32>::zeros((2, 4, 8));
+    a.reserve(Axis(1), 10);
+    a.reserve(Axis(2), 30);
+    assert!(a.into_raw_vec().capacity() >= 2 * 4 * 30);
+}
+
+#[test]
+fn reserve_2d_with_data()
+{
+    let mut a = array![[1, 2], [3, 4], [5, 6]];
+    a.reserve(Axis(1), 100);
+    assert_eq!(a, array![[1, 2], [3, 4], [5, 6]]);
+    assert!(a.into_raw_vec().capacity() >= 3 * 100);
+}