@@ -414,44 +414,49 @@ where
414
414
where
415
415
I : CanSlice < D > + ?Sized ,
416
416
{
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
+ ) ;
420
422
let out_ndim = info. out_ndim ( ) ;
421
423
let mut new_dim = I :: OutDim :: zeros ( out_ndim) ;
422
424
let mut new_strides = I :: OutDim :: zeros ( out_ndim) ;
423
425
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) ;
450
457
451
458
// 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) }
455
460
}
456
461
457
462
/// Slice the array in place without changing the number of dimensions.
0 commit comments