ENH: nditer: Add the basic checking and input of the MASK-related flags

5 files changed, 105 insertions, 46 deletions

diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
index 4e8b2a401..cc57dbf60 100644
--- a/numpy/core/include/numpy/ndarraytypes.h
+++ b/numpy/core/include/numpy/ndarraytypes.h
@@ -889,14 +889,14 @@ typedef void (NpyIter_GetMultiIndexFunc)(NpyIter *iter,
 #define NPY_ITER_ALLOCATE                   0x01000000
 /* If an operand is allocated, don't use any subtype */
 #define NPY_ITER_NO_SUBTYPE                 0x02000000
+/* This is a virtual array slot, operand is NULL but temporary data is there */
+#define NPY_ITER_VIRTUAL                    0x04000000
 /* Require that the dimension match the iterator dimensions exactly */
 #define NPY_ITER_NO_BROADCAST               0x08000000
 /* A mask is being used on this array, affects buffer -> array copy */
-#define NPY_ITER_WRITEMASKED                0x04000000
+#define NPY_ITER_WRITEMASKED                0x10000000
 /* This array is the mask for all WRITEMASKED operands */
-#define NPY_ITER_ARRAYMASK                  0x10000000
-/* This is a virtual array slot of the mask for all WRITEMASKED operands */
-#define NPY_ITER_VIRTUALMASK                0x20000000
+#define NPY_ITER_ARRAYMASK                  0x20000000
 
 #define NPY_ITER_GLOBAL_FLAGS               0x0000ffff
 #define NPY_ITER_PER_OP_FLAGS               0xffff0000
diff --git a/numpy/core/src/multiarray/nditer_api.c b/numpy/core/src/multiarray/nditer_api.c
index 30990fff1..1328ebc38 100644
--- a/numpy/core/src/multiarray/nditer_api.c
+++ b/numpy/core/src/multiarray/nditer_api.c
@@ -1337,6 +1337,9 @@ NpyIter_DebugPrint(NpyIter *iter)
     printf("\n");
     printf("| NDim: %d\n", (int)ndim);
     printf("| NOp: %d\n", (int)nop);
+    if (NIT_MASKOP(iter) >= 0) {
+        printf("| MaskOp: %d\n", (int)NIT_MASKOP(iter));
+    }
     printf("| IterSize: %d\n", (int)NIT_ITERSIZE(iter));
     printf("| IterStart: %d\n", (int)NIT_ITERSTART(iter));
     printf("| IterEnd: %d\n", (int)NIT_ITEREND(iter));
@@ -1418,6 +1421,8 @@ NpyIter_DebugPrint(NpyIter *iter)
             printf("ALIGNED ");
         if ((NIT_OPITFLAGS(iter)[iop])&NPY_OP_ITFLAG_REDUCE)
             printf("REDUCE ");
+        if ((NIT_OPITFLAGS(iter)[iop])&NPY_OP_ITFLAG_VIRTUAL)
+            printf("VIRTUAL ");
         if ((NIT_OPITFLAGS(iter)[iop])&NPY_OP_ITFLAG_WRITEMASKED)
             printf("WRITEMASKED ");
         printf("\n");
diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c
index 7120d44d2..fc13aa0f5 100644
--- a/numpy/core/src/multiarray/nditer_constr.c
+++ b/numpy/core/src/multiarray/nditer_constr.c
@@ -39,7 +39,8 @@ npyiter_prepare_operands(int nop, PyArrayObject **op_in,
                     PyArray_Descr **op_request_dtypes,
                     PyArray_Descr **op_dtype,
                     npy_uint32 flags,
-                    npy_uint32 *op_flags, char *op_itflags);
+                    npy_uint32 *op_flags, char *op_itflags,
+                    npy_int8 *out_maskop);
 static int
 npyiter_check_casting(int nop, PyArrayObject **op,
                     PyArray_Descr **op_dtype,
@@ -200,7 +201,7 @@ NpyIter_AdvancedNew(int nop, PyArrayObject **op_in, npy_uint32 flags,
     if (!npyiter_prepare_operands(nop, op_in, op, op_dataptr,
                         op_request_dtypes, op_dtype,
                         flags,
-                        op_flags, op_itflags)) {
+                        op_flags, op_itflags, &NIT_MASKOP(iter))) {
         PyArray_free(iter);
         return NULL;
     }
@@ -907,21 +908,23 @@ npyiter_check_per_op_flags(npy_uint32 op_flags, char *op_itflags)
                 "be used with READWRITE or WRITEONLY");
             return 0;
         }
