Add .shrink_to_fit() for Array

Add `.shrink_to_fit()` for owned storage arrays.
This initial version does not change array strides, this is to avoid the
most tricky part of the implementation (and maybe take it step by step).

Care must be taken since `self.ptr` points inside the allocation - which
for `Array` is still `Vec`-allocated, and several Vec methods including
its `shrink_to_fit` can reallocate the array.

Author: bluss

src/data_repr.rs

@@ -4,6 +4,7 @@ use alloc::borrow::ToOwned;
 use alloc::slice;
 #[cfg(not(feature = "std"))]
 use alloc::vec::Vec;
+use core::ops::Range;
 use std::mem;
 use std::mem::ManuallyDrop;
 use std::ptr::NonNull;
@@ -28,6 +29,9 @@ pub struct OwnedRepr<A>
     capacity: usize,
 }
 
+// OwnedRepr is a wrapper for a uniquely held allocation. Currently it is allocated by using a Vec
+// (from/to raw parts) which gives the benefit that it can always be converted to/from a Vec
+// cheaply.
 impl<A> OwnedRepr<A>
 {
     pub(crate) fn from(v: Vec<A>) -> Self
@@ -54,6 +58,14 @@ impl<A> OwnedRepr<A>
         self.len
     }
 
+    #[cfg(test)]
+    /// Note: Capacity comes from OwnedRepr (Vec)'s allocation strategy and cannot be absolutely
+    /// guaranteed.
+    pub(crate) fn capacity(&self) -> usize
+    {
+        self.capacity
+    }
+
     pub(crate) fn as_ptr(&self) -> *const A
     {
         self.ptr.as_ptr()
@@ -85,6 +97,36 @@ impl<A> OwnedRepr<A>
         self.as_nonnull_mut()
     }
 
+    /// Return the new lowest address pointer of the allocation.
+    #[must_use = "must use new pointer to update existing pointers"]
+    pub(crate) fn shrink_to_fit(&mut self) -> NonNull<A>
+    {
+        // Vec::shrink_to_fit is allowed to reallocate and invalidate pointers
+        self.modify_as_vec(|mut v| {
+            v.shrink_to_fit();
+            v
+        });
+        self.as_nonnull_mut()
+    }
+
+    /// Truncate "at front and back", preserve only elements inside the range,
+    /// then call shrink_to_fit().
+    /// Moving elements will invalidate existing pointers.
+    ///
+    /// Return the new lowest address pointer of the allocation.
+    #[must_use = "must use new pointer to update existing pointers"]
+    pub(crate) fn preserve_range_and_shrink(&mut self, span: Range<usize>) -> NonNull<A>
+    {
+        self.modify_as_vec(|mut v| {
+            v.truncate(span.end);
+            if span.start > 0 {
+                v.drain(..span.start);
+            }
+            v
+        });
+        self.shrink_to_fit()
+    }
+
     /// Set the valid length of the data
     ///
     /// ## Safety

src/dimension/mod.rs

@@ -428,7 +428,7 @@ fn to_abs_slice(axis_len: usize, slice: Slice) -> (usize, usize, isize)
 
 /// This function computes the offset from the lowest address element to the
 /// logically first element.
-pub fn offset_from_low_addr_ptr_to_logical_ptr<D: Dimension>(dim: &D, strides: &D) -> usize
+pub(crate) 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)| {
         let s = s as isize;
@@ -442,6 +442,22 @@ pub fn offset_from_low_addr_ptr_to_logical_ptr<D: Dimension>(dim: &D, strides: &
     offset as usize
 }
 
+/// This function computes the offset from the logically first element to the highest address
+/// element.
+pub(crate) fn offset_from_logical_ptr_to_high_addr_ptr<D: Dimension>(dim: &D, strides: &D) -> usize
+{
+    let offset = izip!(dim.slice(), strides.slice()).fold(0, |_offset, (&d, &s)| {
+        let s = s as isize;
+        if s > 0 && d > 1 {
+            _offset + s * (d as isize - 1)
+        } else {
+            _offset
+        }
+    });
+    debug_assert!(offset >= 0);
+    offset as usize
+}
+
 /// Modify dimension, stride and return data pointer offset
 ///
 /// **Panics** if stride is 0 or if any index is out of bounds.

src/impl_owned_array.rs

@@ -859,6 +859,61 @@ where D: Dimension
 
         Ok(())
     }
