Add shrink_to_fit to Array

src/data_repr.rs

@@ -1,10 +1,10 @@
+use crate::extension::nonnull;
+use alloc::borrow::ToOwned;
+use alloc::slice;
+use alloc::vec::Vec;
 use std::mem;
 use std::mem::ManuallyDrop;
 use std::ptr::NonNull;
-use alloc::slice;
-use alloc::borrow::ToOwned;
-use alloc::vec::Vec;
-use crate::extension::nonnull;
 
 use rawpointer::PointerExt;
 
@@ -30,24 +30,20 @@ impl<A> OwnedRepr<A> {
         let len = v.len();
         let capacity = v.capacity();
         let ptr = nonnull::nonnull_from_vec_data(&mut v);
-        Self {
-            ptr,
-            len,
-            capacity,
-        }
+        Self { ptr, len, capacity }
     }
 
     pub(crate) fn into_vec(self) -> Vec<A> {
         ManuallyDrop::new(self).take_as_vec()
     }
 
     pub(crate) fn as_slice(&self) -> &[A] {
-        unsafe {
-            slice::from_raw_parts(self.ptr.as_ptr(), self.len)
-        }
+        unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
     }
 
-    pub(crate) fn len(&self) -> usize { self.len }
+    pub(crate) fn len(&self) -> usize {
+        self.len
+    }
 
     pub(crate) fn as_ptr(&self) -> *const A {
         self.ptr.as_ptr()
@@ -63,13 +59,11 @@ impl<A> OwnedRepr<A> {
 
     /// Return end pointer
     pub(crate) fn as_end_nonnull(&self) -> NonNull<A> {
-        unsafe {
-            self.ptr.add(self.len)
-        }
+        unsafe { self.ptr.add(self.len) }
     }
 
     /// Reserve `additional` elements; return the new pointer
-    /// 
+    ///
     /// ## Safety
     ///
     /// Note that existing pointers into the data are invalidated
@@ -82,6 +76,21 @@ impl<A> OwnedRepr<A> {
         self.as_nonnull_mut()
     }
 
+    /// Shrink the capacity of the array with lower bound.
+    /// The capacity will remain at least as large as both the length and
+    /// supplied value.
+    /// If the current capacity is less than the lower limit, this is a no-op.
+    pub(crate) fn shrink_to_fit(&mut self, len: usize) {
+        if len < self.len {
+            self.len = len;
+            self.modify_as_vec(|mut v| {
+                v.shrink_to_fit();
+                v
+            });
+            self.capacity = len;
+        }
+    }
+
     /// Set the valid length of the data
     ///
     /// ## Safety
@@ -126,14 +135,13 @@ impl<A> OwnedRepr<A> {
         let len = self.len;
         self.len = 0;
         self.capacity = 0;
-        unsafe {
-            Vec::from_raw_parts(self.ptr.as_ptr(), len, capacity)
-        }
+        unsafe { Vec::from_raw_parts(self.ptr.as_ptr(), len, capacity) }
     }
 }
 
 impl<A> Clone for OwnedRepr<A>
-    where A: Clone
+where
+    A: Clone,
 {
     fn clone(&self) -> Self {
         Self::from(self.as_slice().to_owned())
@@ -174,6 +182,5 @@ impl<A> Drop for OwnedRepr<A> {
     }
 }
 
-unsafe impl<A> Sync for OwnedRepr<A> where A: Sync { }
-unsafe impl<A> Send for OwnedRepr<A> where A: Send { }
-
+unsafe impl<A> Sync for OwnedRepr<A> where A: Sync {}
+unsafe impl<A> Send for OwnedRepr<A> where A: Send {}

src/impl_owned_array.rs

@@ -1,5 +1,5 @@
-
 use alloc::vec::Vec;
+use std::collections::HashSet;
 use std::mem;
 use std::mem::MaybeUninit;
 
@@ -11,6 +11,7 @@ use crate::dimension;
 use crate::error::{ErrorKind, ShapeError};
 use crate::iterators::Baseiter;
 use crate::low_level_util::AbortIfPanic;
+use crate::ArrayViewMut;
 use crate::OwnedRepr;
 use crate::Zip;
 
@@ -166,7 +167,8 @@ impl<A> Array<A, Ix2> {
 }
 
 impl<A, D> Array<A, D>
-    where D: Dimension
+where
+    D: Dimension,
 {
     /// Move all elements from self into `new_array`, which must be of the same shape but
     /// can have a different memory layout. The destination is overwritten completely.
@@ -198,9 +200,7 @@ impl<A, D> Array<A, D>
         } else {
             // If `A` doesn't need drop, we can overwrite the destination.
             // Safe because: move_into_uninit only writes initialized values
-            unsafe {
-                self.move_into_uninit(new_array.into_maybe_uninit())
-            }
+            unsafe { self.move_into_uninit(new_array.into_maybe_uninit()) }
         }
     }
 
@@ -209,7 +209,8 @@ impl<A, D> Array<A, D>
         // Afterwards, `self` drops full of initialized values and dropping works as usual.
         // This avoids moving out of owned values in `self` while at the same time managing
         // the dropping if the values being overwritten in `new_array`.
-        Zip::from(&mut self).and(new_array)
+        Zip::from(&mut self)
+            .and(new_array)
             .for_each(|src, dst| mem::swap(src, dst));
     }
 
