Skip to content

Commit c68d2b1

Browse files
flexatoneflexatone
committed
updated
1 parent ae90079 commit c68d2b1

File tree

1 file changed

+96
-34
lines changed

1 file changed

+96
-34
lines changed

‎src/lib.rs

Lines changed: 96 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,15 @@
1+
use numpy::PyReadonlyArray1;
12
use pyo3::prelude::*;
23
use wide::*;
3-
use numpy::{PyReadonlyArray1};
44

55
#[pyfunction]
66
fn first_true_1d_a(array: PyReadonlyArray1<bool>) -> isize {
77
match array.as_slice() {
8-
Ok(slice) => slice.iter().position(|&v| v).map(|i| i as isize).unwrap_or(-1),
8+
Ok(slice) => slice
9+
.iter()
10+
.position(|&v| v)
11+
.map(|i| i as isize)
12+
.unwrap_or(-1),
913
Err(_) => -1, // Should not happen for 1D arrays, but fallback to -1
1014
}
1115
}
@@ -43,29 +47,51 @@ fn first_true_1d_c(array: PyReadonlyArray1<bool>) -> isize {
4347

4448
// Process 8 elements at a time
4549
while i + 8 <= len {
46-
if slice[i] { return i as isize; }
47-
if slice[i+1] { return (i+1) as isize; }
48-
if slice[i+2] { return (i+2) as isize; }
49-
if slice[i+3] { return (i+3) as isize; }
50-
if slice[i+4] { return (i+4) as isize; }
51-
if slice[i+5] { return (i+5) as isize; }
52-
if slice[i+6] { return (i+6) as isize; }
53-
if slice[i+7] { return (i+7) as isize; }
50+
if slice[i] {
51+
return i as isize;
52+
}
53+
if slice[i + 1] {
54+
return (i + 1) as isize;
55+
}
56+
if slice[i + 2] {
57+
return (i + 2) as isize;
58+
}
59+
if slice[i + 3] {
60+
return (i + 3) as isize;
61+
}
62+
if slice[i + 4] {
63+
return (i + 4) as isize;
64+
}
65+
if slice[i + 5] {
66+
return (i + 5) as isize;
67+
}
68+
if slice[i + 6] {
69+
return (i + 6) as isize;
70+
}
71+
if slice[i + 7] {
72+
return (i + 7) as isize;
73+
}
5474
i += 8;
5575
}
5676

5777
// Handle remainder
5878
while i < len {
59-
if slice[i] { return i as isize; }
79+
if slice[i] {
80+
return i as isize;
81+
}
6082
i += 1;
6183
}
6284
-1
6385
} else {
64-
array.as_array().iter().position(|&v| v).map(|i| i as isize).unwrap_or(-1)
86+
array
87+
.as_array()
88+
.iter()
89+
.position(|&v| v)
90+
.map(|i| i as isize)
91+
.unwrap_or(-1)
6592
}
6693
}
6794

