extended tests

Author: flexatone

src/lib.rs

@@ -180,6 +180,80 @@ fn first_true_1d_f(py: Python, array: PyReadonlyArray1<bool>) -> isize {
 }
 
 
+#[pyfunction]
+#[pyo3(signature = (array, forward=true))]
+fn first_true_1d(py: Python, array: PyReadonlyArray1<bool>, forward: bool) -> isize {
+    if let Ok(slice) = array.as_slice() {
+        py.allow_threads(|| {
+            let len = slice.len();
+            let ptr = slice.as_ptr() as *const u8;
+            let ones = u8x32::splat(1);
+
+            if forward {
+                let mut i = 0;
+                unsafe {
+                    // Process 32 bytes at a time with SIMD
+                    while i + 32 <= len {
+                        let bytes = &*(ptr.add(i) as *const [u8; 32]);
+                        let chunk = u8x32::from(*bytes);
+                        let equal_one = chunk.cmp_eq(ones);
+                        if equal_one.any() {
+                            break;
+                        }
+                        i += 32;
+                    }
+                    // Handle final remainder
+                    while i < len.min(i + 32) {
+                        if *ptr.add(i) != 0 {
+                            return i as isize;
+                        }
+                        i += 1;
+                    }
+                }
+            } else {
+                // Backward search
+                let mut i = len;
+                unsafe {
+                    // Process 32 bytes at a time with SIMD (backwards)
+                    while i >= 32 {
+                        i -= 32;
+                        let bytes = &*(ptr.add(i) as *const [u8; 32]);
+                        let chunk = u8x32::from(*bytes);
+                        let equal_one = chunk.cmp_eq(ones);
+                        if equal_one.any() {
+                            // Found a true in this chunk, search backwards within it
+                            for j in (0..32).rev() {
+                                if i + j < len && *ptr.add(i + j) != 0 {
+                                    return (i + j) as isize;
+                                }
+                            }
+                        }
+                    }
+                    // Handle remaining bytes at the beginning
+                    if i > 0 {
+                        for j in (0..i).rev() {
+                            if *ptr.add(j) != 0 {
+                                return j as isize;
+                            }
+                        }
+                    }
+                }
+            }
+            -1
+        })
+    } else {
+        let array_view = array.as_array();
+        py.allow_threads(|| {
+            if forward {
+                array_view.iter().position(|&v| v).map(|i| i as isize).unwrap_or(-1)
+            } else {
+                array_view.iter().rposition(|&v| v).map(|i| i as isize).unwrap_or(-1)
+            }
+        })
+    }
+}
+
+
 #[pymodule]
 fn arrayredox(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(first_true_1d_a, m)?)?;
@@ -188,5 +262,6 @@ fn arrayredox(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(first_true_1d_d, m)?)?;
     m.add_function(wrap_pyfunction!(first_true_1d_e, m)?)?;
     m.add_function(wrap_pyfunction!(first_true_1d_f, m)?)?;
+    m.add_function(wrap_pyfunction!(first_true_1d, m)?)?;
     Ok(())
 }

tests/test_basic.py

@@ -10,4 +10,71 @@ def test_first_true_1d():
     assert first_true_1d(b) == -1
 
 
+#--------------------------------------------------------------------------
+def test_first_true_1d_a() -> None:
+    a1 = np.arange(100) == 50
+    post = first_true_1d(a1, forward=True)
+    assert post == 50
 
+# def test_first_true_1d_b() -> None:
+#     with self.assertRaises(TypeError):
+#         a1 = [2, 4, 5,]
+#         first_true_1d(a1, forward=True)
+
+# def test_first_true_1d_c() -> None:
+#     with self.assertRaises(ValueError):
+#         a1 = np.arange(100) == 50
+#         first_true_1d(a1, forward=a1)
+
+def test_first_true_1d_d() -> None:
+    a1 = np.arange(100) < 0
+    post = first_true_1d(a1, forward=True)
+    assert post == -1
+
+# def test_first_true_1d_e() -> None:
+#     a1 = np.arange(100)
+#     # only a Boolean array
+#     with self.assertRaises(ValueError):
+#         post = first_true_1d(a1, forward=True)
+
+# def test_first_true_1d_f() -> None:
+#     a1 = (np.arange(100) == 0)[:50:2]
+#     # only a contiguous array
+#     with self.assertRaises(ValueError):
+#         post = first_true_1d(a1, forward=True)
+
+# def test_first_true_1d_g() -> None:
+#     a1 = (np.arange(100) == 0).reshape(10, 10)
+#     # only a contiguous array
+#     with self.assertRaises(ValueError):
+#         post = first_true_1d(a1, forward=True)
+
+def test_first_true_1d_reverse_a() -> None:
+    a1 = np.arange(100) == 50
+    post = first_true_1d(a1, forward=False)
+    assert post == 50
+
+def test_first_true_1d_reverse_b() -> None:
+    a1 = np.arange(100) == 0
+    post = first_true_1d(a1, forward=False)
+    assert post == 0
+
+def test_first_true_1d_reverse_c() -> None:
+    a1 = np.arange(100) == -1
+    post = first_true_1d(a1, forward=False)
+    assert post == -1
+
+def test_first_true_1d_reverse_d() -> None:
+    a1 = np.arange(100) == 99
+    post = first_true_1d(a1, forward=False)
+    assert post == 99
+
+def test_first_true_1d_multi_a() -> None:
+    a1 = np.isin(np.arange(100), (50, 70, 90))
+    assert first_true_1d(a1, forward=True) == 50
+    assert first_true_1d(a1, forward=False) == 90
+
+def test_first_true_1d_multi_b() -> None:
+    a1 = np.isin(np.arange(100), (10, 30, 50))
+    assert first_true_1d(a1, forward=True) == 10
+    assert first_true_1d(a1, forward=False) == 50