-        if ((op_flags&(NPY_ITER_ARRAYMASK|NPY_ITER_VIRTUALMASK)) != 0) {
+        if ((op_flags&NPY_ITER_ARRAYMASK) != 0) {
             PyErr_SetString(PyExc_ValueError,
                 "The iterator flag WRITEMASKED may not "
-                "be used together with ARRAYMASK or VIRTUALMASK");
+                "be used together with ARRAYMASK");
             return 0;
         }
         *op_itflags |= NPY_OP_ITFLAG_WRITEMASKED;
     }
 
-    if ((op_flags&(NPY_ITER_ARRAYMASK|NPY_ITER_VIRTUALMASK)) ==
-                        (NPY_ITER_ARRAYMASK|NPY_ITER_VIRTUALMASK)) {
-        PyErr_SetString(PyExc_ValueError,
-            "The iterator flag ARRAYMASK may not "
-            "be used together with VIRTUALMASK");
-        return 0;
+    if ((op_flags&NPY_ITER_VIRTUAL) != 0) {
+        if ((op_flags&NPY_ITER_READWRITE) == 0) {
+            PyErr_SetString(PyExc_ValueError,
+                "The iterator flag VIRTUAL should be "
+                "be used together with READWRITE");
+            return 0;
+        }
+        *op_itflags |= NPY_OP_ITFLAG_VIRTUAL;
     }
 
     return 1;
@@ -944,43 +947,69 @@ npyiter_prepare_one_operand(PyArrayObject **op,
 {
     /* NULL operands must be automatically allocated outputs */
     if (*op == NULL) {
-        /* ALLOCATE should be enabled */
-        if (!(op_flags&NPY_ITER_ALLOCATE)) {
+        /* ALLOCATE or VIRTUAL should be enabled */
+        if ((op_flags&(NPY_ITER_ALLOCATE|NPY_ITER_VIRTUAL)) == 0) {
             PyErr_SetString(PyExc_ValueError,
-                    "Iterator operand was NULL, but automatic allocation as an "
-                    "output wasn't requested");
+                    "Iterator operand was NULL, but neither the "
+                    "ALLOCATE nor the VIRTUAL flag was specified");
             return 0;
         }
-        /* Writing should be enabled */
-        if (!((*op_itflags)&NPY_OP_ITFLAG_WRITE)) {
-            PyErr_SetString(PyExc_ValueError,
-                    "Automatic allocation was requested for an iterator "
-                    "operand, but it wasn't flagged for writing");
-            return 0;
+
+        if (op_flags&NPY_ITER_ALLOCATE) {
+            /* Writing should be enabled */
+            if (!((*op_itflags)&NPY_OP_ITFLAG_WRITE)) {
+                PyErr_SetString(PyExc_ValueError,
+                        "Automatic allocation was requested for an iterator "
+                        "operand, but it wasn't flagged for writing");
+                return 0;
+            }
+            /*
+             * Reading should be disabled if buffering is enabled without
+             * also enabling NPY_ITER_DELAY_BUFALLOC.  In all other cases,
+             * the caller may initialize the allocated operand to a value
+             * before beginning iteration.
+             */
+            if (((flags&(NPY_ITER_BUFFERED|
+                            NPY_ITER_DELAY_BUFALLOC)) == NPY_ITER_BUFFERED) &&
+                    ((*op_itflags)&NPY_OP_ITFLAG_READ)) {
+                PyErr_SetString(PyExc_ValueError,
+                        "Automatic allocation was requested for an iterator "
+                        "operand, and it was flagged as readable, but "
+                        "buffering  without delayed allocation was enabled");
+                return 0;
+            }
+
+            /* If a requested dtype was provided, use it, otherwise NULL */
+            Py_XINCREF(op_request_dtype);
+            *op_dtype = op_request_dtype;
         }
-        /*
-         * Reading should be disabled if buffering is enabled without
-         * also enabling NPY_ITER_DELAY_BUFALLOC.  In all other cases,
-         * the caller may initialize the allocated operand to a value
-         * before beginning iteration.
-         */
-        if (((flags&(NPY_ITER_BUFFERED|
-                        NPY_ITER_DELAY_BUFALLOC)) == NPY_ITER_BUFFERED) &&
-                ((*op_itflags)&NPY_OP_ITFLAG_READ)) {
-            PyErr_SetString(PyExc_ValueError,
-                    "Automatic allocation was requested for an iterator "
-                    "operand, and it was flagged as readable, but buffering "
-                    " without delayed allocation was enabled");
-            return 0;
+        else {
+            *op_dtype = NULL;
         }
+
+        /* Specify uint8 if no dtype was requested for the mask */
+        if (op_flags&NPY_ITER_ARRAYMASK) {
+            if (*op_dtype == NULL) {
+                *op_dtype = PyArray_DescrFromType(NPY_UINT8);
+                if (*op_dtype == NULL) {
+                    return 0;
+                }
+            }
+        }
+
         *op_dataptr = NULL;
-        /* If a requested dtype was provided, use it, otherwise NULL */
-        Py_XINCREF(op_request_dtype);
-        *op_dtype = op_request_dtype;
 
         return 1;
     }
 
+    /* VIRTUAL operands must be NULL */
+    if (op_flags&NPY_ITER_VIRTUAL) {
+        PyErr_SetString(PyExc_ValueError,
+                "Iterator operand flag VIRTUAL was specified, "
+                "but the operand was not NULL");
+        return 0;
+    }
+
     if (PyArray_Check(*op)) {
         if (((*op_itflags)&NPY_OP_ITFLAG_WRITE) &&
                     (!PyArray_CHKFLAGS(*op, NPY_ARRAY_WRITEABLE))) {
@@ -1091,9 +1120,11 @@ npyiter_prepare_operands(int nop, PyArrayObject **op_in,
                     PyArray_Descr **op_request_dtypes,
                     PyArray_Descr **op_dtype,
                     npy_uint32 flags,
-                    npy_uint32 *op_flags, char *op_itflags)
+                    npy_uint32 *op_flags, char *op_itflags,
+                    npy_int8 *out_maskop)
 {
     int iop, i;
+    npy_int8 maskop = -1;
 
     for (iop = 0; iop < nop; ++iop) {
         op[iop] = op_in[iop];
@@ -1109,6 +1140,23 @@ npyiter_prepare_operands(int nop, PyArrayObject **op_in,
             return 0;
         }
 
+        /* Extract the operand which is for masked iteration */
+        if ((op_flags[iop]&NPY_ITER_ARRAYMASK) != 0) {
+            if (maskop != -1) {
+                PyErr_SetString(PyExc_ValueError,
+                        "Only one iterator operand may receive an "
+                        "ARRAYMASK flag");
+                for (i = 0; i <= iop; ++i) {
+                    Py_XDECREF(op[i]);
+                    Py_XDECREF(op_dtype[i]);
+                }
+                return 0;
+            }
+
+            maskop = iop;
+            *out_maskop = iop;
+        }
+
         /*
          * Prepare the operand.  This produces an op_dtype[iop] reference
          * on success.
diff --git a/numpy/core/src/multiarray/nditer_impl.h b/numpy/core/src/multiarray/nditer_impl.h
index 3494d3937..e5ec487f8 100644
--- a/numpy/core/src/multiarray/nditer_impl.h
+++ b/numpy/core/src/multiarray/nditer_impl.h
@@ -116,8 +116,10 @@
 #define NPY_OP_ITFLAG_ALIGNED      0x10
 /* The operand is being reduced */
 #define NPY_OP_ITFLAG_REDUCE       0x20
+/* The operand is for temporary use, does not have a backing array */
+#define NPY_OP_ITFLAG_VIRTUAL      0x40
 /* The operand requires masking when copying buffer -> array */
-#define NPY_OP_ITFLAG_WRITEMASKED  0x40
+#define NPY_OP_ITFLAG_WRITEMASKED  0x80
 
 /*
  * The data layout of the iterator is fully specified by
@@ -131,7 +133,9 @@
 struct NpyIter_InternalOnly {
     /* Initial fixed position data */
     npy_uint32 itflags;
-    npy_uint16 ndim, nop;
+    npy_uint8 ndim, nop;
+    npy_int8 maskop;
+    npy_uint8 unused_padding;
     npy_intp itersize, iterstart, iterend;
     /* iterindex is only used if RANGED or BUFFERED is set */
     npy_intp iterindex;
@@ -190,6 +194,8 @@ typedef struct NpyIter_BD NpyIter_BufferData;
         ((iter)->ndim)
 #define NIT_NOP(iter) \
         ((iter)->nop)
+#define NIT_MASKOP(iter) \
+        ((iter)->maskop)
 #define NIT_ITERSIZE(iter) \
         (iter->itersize)
 #define NIT_ITERSTART(iter) \
diff --git a/numpy/core/src/multiarray/nditer_pywrap.c b/numpy/core/src/multiarray/nditer_pywrap.c
index 7f902d6fa..1e86487f5 100644
--- a/numpy/core/src/multiarray/nditer_pywrap.c
+++ b/numpy/core/src/multiarray/nditer_pywrap.c
@@ -433,8 +433,8 @@ NpyIter_OpFlagsConverter(PyObject *op_flags_in,
                 }
                 break;
             case 'v':
-                if (strcmp(str, "virtualmask") == 0) {
-                    flag = NPY_ITER_VIRTUALMASK;
+                if (strcmp(str, "virtual") == 0) {
+                    flag = NPY_ITER_VIRTUAL;
                 }
                 break;
             case 'w':