diff --git a/benches/reserve.rs b/benches/reserve.rs index 9eef73b26..9e1ea8668 100644 --- a/benches/reserve.rs +++ b/benches/reserve.rs @@ -11,7 +11,7 @@ fn push_reserve(bench: &mut Bencher) { let ones: Array = array![1f32]; bench.iter(|| { let mut a: Array = array![]; - a.reserve(Axis(0), 100); + a.reserve(Axis(0), 100).unwrap(); for _ in 0..100 { a.append(Axis(0), ones.view()).unwrap(); } diff --git a/src/impl_owned_array.rs b/src/impl_owned_array.rs index 6bac049f6..42c7d4051 100644 --- a/src/impl_owned_array.rs +++ b/src/impl_owned_array.rs @@ -173,14 +173,18 @@ impl Array { /// This is useful when pushing or appending repeatedly to an array to avoid multiple /// allocations. /// + /// ***Errors*** with a shape error if the resultant capacity is larger than the addressable + /// bounds; that is, the product of non-zero axis lengths once `axis` has been extended by + /// `additional` exceeds `isize::MAX`. + /// /// ```rust /// use ndarray::Array2; /// let mut a = Array2::::zeros((2,4)); - /// a.reserve_rows(1000); + /// a.reserve_rows(1000).unwrap(); /// assert!(a.into_raw_vec().capacity() >= 4*1002); /// ``` - pub fn reserve_rows(&mut self, additional: usize) { - self.reserve(Axis(0), additional); + pub fn reserve_rows(&mut self, additional: usize) -> Result<(), ShapeError> { + self.reserve(Axis(0), additional) } /// Reserve capacity to grow array by at least `additional` columns. @@ -191,14 +195,18 @@ impl Array { /// This is useful when pushing or appending repeatedly to an array to avoid multiple /// allocations. /// + /// ***Errors*** with a shape error if the resultant capacity is larger than the addressable + /// bounds; that is, the product of non-zero axis lengths once `axis` has been extended by + /// `additional` exceeds `isize::MAX`. + /// /// ```rust /// use ndarray::Array2; /// let mut a = Array2::::zeros((2,4)); - /// a.reserve_columns(1000); + /// a.reserve_columns(1000).unwrap(); /// assert!(a.into_raw_vec().capacity() >= 2*1002); /// ``` - pub fn reserve_columns(&mut self, additional: usize) { - self.reserve(Axis(1), additional); + pub fn reserve_columns(&mut self, additional: usize) -> Result<(), ShapeError> { + self.reserve(Axis(1), additional) } } @@ -604,7 +612,7 @@ impl Array }; // grow backing storage and update head ptr - self.reserve(axis, array_dim[axis.index()]); + self.reserve(axis, array_dim[axis.index()])?; unsafe { // clone elements from view to the array now @@ -691,26 +699,43 @@ impl Array /// Reserve capacity to grow array along `axis` by at least `additional` elements. /// + /// The axis should be in the range `Axis(` 0 .. *n* `)` where *n* is the + /// number of dimensions (axes) of the array. + /// /// Existing elements of `array` are untouched and the backing storage is grown by /// calling the underlying `reserve` method of the `OwnedRepr`. /// /// This is useful when pushing or appending repeatedly to an array to avoid multiple /// allocations. /// + /// ***Panics*** if the axis is out of bounds. + /// + /// ***Errors*** with a shape error if the resultant capacity is larger than the addressable + /// bounds; that is, the product of non-zero axis lengths once `axis` has been extended by + /// `additional` exceeds `isize::MAX`. + /// /// ```rust /// use ndarray::{Array3, Axis}; /// let mut a = Array3::::zeros((0,2,4)); - /// a.reserve(Axis(0), 1000); + /// a.reserve(Axis(0), 1000).unwrap(); /// assert!(a.into_raw_vec().capacity() >= 2*4*1000); /// ``` - pub fn reserve(&mut self, axis: Axis, additional: usize) + /// + pub fn reserve(&mut self, axis: Axis, additional: usize) -> Result<(), ShapeError> where D: RemoveAxis, { + debug_assert!(axis.index() < self.ndim()); let self_dim = self.raw_dim(); let remaining_shape = self_dim.remove_axis(axis); let len_to_append = remaining_shape.size() * additional; + // Make sure new capacity is still in bounds + let mut res_dim = self_dim; + res_dim[axis.index()] += additional; + let new_len = dimension::size_of_shape_checked(&res_dim)?; + debug_assert_eq!(self.len() + len_to_append, new_len); + unsafe { // grow backing storage and update head ptr let data_to_array_offset = if std::mem::size_of::() != 0 { @@ -724,6 +749,10 @@ impl Array .reserve(len_to_append) .offset(data_to_array_offset); } + + debug_assert!(self.pointer_is_inbounds()); + + Ok(()) } } diff --git a/tests/reserve.rs b/tests/reserve.rs index 64bfc2865..cc6b22b1f 100644 --- a/tests/reserve.rs +++ b/tests/reserve.rs @@ -3,7 +3,7 @@ use ndarray::prelude::*; #[test] fn reserve_1d() { let mut a = Array1::::zeros((4,)); - a.reserve(Axis(0), 1000); + a.reserve(Axis(0), 1000).unwrap(); assert_eq!(a.shape(), &[4]); assert!(a.into_raw_vec().capacity() >= 1004); } @@ -11,7 +11,7 @@ fn reserve_1d() { #[test] fn reserve_3d() { let mut a = Array3::::zeros((0, 4, 8)); - a.reserve(Axis(0), 10); + a.reserve(Axis(0), 10).unwrap(); assert_eq!(a.shape(), &[0, 4, 8]); assert!(a.into_raw_vec().capacity() >= 4 * 8 * 10); } @@ -19,28 +19,28 @@ fn reserve_3d() { #[test] fn reserve_empty_3d() { let mut a = Array3::::zeros((0, 0, 0)); - a.reserve(Axis(0), 10); + a.reserve(Axis(0), 10).unwrap(); } #[test] fn reserve_3d_axis1() { let mut a = Array3::::zeros((2, 4, 8)); - a.reserve(Axis(1), 10); + a.reserve(Axis(1), 10).unwrap(); assert!(a.into_raw_vec().capacity() >= 2 * 8 * 10); } #[test] fn reserve_3d_repeat() { let mut a = Array3::::zeros((2, 4, 8)); - a.reserve(Axis(1), 10); - a.reserve(Axis(2), 30); + a.reserve(Axis(1), 10).unwrap(); + a.reserve(Axis(2), 30).unwrap(); assert!(a.into_raw_vec().capacity() >= 2 * 4 * 30); } #[test] fn reserve_2d_with_data() { let mut a = array![[1, 2], [3, 4], [5, 6]]; - a.reserve(Axis(1), 100); + a.reserve(Axis(1), 100).unwrap(); assert_eq!(a, array![[1, 2], [3, 4], [5, 6]]); assert!(a.into_raw_vec().capacity() >= 3 * 100); }