Add .shrink_to_fit() for Array

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,26 @@ impl<A> OwnedRepr<A>
         self.as_nonnull_mut()
     }
 
+    /// Truncate "at front and back", preserve only elements inside the range,
+    /// then call Vec::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);
+            }
+            // Vec::shrink_to_fit is allowed to reallocate and invalidate pointers
+            v.shrink_to_fit();
+            v
+        });
+        self.as_nonnull_mut()
+    }
+
     /// 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