@@ -401,16 +402,88 @@ impl<A, D> Array<A, D>
     ///            [0., 0., 0., 0.],
     ///            [1., 1., 1., 1.]]);
     /// ```
-    pub fn push(&mut self, axis: Axis, array: ArrayView<A, D::Smaller>)
-        -> Result<(), ShapeError>
+    pub fn push(&mut self, axis: Axis, array: ArrayView<A, D::Smaller>) -> Result<(), ShapeError>
     where
         A: Clone,
         D: RemoveAxis,
     {
         // same-dimensionality conversion
-        self.append(axis, array.insert_axis(axis).into_dimensionality::<D>().unwrap())
+        self.append(
+            axis,
+            array.insert_axis(axis).into_dimensionality::<D>().unwrap(),
+        )
     }
 
+    /// Return the offset between the pointer to the beginning of the heap
+    /// allocated by Array (`self.data.ptr`) and `self.ptr`.
+    unsafe fn offset_from_data(&self) -> isize {
+        if std::mem::size_of::<A>() != 0 {
+            self.as_ptr().offset_from(self.data.as_ptr())
+        } else {
+            0
+        }
+    }
+
+    /// Convert from any stride to the default stride
+    pub fn to_default_stride(&mut self) {
+        let view_mut =
+            unsafe { ArrayViewMut::new(self.ptr, self.dim.clone(), self.strides.clone()) };
+        let offset = unsafe { self.offset_from_data() };
+        unsafe { self.ptr.offset(offset) };
+        self.strides = self.dim.default_strides();
+        let mut index_ = match self.dim.first_index() {
+            Some(x) => x,
+            None => unreachable!(),
+        };
+        let mut swap_idx: Vec<(isize, isize)> = Vec::new();
+        loop {
+            let self_index = self
+                .dim
+                .stride_offset_checked(&self.strides, &index_)
+                .unwrap();
+            let view_mut_index = view_mut
+                .dim
+                .stride_offset_checked(&view_mut.strides, &index_)
+                .unwrap()
+                + offset;
+            swap_idx.push((
+                std::cmp::min(self_index, view_mut_index),
+                std::cmp::max(self_index, view_mut_index),
+            ));
+
+            index_ = match self.dim.next_for(index_) {
+                Some(x) => x,
+                None => {
+                    break;
+                }
+            };
+        }
+        let swap_idx = swap_idx
+            .into_iter()
+            .collect::<HashSet<_>>()
+            .into_iter()
+            .collect::<Vec<_>>();
+        for (x, y) in swap_idx.iter() {
+            unsafe { mem::swap(self.ptr.offset(*x).as_mut(), self.ptr.offset(*y).as_mut()) };
+        }
+    }
+
+    /// Shrinks the capacity of the array as much as possible.
+    ///
+    /// ```
+    /// use ndarray::array;
+    /// use ndarray::s;
+    ///
+    /// let a = array![[0, 1, 2], [3, 4, 5], [6, 7,8]];
+    /// let mut a = a.slice_move(s![.., 0..2]);
+    /// let b = a.clone();
+    /// a.shrink_to_fit();
+    /// assert_eq!(a, b);
+    /// ```
+    pub fn shrink_to_fit(&mut self) {
+        self.to_default_stride();
+        self.data.shrink_to_fit(self.len());
+    }
 
     /// Append an array to the array along an axis.
     ///
