Skip to content

Commit 815e708

Browse files
jturner314bluss
authored andcommitted
Make slice_move not call slice_collapse
This isn't much more code and simplifies the logic somewhat.
1 parent 3ba6ceb commit 815e708

File tree

1 file changed

+37
-32
lines changed

1 file changed

+37
-32
lines changed

src/impl_methods.rs

Lines changed: 37 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -414,44 +414,49 @@ where
414414
where
415415
I: CanSlice<D> + ?Sized,
416416
{
417-
// Slice and collapse in-place without changing the number of dimensions.
418-
self.slice_collapse(info);
419-
417+
assert_eq!(
418+
info.in_ndim(),
419+
self.ndim(),
420+
"The input dimension of `info` must match the array to be sliced.",
421+
);
420422
let out_ndim = info.out_ndim();
421423
let mut new_dim = I::OutDim::zeros(out_ndim);
422424
let mut new_strides = I::OutDim::zeros(out_ndim);
423425

424-
// Write the dim and strides to the correct new axes.
425-
{
426-
let mut old_axis = 0;
427-
let mut new_axis = 0;
428-
info.as_ref().iter().for_each(|ax_info| match ax_info {
429-
AxisSliceInfo::Slice { .. } => {
430-
// Copy the old dim and stride to corresponding axis.
431-
new_dim[new_axis] = self.dim[old_axis];
432-
new_strides[new_axis] = self.strides[old_axis];
433-
old_axis += 1;
434-
new_axis += 1;
435-
}
436-
AxisSliceInfo::Index(_) => {
437-
// Skip the old axis since it should be removed.
438-
old_axis += 1;
439-
}
440-
AxisSliceInfo::NewAxis => {
441-
// Set the dim and stride of the new axis.
442-
new_dim[new_axis] = 1;
443-
new_strides[new_axis] = 0;
444-
new_axis += 1;
445-
}
446-
});
447-
debug_assert_eq!(old_axis, self.ndim());
448-
debug_assert_eq!(new_axis, out_ndim);
449-
}
426+
let mut old_axis = 0;
427+
let mut new_axis = 0;
428+
info.as_ref().iter().for_each(|&ax_info| match ax_info {
429+
AxisSliceInfo::Slice { start, end, step } => {
430+
// Slice the axis in-place to update the `dim`, `strides`, and `ptr`.
431+
self.slice_axis_inplace(Axis(old_axis), Slice { start, end, step });
432+
// Copy the sliced dim and stride to corresponding axis.
433+
new_dim[new_axis] = self.dim[old_axis];
434+
new_strides[new_axis] = self.strides[old_axis];
435+
old_axis += 1;
436+
new_axis += 1;
437+
}
438+
AxisSliceInfo::Index(index) => {
439+
// Collapse the axis in-place to update the `ptr`.
440+
let i_usize = abs_index(self.len_of(Axis(old_axis)), index);
441+
self.collapse_axis(Axis(old_axis), i_usize);
442+
// Skip copying the axis since it should be removed. Note that
443+
// removing this axis is safe because `.collapse_axis()` panics
444+
// if the index is out-of-bounds, so it will panic if the axis
445+
// is zero length.
446+
old_axis += 1;
447+
}
448+
AxisSliceInfo::NewAxis => {
449+
// Set the dim and stride of the new axis.
450+
new_dim[new_axis] = 1;
451+
new_strides[new_axis] = 0;
452+
new_axis += 1;
453+
}
454+
});
455+
debug_assert_eq!(old_axis, self.ndim());
456+
debug_assert_eq!(new_axis, out_ndim);
450457

451458
// safe because new dimension, strides allow access to a subset of old data
452-
unsafe {
453-
self.with_strides_dim(new_strides, new_dim)
454-
}
459+
unsafe { self.with_strides_dim(new_strides, new_dim) }
455460
}
456461

457462
/// Slice the array in place without changing the number of dimensions.

0 commit comments

Comments
 (0)