Merge branch 'main' into nonzero_optmize

numpy · eendebakpt · Oct 6, 2024 · Oct 11, 2024 · Oct 11, 2024 · Oct 11, 2024
commit 5444f08274c3f2a2a0e56a28c78ede84105a34db
diff --git a/numpy/_core/src/multiarray/item_selection.c b/numpy/_core/src/multiarray/item_selection.c
@@ -2810,10 +2810,12 @@ nonzero_idxs_1D_bool(npy_intp count, npy_intp nonzero_count, char *data,
    * the fast bool count is followed by this sparse path is faster
    * than combining the two loops, even for larger arrays
    */
+   npy_intp * multi_index_end = multi_index + nonzero_count;
+
    if (((double)nonzero_count / count) <= 0.1) {
        npy_intp subsize;
        npy_intp j = 0;
-        while (1) {
+        while (multi_index < multi_index_end) {
            npy_memchr(data + j * stride, 0, stride, count - j,
                        &subsize, 1);
            j += subsize;
@@ -2832,7 +2834,7 @@ nonzero_idxs_1D_bool(npy_intp count, npy_intp nonzero_count, char *data,
        npy_intp j = 0;

        /* Manually unroll for GCC and maybe other compilers */
-        while (multi_index + 4 < multi_index_end) {
+        while (multi_index + 4 < multi_index_end && (j < count - 4) ) {
            *multi_index = j;
            multi_index += data[0] != 0;
            *multi_index = j + 1;
@@ -2845,7 +2847,7 @@ nonzero_idxs_1D_bool(npy_intp count, npy_intp nonzero_count, char *data,
            j += 4;
        }

-        while (multi_index < multi_index_end) {
+        while (multi_index < multi_index_end && (j < count) ) {
            *multi_index = j;
            multi_index += *data != 0;
            data += stride;
@@ -2956,6 +2958,7 @@ PyArray_Nonzero(PyArrayObject *self)
        for (npy_intp j = 0; j < count; ++j) {
            if (nonzero(data, self)) {
                if (++added_count > nonzero_count) {
+
                    break;
                }
                *multi_index++ = j;