68-
6995
#[pyfunction]
7096
fn first_true_1d_d(array: PyReadonlyArray1<bool>) -> isize {
7197
if let Ok(slice) = array.as_slice() {
@@ -75,30 +101,52 @@ fn first_true_1d_d(array: PyReadonlyArray1<bool>) -> isize {
75101
unsafe {
76102
// Process 8 elements at a time
77103
while i + 8 <= len {
78-
if *slice.get_unchecked(i) { return i as isize; }
79-
if *slice.get_unchecked(i+1) { return (i+1) as isize; }
80-
if *slice.get_unchecked(i+2) { return (i+2) as isize; }
81-
if *slice.get_unchecked(i+3) { return (i+3) as isize; }
82-
if *slice.get_unchecked(i+4) { return (i+4) as isize; }
83-
if *slice.get_unchecked(i+5) { return (i+5) as isize; }
84-
if *slice.get_unchecked(i+6) { return (i+6) as isize; }
85-
if *slice.get_unchecked(i+7) { return (i+7) as isize; }
104+
if *slice.get_unchecked(i) {
105+
return i as isize;
106+
}
107+
if *slice.get_unchecked(i + 1) {
108+
return (i + 1) as isize;
109+
}
110+
if *slice.get_unchecked(i + 2) {
111+
return (i + 2) as isize;
112+
}
113+
if *slice.get_unchecked(i + 3) {
114+
return (i + 3) as isize;
115+
}
116+
if *slice.get_unchecked(i + 4) {
117+
return (i + 4) as isize;
118+
}
119+
if *slice.get_unchecked(i + 5) {
120+
return (i + 5) as isize;
121+
}
122+
if *slice.get_unchecked(i + 6) {
123+
return (i + 6) as isize;
124+
}
125+
if *slice.get_unchecked(i + 7) {
126+
return (i + 7) as isize;
127+
}
86128
i += 8;
87129
}
88130

89131
// Handle remainder
90132
while i < len {
91-
if *slice.get_unchecked(i) { return i as isize; }
133+
if *slice.get_unchecked(i) {
134+
return i as isize;
135+
}
92136
i += 1;
93137
}
94138
}
95139
-1
96140
} else {
97-
array.as_array().iter().position(|&v| v).map(|i| i as isize).unwrap_or(-1)
141+
array
142+
.as_array()
143+
.iter()
144+
.position(|&v| v)
145+
.map(|i| i as isize)
146+
.unwrap_or(-1)
98147
}
99148
}
100149

101-
102150
#[pyfunction]
103151
fn first_true_1d_e(array: PyReadonlyArray1<bool>) -> isize {
104152
if let Ok(slice) = array.as_slice() {
@@ -132,11 +180,15 @@ fn first_true_1d_e(array: PyReadonlyArray1<bool>) -> isize {
132180
}
133181
-1
134182
} else {
135-
array.as_array().iter().position(|&v| v).map(|i| i as isize).unwrap_or(-1)
183+
array
184+
.as_array()
185+
.iter()
186+
.position(|&v| v)
187+
.map(|i| i as isize)
188+
.unwrap_or(-1)
136189
}
137190
}
138191

139-
140192
#[pyfunction]
141193
fn first_true_1d_f(py: Python, array: PyReadonlyArray1<bool>) -> isize {
142194
if let Ok(slice) = array.as_slice() {
@@ -174,12 +226,15 @@ fn first_true_1d_f(py: Python, array: PyReadonlyArray1<bool>) -> isize {
174226
} else {
175227
let array_view = array.as_array();
176228
py.allow_threads(|| {
177-
array_view.iter().position(|&v| v).map(|i| i as isize).unwrap_or(-1)
229+
array_view
230+
.iter()
231+
.position(|&v| v)
232+
.map(|i| i as isize)
233+
.unwrap_or(-1)
178234
})
179235
}
180236
}
181237

182-
183238
#[pyfunction]
184239
#[pyo3(signature = (array, forward=true))]
185240
fn first_true_1d(py: Python, array: PyReadonlyArray1<bool>, forward: bool) -> isize {
@@ -222,9 +277,9 @@ fn first_true_1d(py: Python, array: PyReadonlyArray1<bool>, forward: bool) -> is
222277
let equal_one = chunk.cmp_eq(ones);
223278
if equal_one.any() {
224279
// Found a true in this chunk, search backwards within it
225-
for j in (0..32).rev() {
226-
if i + j < len && *ptr.add(i + j) != 0 {
227-
return (i + j) as isize;
280+
for j in (i..i + 32).rev() {
281+
if *ptr.add(j) != 0 {
282+
return j as isize;
228283
}
229284
}
230285
}
@@ -245,15 +300,22 @@ fn first_true_1d(py: Python, array: PyReadonlyArray1<bool>, forward: bool) -> is
245300
let array_view = array.as_array();
246301
py.allow_threads(|| {
247302
if forward {
248-
array_view.iter().position(|&v| v).map(|i| i as isize).unwrap_or(-1)
303+
array_view
304+
.iter()
305+
.position(|&v| v)
306+
.map(|i| i as isize)
307+
.unwrap_or(-1)
249308
} else {
250-
array_view.iter().rposition(|&v| v).map(|i| i as isize).unwrap_or(-1)
309+
array_view
310+
.iter()
311+
.rposition(|&v| v)
312+
.map(|i| i as isize)
313+
.unwrap_or(-1)
251314
}
252315
})
253316
}
254317
}
255318

256-
257319
#[pymodule]
258320
fn arrayredox(m: &Bound<'_, PyModule>) -> PyResult<()> {
259321
m.add_function(wrap_pyfunction!(first_true_1d_a, m)?)?;

0 commit comments

Comments
 (0)