From c4b0e7a1077fa3788517e18698daddf6b73a8f46 Mon Sep 17 00:00:00 2001
From: Ulrik Sverdrup <bluss@users.noreply.github.com>
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 7a3955fb7..c908766fc 100644
--- a/src/dimension/dimension_trait.rs
+++ b/src/dimension/dimension_trait.rs
@@ -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
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];