Skip to content

Commit c033453

Browse files
committed
dimension: Fix contig check for single element arrays
When an array has 0 or 1 elements, strides don't matter anymore. The general case of this function handled this correctly, but the special case for ndim == 1 did not.
1 parent e8a874e commit c033453

File tree

2 files changed

+33
-13
lines changed

2 files changed

+33
-13
lines changed

src/dimension/dimension_trait.rs

+17-13
Original file line numberDiff line numberDiff line change
@@ -285,21 +285,25 @@ pub trait Dimension:
285285
return true;
286286
}
287287
if dim.ndim() == 1 {
288-
return strides[0] as isize == -1;
289-
}
290-
let order = strides._fastest_varying_stride_order();
291-
let strides = strides.slice();
292-
293-
let dim_slice = dim.slice();
294-
let mut cstride = 1;
295-
for &i in order.slice() {
296-
// a dimension of length 1 can have unequal strides
297-
if dim_slice[i] != 1 && (strides[i] as isize).unsigned_abs() != cstride {
298-
return false;
288+
// fast case for ndim == 1:
289+
// Either we have length <= 1, then stride is arbitrary,
290+
// or we have stride == 1 or stride == -1, but +1 case is already handled above.
291+
dim[0] <= 1 || strides[0] as isize == -1
292+
} else {
293+
let order = strides._fastest_varying_stride_order();
294+
let strides = strides.slice();
295+
296+
let dim_slice = dim.slice();
297+
let mut cstride = 1;
298+
for &i in order.slice() {
299+
// a dimension of length 1 can have unequal strides
300+
if dim_slice[i] != 1 && (strides[i] as isize).unsigned_abs() != cstride {
301+
return false;
302+
}
303+
cstride *= dim_slice[i];
299304
}
300-
cstride *= dim_slice[i];
305+
true
301306
}
302-
true
303307
}
304308

305309
/// Return the axis ordering corresponding to the fastest variation

tests/array.rs

+16
Original file line numberDiff line numberDiff line change
@@ -1957,6 +1957,22 @@ fn test_contiguous() {
19571957
assert!(b.as_slice_memory_order().is_some());
19581958
}
19591959

1960+
#[test]
1961+
fn test_contiguous_single_element()
1962+
{
1963+
assert_matches!(array![1].as_slice_memory_order(), Some(&[1]));
1964+
1965+
let arr1 = array![1, 2, 3];
1966+
assert_matches!(arr1.slice(s![0..1]).as_slice_memory_order(), Some(&[1]));
1967+
assert_matches!(arr1.slice(s![1..2]).as_slice_memory_order(), Some(&[2]));
1968+
assert_matches!(arr1.slice(s![2..3]).as_slice_memory_order(), Some(&[3]));
1969+
assert_matches!(arr1.slice(s![0..0]).as_slice_memory_order(), Some(&[]));
1970+
1971+
let arr2 = array![[1, 2, 3], [4, 5, 6]];
1972+
assert_matches!(arr2.slice(s![.., 2..3]).as_slice_memory_order(), None);
1973+
assert_matches!(arr2.slice(s![1, 2..3]).as_slice_memory_order(), Some(&[6]));
1974+
}
1975+
19601976
#[test]
19611977
fn test_contiguous_neg_strides() {
19621978
let s = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13];

0 commit comments

Comments
 (0)