Skip to content

Commit

Permalink
Update into_raw_vec
Browse files Browse the repository at this point in the history
Without the offset from start of allocation to the logically first
element of the array, correctly reinterpreting the results of
.into_raw_vec() as an n-D array is tricky.
  • Loading branch information
jturner314 authored and bluss committed Apr 1, 2024
1 parent 77332b1 commit 6d04ebd
Show file tree
Hide file tree
Showing 5 changed files with 88 additions and 18 deletions.
2 changes: 1 addition & 1 deletion examples/sort-axis.rs
Expand Up @@ -157,7 +157,7 @@ where D: Dimension
});
debug_assert_eq!(result.len(), moved_elements);
// forget the old elements but not the allocation
let mut old_storage = self.into_raw_vec();
let mut old_storage = self.into_raw_vec().0;
old_storage.set_len(0);

// transfer ownership of the elements into the result
Expand Down
4 changes: 2 additions & 2 deletions src/dimension/mod.rs
Expand Up @@ -409,8 +409,8 @@ fn to_abs_slice(axis_len: usize, slice: Slice) -> (usize, usize, isize)
(start, end, step)
}

/// Returns the offset from the lowest-address element to the logically first
/// element.
/// This function computes the offset from the lowest address element to the
/// logically first element. The result is always >= 0.
pub 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)| {
Expand Down
90 changes: 80 additions & 10 deletions src/impl_owned_array.rs
Expand Up @@ -59,14 +59,89 @@ impl<A> Array<A, Ix0>
impl<A, D> Array<A, D>
where D: Dimension
{
/// Returns the offset (in units of `A`) from the start of the allocation
/// to the first element, or `None` if the array is empty.
fn offset_from_alloc_to_logical_ptr(&self) -> Option<usize>
{
if self.is_empty() {
return None;
}
if std::mem::size_of::<A>() == 0 {
Some(dimension::offset_from_low_addr_ptr_to_logical_ptr(&self.dim, &self.strides))
} else {
let offset = unsafe { self.as_ptr().offset_from(self.data.as_ptr()) };
debug_assert!(offset >= 0);
Some(offset as usize)
}
}

/// Return a vector of the elements in the array, in the way they are
/// stored internally.
/// stored internally, and the index in the vector corresponding to the
/// logically first element of the array (or `None` if the array is empty).
///
/// If the array is in standard memory layout, the logical element order
/// of the array (`.iter()` order) and of the returned vector will be the same.
pub fn into_raw_vec(self) -> Vec<A>
///
/// ```
/// use ndarray::{array, Array2, Axis};
///
/// let mut arr: Array2<f64> = array![[1., 2.], [3., 4.], [5., 6.]];
/// arr.slice_axis_inplace(Axis(0), (1..).into());
/// assert_eq!(arr[[0, 0]], 3.);
/// let copy = arr.clone();
///
/// let shape = arr.shape().to_owned();
/// let strides = arr.strides().to_owned();
/// let (v, offset) = arr.into_raw_vec();
///
/// assert_eq!(v, &[1., 2., 3., 4., 5., 6.]);
/// assert_eq!(offset, Some(2));
/// assert_eq!(v[offset.unwrap()], 3.);
/// for row in 0..shape[0] {
/// for col in 0..shape[1] {
/// let index = (
/// offset.unwrap() as isize
/// + row as isize * strides[0]
/// + col as isize * strides[1]
/// ) as usize;
/// assert_eq!(v[index], copy[[row, col]]);
/// }
/// }
/// ```
///
/// In the case of zero-sized elements, the offset to the logically first
/// element is somewhat meaningless. For convenience, an offset will be
/// returned such that all indices computed using the offset, shape, and
/// strides will be in-bounds for the `Vec<A>`. Note that this offset won't
/// necessarily be the same as the offset for an array of nonzero-sized
/// elements sliced in the same way.
///
/// ```
/// use ndarray::{array, Array2, Axis};
///
/// let mut arr: Array2<()> = array![[(), ()], [(), ()], [(), ()]];
/// arr.slice_axis_inplace(Axis(0), (1..).into());
///
/// let shape = arr.shape().to_owned();
/// let strides = arr.strides().to_owned();
/// let (v, offset) = arr.into_raw_vec();
///
/// assert_eq!(v, &[(), (), (), (), (), ()]);
/// for row in 0..shape[0] {
/// for col in 0..shape[1] {
/// let index = (
/// offset.unwrap() as isize
/// + row as isize * strides[0]
/// + col as isize * strides[1]
/// ) as usize;
/// assert_eq!(v[index], ());
/// }
/// }
/// ```
pub fn into_raw_vec(self) -> (Vec<A>, Option<usize>)
{
self.data.into_vec()
let offset = self.offset_from_alloc_to_logical_ptr();
(self.data.into_vec(), offset)
}
}

Expand Down Expand Up @@ -575,16 +650,11 @@ where D: Dimension

unsafe {
// grow backing storage and update head ptr
let data_to_array_offset = if std::mem::size_of::<A>() != 0 {
self.as_ptr().offset_from(self.data.as_ptr())
} else {
0
};
debug_assert!(data_to_array_offset >= 0);
let offset_from_alloc_to_logical = self.offset_from_alloc_to_logical_ptr().unwrap_or(0);
self.ptr = self
.data
.reserve(len_to_append)
.offset(data_to_array_offset);
.add(offset_from_alloc_to_logical);

// clone elements from view to the array now
//
Expand Down
4 changes: 2 additions & 2 deletions tests/array-construct.rs
Expand Up @@ -19,9 +19,9 @@ fn test_from_shape_fn()
fn test_dimension_zero()
{
let a: Array2<f32> = Array2::from(vec![[], [], []]);
assert_eq!(vec![0.; 0], a.into_raw_vec());
assert_eq!(vec![0.; 0], a.into_raw_vec().0);
let a: Array3<f32> = Array3::from(vec![[[]], [[]], [[]]]);
assert_eq!(vec![0.; 0], a.into_raw_vec());
assert_eq!(vec![0.; 0], a.into_raw_vec().0);
}

#[test]
Expand Down
6 changes: 3 additions & 3 deletions tests/array.rs
Expand Up @@ -1157,7 +1157,7 @@ fn array0_into_scalar()
// With this kind of setup, the `Array`'s pointer is not the same as the
// underlying `Vec`'s pointer.
let a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
assert_ne!(a.as_ptr(), a.into_raw_vec().as_ptr());
assert_ne!(a.as_ptr(), a.into_raw_vec().0.as_ptr());
// `.into_scalar()` should still work correctly.
let a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
assert_eq!(a.into_scalar(), 6);
Expand All @@ -1173,7 +1173,7 @@ fn array_view0_into_scalar()
// With this kind of setup, the `Array`'s pointer is not the same as the
// underlying `Vec`'s pointer.
let a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
assert_ne!(a.as_ptr(), a.into_raw_vec().as_ptr());
assert_ne!(a.as_ptr(), a.into_raw_vec().0.as_ptr());
// `.into_scalar()` should still work correctly.
let a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
assert_eq!(a.view().into_scalar(), &6);
Expand All @@ -1189,7 +1189,7 @@ fn array_view_mut0_into_scalar()
// With this kind of setup, the `Array`'s pointer is not the same as the
// underlying `Vec`'s pointer.
let a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
assert_ne!(a.as_ptr(), a.into_raw_vec().as_ptr());
assert_ne!(a.as_ptr(), a.into_raw_vec().0.as_ptr());
// `.into_scalar()` should still work correctly.
let mut a: Array0<i32> = array![4, 5, 6, 7].index_axis_move(Axis(0), 2);
assert_eq!(a.view_mut().into_scalar(), &6);
Expand Down

0 comments on commit 6d04ebd

Please sign in to comment.