+
+    /// Shrink Array allocation capacity to be as small as it can be.
+    pub fn shrink_to_fit(&mut self)
+    {
+        // Example:
+        //                        (1)         (2)                 (3)              .- len
+        // Vector:     [ x x x x x V x V x V x V x V x V x V x V x V x x x x x x x ]       .- capacity
+        // Allocation: [ m m m m m m m m m m m m m m m m m m m m m m m m m m m m m m m m m ]
+        //
+        // x: valid data in OwnedRepr but outside current array slicing
+        // V: valid data in OwnedRepr and visible in current array slicing
+        // m: allocated memory
+        // (1): Lowest address element
+        // (2): Logical pointer (Element at index zero; normally (1) == (2) but can be
+        // located anywhere (1) <= (2) <= (3))
+        // (3): Highest address element
+        //
+        // Span: From (1) to (3).
+        //
+        // Algorithm: Compute 1, 2, 3.
+        // Move data so that unused areas before (1) and after (3) are removed from the storage/vector.
+        // Then shrink the vector's allocation to fit the valid elements.
+        //
+        // After:
+        //              (1)         (2)                 (3).- len == capacity
+        // Vector:     [ V x V x V x V x V x V x V x V x V ]
+        // Allocation: [ m m m m m m m m m m m m m m m m m ]
+        //
+
+        if mem::size_of::<A>() == 0 {
+            return;
+        }
+
+        let data_ptr = self.data.as_ptr();
+        let logical_ptr = self.as_ptr();
+        let offset_to_logical = dimension::offset_from_low_addr_ptr_to_logical_ptr(&self.dim, &self.strides);
+        let offset_to_high = dimension::offset_from_logical_ptr_to_high_addr_ptr(&self.dim, &self.strides);
+
+        let span = offset_to_logical + offset_to_high + 1;
+        debug_assert!(span >= self.len());
+
+        // We are in a panic critical section because: Array/OwnedRepr's destructors rely on
+        // dimension, strides, and self.ptr to deallocate correctly.
+        // We could panic here because custom user code is running when removing elements
+        // (destructors running).
+        let guard = AbortIfPanic(&"shrink_to_fit: owned repr not in consistent state");
+        unsafe {
+            let front_slop = logical_ptr.offset_from(data_ptr) as usize - offset_to_logical;
+            let new_low_ptr = self
+                .data
+                .preserve_range_and_shrink(front_slop..(front_slop + span));
+            self.ptr = new_low_ptr.add(offset_to_logical);
+        }
+        guard.defuse();
+    }
 }
 
 /// This drops all "unreachable" elements in `self_` given the data pointer and data length.
@@ -1016,3 +1071,70 @@ where D: Dimension
         }
     }
 }
+
+#[cfg(test)]
+mod tests
+{
+    use crate::Array;
+    use crate::Array2;
+    use crate::Slice;
+    use core::fmt::Debug;
+    use core::mem::size_of;
+
+    #[test]
+    fn test_shrink_to_fit()
+    {
+        fn assert_shrink_before_after<T>(mut a: Array2<T>, s1: Slice, s2: Slice, new_capacity: usize)
+        where T: Debug + Clone + Eq
+        {
+            let initial_len = a.len();
+            if size_of::<T>() > 0 {
+                assert_eq!(a.data.capacity(), initial_len);
+            }
+            a = a.slice_move(s![s1, s2]);
+            let before_value = a.clone();
+            let before_strides = a.strides().to_vec();
+            #[cfg(feature = "std")]
+            println!("{:?}, {}, {:?}", a, a.len(), a.data);
+            a.shrink_to_fit();
+            #[cfg(feature = "std")]
+            println!("{:?}, {}, {:?}", a, a.len(), a.data);
+
+            assert_eq!(before_value, a);
+            assert_eq!(before_strides, a.strides());
+
+            if size_of::<T>() > 0 {
+                assert!(a.data.capacity() < initial_len);
+                assert!(a.data.capacity() >= a.len());
+            }
+            assert_eq!(a.data.capacity(), new_capacity);
+        }
+
+        let a = Array::from_iter(0..56)
+            .into_shape_with_order((8, 7))
+            .unwrap();
+        assert_shrink_before_after(a, Slice::new(1, Some(-1), 1), Slice::new(0, None, 2), 42);
+
+        let a = Array::from_iter(0..56)
+            .into_shape_with_order((8, 7))
+            .unwrap();
+        assert_shrink_before_after(a, Slice::new(1, Some(-1), -1), Slice::new(0, None, -1), 42);
+
+        let a = Array::from_iter(0..56)
+            .into_shape_with_order((8, 7))
+            .unwrap();
+        assert_shrink_before_after(a, Slice::new(1, Some(3), 1), Slice::new(1, None, -2), 12);
+
+        // empty but still has some allocation to allow offsetting along each stride
+        let a = Array::from_iter(0..56)
+            .into_shape_with_order((8, 7))
+            .unwrap();
+        assert_shrink_before_after(a, Slice::new(1, Some(1), 1), Slice::new(1, None, 1), 6);
+
+        // Test ZST
+        let a = Array::from_iter((0..56).map(|_| ()))
+            .into_shape_with_order((8, 7))
+            .unwrap();
+        assert_shrink_before_after(a, Slice::new(1, Some(3), 1), Slice::new(1, None, -2), usize::MAX);
+    }
+}

tests/assign.rs

@@ -232,6 +232,39 @@ fn move_into()
     }
 }
 
+#[test]
+fn shrink_to_fit_slicing()
+{
+    // Count correct number of drops when using shrink_to_fit and discontiguous arrays (with holes).
+    for &use_f_order in &[false, true] {
+        for &invert_axis in &[0b00, 0b01, 0b10, 0b11] {
+            // bitmask for axis to invert
+            let counter = DropCounter::default();
+            {
+                let (m, n) = (5, 4);
+
+                let mut a = Array::from_shape_fn((m, n).set_f(use_f_order), |_idx| counter.element());
+                a.slice_collapse(s![1..-1, ..;2]);
+                if invert_axis & 0b01 != 0 {
+                    a.invert_axis(Axis(0));
+                }
+                if invert_axis & 0b10 != 0 {
+                    a.invert_axis(Axis(1));
+                }
+
+                a.shrink_to_fit();
+
+                let total = m * n;
+                let dropped_1 = if use_f_order { n * 2 - 1 } else { m * 2 - 1 };
+                assert_eq!(counter.created(), total);
+                assert_eq!(counter.dropped(), dropped_1 as usize);
+                drop(a);
+            }
+            counter.assert_drop_count();
+        }
+    }
+}
+
 /// This counter can create elements, and then count and verify
 /// the number of which have actually been dropped again.
 #[derive(Default)]
@@ -241,6 +274,7 @@ struct DropCounter
     dropped: AtomicUsize,
 }
 
+#[derive(Debug)]
 struct Element<'a>(&'a AtomicUsize);
 
 impl DropCounter