Skip to content

Commit

Permalink
dimension: Fix contig check for single element arrays
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
bluss committed Feb 26, 2024
1 parent 5993456 commit c4b0e7a
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 13 deletions.
30 changes: 17 additions & 13 deletions src/dimension/dimension_trait.rs
Original file line number Diff line number Diff line change
Expand Up @@ -284,21 +284,25 @@ pub trait Dimension:
return true;
}
if dim.ndim() == 1 {
return strides[0] as isize == -1;
}
let order = strides._fastest_varying_stride_order();
let strides = strides.slice();

let dim_slice = dim.slice();
let mut cstride = 1;
for &i in order.slice() {
// a dimension of length 1 can have unequal strides
if dim_slice[i] != 1 && (strides[i] as isize).unsigned_abs() != cstride {
return false;
// fast case for ndim == 1:
// Either we have length <= 1, then stride is arbitrary,
// or we have stride == 1 or stride == -1, but +1 case is already handled above.
dim[0] <= 1 || strides[0] as isize == -1
} else {
let order = strides._fastest_varying_stride_order();
let strides = strides.slice();

let dim_slice = dim.slice();
let mut cstride = 1;
for &i in order.slice() {
// a dimension of length 1 can have unequal strides
if dim_slice[i] != 1 && (strides[i] as isize).unsigned_abs() != cstride {
return false;
}
cstride *= dim_slice[i];
}
cstride *= dim_slice[i];
true
}
true
}

/// Return the axis ordering corresponding to the fastest variation
Expand Down
16 changes: 16 additions & 0 deletions tests/array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1957,6 +1957,22 @@ fn test_contiguous() {
assert!(b.as_slice_memory_order().is_some());
}

#[test]
fn test_contiguous_single_element()
{
assert_matches!(array![1].as_slice_memory_order(), Some(&[1]));

let arr1 = array![1, 2, 3];
assert_matches!(arr1.slice(s![0..1]).as_slice_memory_order(), Some(&[1]));
assert_matches!(arr1.slice(s![1..2]).as_slice_memory_order(), Some(&[2]));
assert_matches!(arr1.slice(s![2..3]).as_slice_memory_order(), Some(&[3]));
assert_matches!(arr1.slice(s![0..0]).as_slice_memory_order(), Some(&[]));

let arr2 = array![[1, 2, 3], [4, 5, 6]];
assert_matches!(arr2.slice(s![.., 2..3]).as_slice_memory_order(), None);
assert_matches!(arr2.slice(s![1, 2..3]).as_slice_memory_order(), Some(&[6]));
}

#[test]
fn test_contiguous_neg_strides() {
let s = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13];
Expand Down

0 comments on commit c4b0e7a

Please sign in to comment.