Skip to content

Commit c360b87

Browse files
committed
Fix slice offset computation for special cases
With the old implementation, this resulted in undefined behavior in release mode and a panic in debug mode: ```rust let mut arr = Array2::<i32>::zeros((5, 5)); arr.slice_axis_inplace(Axis(0), Slice::new(0, Some(0), -1)); ``` as did this: ```rust let mut arr = Array2::from_shape_vec((1, 1).strides((10, 1)), vec![5]).unwrap(); arr.slice_axis_inplace(Axis(0), Slice::new(1, Some(1), 1)); ``` Now, both examples operate correctly.
1 parent cedbdf9 commit c360b87

File tree

2 files changed

+33
-12
lines changed

2 files changed

+33
-12
lines changed

src/dimension/mod.rs

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -367,18 +367,29 @@ pub fn do_slice(dim: &mut usize, stride: &mut usize, slice: Slice) -> isize {
367367
let m = end - start;
368368
let s = (*stride) as isize;
369369

370-
// Data pointer offset
371-
let mut offset = stride_offset(start, *stride);
372-
// Adjust for strides
373-
//
374-
// How to implement negative strides:
375-
//
376-
// Increase start pointer by
377-
// old stride * (old dim - 1)
378-
// to put the pointer completely in the other end
379-
if step < 0 {
380-
offset += stride_offset(m - 1, *stride);
381-
}
370+
// Compute data pointer offset.
371+
let offset = if m == 0 {
372+
// In this case, the resulting array is empty, so we *can* avoid performing a nonzero
373+
// offset.
374+
//
375+
// In two special cases (which are the true reason for this `m == 0` check), we *must* avoid
376+
// the nonzero offset corresponding to the general case.
377+
//
378+
// * When `end == 0 && step < 0`. (These conditions imply that `m == 0` since `to_abs_slice`
379+
// ensures that `0 <= start <= end`.) We cannot execute `stride_offset(end - 1, *stride)`
380+
// because the `end - 1` would underflow.
381+
//
382+
// * When `start == *dim && step > 0`. (These conditions imply that `m == 0` since
383+
// `to_abs_slice` ensures that `start <= end <= *dim`.) We cannot use the offset returned
384+
// by `stride_offset(start, *stride)` because that would be past the end of the axis.
385+
0
386+
} else if step < 0 {
387+
// When the step is negative, the new first element is `end - 1`, not `start`, since the
388+
// direction is reversed.
389+
stride_offset(end - 1, *stride)
390+
} else {
391+
stride_offset(start, *stride)
392+
};
382393

383394
// Update dimension.
384395
let abs_step = step.abs() as usize;

tests/array.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,16 @@ fn test_slice()
9696
assert!(vi.iter().zip(A.iter()).all(|(a, b)| a == b));
9797
}
9898

99+
#[test]
100+
fn test_slice_edge_cases() {
101+
let mut arr = Array3::<u8>::zeros((3, 4, 5));
102+
arr.slice_collapse(s![0..0;-1, .., ..]);
103+
assert_eq!(arr.shape(), &[0, 4, 5]);
104+
let mut arr = Array2::<u8>::from_shape_vec((1, 1).strides((10, 1)), vec![5]).unwrap();
105+
arr.slice_collapse(s![1..1, ..]);
106+
assert_eq!(arr.shape(), &[0, 1]);
107+
}
108+
99109
#[test]
100110
fn test_slice_inclusive_range() {
101111
let arr = array![[1, 2, 3], [4, 5, 6]];

0 commit comments

Comments
 (0)