From c03345377227604ec5a26a2ba22d4b5a2ebe3070 Mon Sep 17 00:00:00 2001 From: Ulrik Sverdrup Date: Mon, 26 Feb 2024 20:20:03 +0100 Subject: [PATCH] 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. --- src/dimension/dimension_trait.rs | 30 +++++++++++++++++------------- tests/array.rs | 16 ++++++++++++++++ 2 files changed, 33 insertions(+), 13 deletions(-) diff --git a/src/dimension/dimension_trait.rs b/src/dimension/dimension_trait.rs index 14ec2cd31..c594859e9 100644 --- a/src/dimension/dimension_trait.rs +++ b/src/dimension/dimension_trait.rs @@ -285,21 +285,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 diff --git a/tests/array.rs b/tests/array.rs index 8fdbb9992..d770d7822 100644 --- a/tests/array.rs +++ b/tests/array.rs @@ -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];