diff --git a/src/data_repr.rs b/src/data_repr.rs
index 4041c192b..09a4a8f94 100644
--- a/src/data_repr.rs
+++ b/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
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 OwnedRepr
{
pub(crate) fn from(v: Vec) -> Self
@@ -54,6 +58,14 @@ impl OwnedRepr
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 OwnedRepr
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) -> NonNull
+ {
+ 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
diff --git a/src/dimension/mod.rs b/src/dimension/mod.rs
index 601f0dc43..4f927654f 100644
--- a/src/dimension/mod.rs
+++ b/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(dim: &D, strides: &D) -> usize
+pub(crate) fn offset_from_low_addr_ptr_to_logical_ptr(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(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(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.
diff --git a/src/impl_owned_array.rs b/src/impl_owned_array.rs
index 44ac12dd4..0776b8485 100644
--- a/src/impl_owned_array.rs
+++ b/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::() == 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(mut a: Array2, s1: Slice, s2: Slice, new_capacity: usize)
+ where T: Debug + Clone + Eq
+ {
+ let initial_len = a.len();
+ if size_of::() > 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::() > 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);
+ }
+}
diff --git a/tests/assign.rs b/tests/assign.rs
index 29a6b851a..4fbcabe2c 100644
--- a/tests/assign.rs
+++ b/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