From 1e561b92f339bc90566301bf7506fdfdcdff6fdd Mon Sep 17 00:00:00 2001 From: Ulrik Sverdrup Date: Wed, 14 Aug 2024 21:55:02 +0200 Subject: [PATCH] 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. --- src/data_repr.rs | 42 ++++++++++++++ src/dimension/mod.rs | 18 +++++- src/impl_owned_array.rs | 118 ++++++++++++++++++++++++++++++++++++++++ tests/assign.rs | 34 ++++++++++++ 4 files changed, 211 insertions(+), 1 deletion(-) diff --git a/src/data_repr.rs b/src/data_repr.rs index 4041c192b..7d2ebc4a8 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,36 @@ impl OwnedRepr 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 + { + // 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) -> NonNull + { + 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 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..97e8dc36f 100644 --- a/src/impl_owned_array.rs +++ b/src/impl_owned_array.rs @@ -859,6 +859,57 @@ 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()); + + 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 +1067,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