@@ -462,8 +535,7 @@ impl<A, D> Array<A, D>
     ///            [1., 1., 1., 1.],
     ///            [1., 1., 1., 1.]]);
     /// ```
-    pub fn append(&mut self, axis: Axis, mut array: ArrayView<A, D>)
-        -> Result<(), ShapeError>
+    pub fn append(&mut self, axis: Axis, mut array: ArrayView<A, D>) -> Result<(), ShapeError>
     where
         A: Clone,
         D: RemoveAxis,
@@ -556,7 +628,11 @@ impl<A, D> Array<A, D>
                     acc
                 } else {
                     let this_ax = ax.len as isize * ax.stride.abs();
-                    if this_ax > acc { this_ax } else { acc }
+                    if this_ax > acc {
+                        this_ax
+                    } else {
+                        acc
+                    }
                 }
             });
             let mut strides = self.strides.clone();
@@ -574,7 +650,10 @@ impl<A, D> Array<A, D>
                 0
             };
             debug_assert!(data_to_array_offset >= 0);
-            self.ptr = self.data.reserve(len_to_append).offset(data_to_array_offset);
+            self.ptr = self
+                .data
+                .reserve(len_to_append)
+                .offset(data_to_array_offset);
 
             // clone elements from view to the array now
             //
@@ -608,10 +687,13 @@ impl<A, D> Array<A, D>
 
             if tail_view.ndim() > 1 {
                 sort_axes_in_default_order_tandem(&mut tail_view, &mut array);
-                debug_assert!(tail_view.is_standard_layout(),
-                              "not std layout dim: {:?}, strides: {:?}",
-                              tail_view.shape(), tail_view.strides());
-            } 
+                debug_assert!(
+                    tail_view.is_standard_layout(),
+                    "not std layout dim: {:?}, strides: {:?}",
+                    tail_view.shape(),
+                    tail_view.strides()
+                );
+            }
 
             // Keep track of currently filled length of `self.data` and update it
             // on scope exit (panic or loop finish). This "indirect" way to
@@ -635,7 +717,6 @@ impl<A, D> Array<A, D>
                 data: &mut self.data,
             };
 
-
             // Safety: tail_view is constructed to have the same shape as array
             Zip::from(tail_view)
                 .and_unchecked(array)
@@ -665,8 +746,11 @@ impl<A, D> Array<A, D>
 ///
 /// This is an internal function for use by move_into and IntoIter only, safety invariants may need
 /// to be upheld across the calls from those implementations.
-pub(crate) unsafe fn drop_unreachable_raw<A, D>(mut self_: RawArrayViewMut<A, D>, data_ptr: *mut A, data_len: usize)
-where
+pub(crate) unsafe fn drop_unreachable_raw<A, D>(
+    mut self_: RawArrayViewMut<A, D>,
+    data_ptr: *mut A,
+    data_len: usize,
+) where
     D: Dimension,
 {
     let self_len = self_.len();
@@ -731,8 +815,11 @@ where
         dropped_elements += 1;
     }
 
-    assert_eq!(data_len, dropped_elements + self_len,
-               "Internal error: inconsistency in move_into");
+    assert_eq!(
+        data_len,
+        dropped_elements + self_len,
+        "Internal error: inconsistency in move_into"
+    );
 }
 
 /// Sort axes to standard order, i.e Axis(0) has biggest stride and Axis(n - 1) least stride
@@ -774,7 +861,6 @@ where
     }
 }
 
-
 /// Sort axes to standard order, i.e Axis(0) has biggest stride and Axis(n - 1) least stride
 ///
 /// Axes in a and b are sorted by the strides of `a`, and `a`'s axes should have stride >= 0 before