ENH: iter: Split flag NBO_ALIGN into NBO and ALIGN, add CONTIG flag

4 files changed, 84 insertions, 37 deletions

diff --git a/numpy/core/include/numpy/ndarraytypes.h b/numpy/core/include/numpy/ndarraytypes.h
index 84a41d558..46dbbefc7 100644
--- a/numpy/core/include/numpy/ndarraytypes.h
+++ b/numpy/core/include/numpy/ndarraytypes.h
@@ -899,20 +899,24 @@ typedef void (*NpyIter_GetCoords_Fn )(NpyIter *iter,
 #define NPY_ITER_READONLY                   0x00020000
 /* The operand will only be written to */
 #define NPY_ITER_WRITEONLY                  0x00040000
-/* The operand's data must be in native byte order and aligned */
-#define NPY_ITER_NBO_ALIGNED                0x00080000
+/* The operand's data must be in native byte order */
+#define NPY_ITER_NBO                        0x00080000
+/* The operand's data must be aligned */
+#define NPY_ITER_ALIGNED                    0x00100000
+/* The operand's data must be contiguous (within the inner loop) */
+#define NPY_ITER_CONTIG                     0x00200000
 /* The operand may be copied to satisfy requirements */
-#define NPY_ITER_COPY                       0x00100000
+#define NPY_ITER_COPY                       0x00400000
 /* The operand may be copied with UPDATEIFCOPY to satisfy requirements */
-#define NPY_ITER_UPDATEIFCOPY               0x00200000
+#define NPY_ITER_UPDATEIFCOPY               0x00800000
 /* Allow writeable operands to have references or pointers */
-#define NPY_ITER_WRITEABLE_REFERENCES       0x00400000
+#define NPY_ITER_WRITEABLE_REFERENCES       0x01000000
 /* Allocate the operand if it is NULL */
-#define NPY_ITER_ALLOCATE                   0x00800000
+#define NPY_ITER_ALLOCATE                   0x02000000
 /* If an operand is allocated, don't use any subtype */
-#define NPY_ITER_NO_SUBTYPE                 0x01000000
+#define NPY_ITER_NO_SUBTYPE                 0x04000000
 /* Require that the dimension match the iterator dimensions exactly */
-#define NPY_ITER_NO_BROADCAST               0x02000000
+#define NPY_ITER_NO_BROADCAST               0x08000000
 
 #define NPY_ITER_GLOBAL_FLAGS               0x0000ffff
 #define NPY_ITER_PER_OP_FLAGS               0xffff0000
diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src
index 154706ec9..cc6a81be0 100644
--- a/numpy/core/src/multiarray/new_iterator.c.src
+++ b/numpy/core/src/multiarray/new_iterator.c.src
@@ -2311,8 +2311,8 @@ npyiter_prepare_one_operand(PyArrayObject **op,
             *op_dtype = op_request_dtype;
         }
 
-        /* Check if the operand is aligned and in the byte order requested */
-        if (op_flags&NPY_ITER_NBO_ALIGNED) {
+        /* Check if the operand is in the byte order requested */
+        if (op_flags&NPY_ITER_NBO) {
             /* Check byte order */
             if (!PyArray_ISNBO((*op_dtype)->byteorder)) {
                 PyArray_Descr *nbo_dtype;
@@ -2325,11 +2325,18 @@ npyiter_prepare_one_operand(PyArrayObject **op,
                 /* Indicate that byte order or alignment needs fixing */
                 *op_itflags |= NPY_OP_ITFLAG_CAST;
             }
+        }
+        /* Check if the operand is aligned */
+        if (op_flags&NPY_ITER_ALIGNED) {
             /* Check alignment */
-            else if (!PyArray_ISALIGNED(*op)) {
+            if (!PyArray_ISALIGNED(*op)) {
                 *op_itflags |= NPY_OP_ITFLAG_CAST;
             }
         }
+        /*
+         * The check for NPY_ITER_CONTIG can only be done later,
+         * once the final iteration order is settled.
+         */
     }
     else {
         PyErr_SetString(PyExc_ValueError,
@@ -3594,7 +3601,7 @@ npyiter_allocate_arrays(NpyIter *iter,
             npyiter_replace_axisdata(iter, iiter, op[iiter], ondim,
                     PyArray_DATA(op[iiter]), op_axes ? op_axes[iiter] : NULL);
 
-            /* New arrays are aligned and need no swapping or casting */
+            /* New arrays are aligned and need no cast */
             op_itflags[iiter] |= NPY_OP_ITFLAG_ALIGNED;
             op_itflags[iiter] &= ~NPY_OP_ITFLAG_CAST;
         }
@@ -3638,7 +3645,7 @@ npyiter_allocate_arrays(NpyIter *iter,
             npyiter_replace_axisdata(iter, iiter, op[iiter], ondim,
                     PyArray_DATA(op[iiter]), op_axes ? op_axes[iiter] : NULL);
 
-            /* The temporary copy is aligned and needs no swap or cast */
+            /* The temporary copy is aligned and needs no cast */
             op_itflags[iiter] |= NPY_OP_ITFLAG_ALIGNED;
             op_itflags[iiter] &= ~NPY_OP_ITFLAG_CAST;
         }
@@ -3664,6 +3671,23 @@ npyiter_allocate_arrays(NpyIter *iter,
             }
         }
 
+        /* Here we can finally check for contiguous iteration */
+        if (op_flags[iiter]&NPY_ITER_CONTIG) {
+            NpyIter_AxisData *axisdata = NIT_AXISDATA(iter);
+            npy_intp stride = NAD_STRIDES(axisdata)[iiter];
+
+            if (stride != op_dtype[iiter]->elsize) {
+                op_itflags[iiter] |= NPY_OP_ITFLAG_CAST;
+                if (!(itflags&NPY_ITFLAG_BUFFER)) {
+                    PyErr_SetString(PyExc_TypeError,
+                            "Iterator operand required buffering, "
+                            "to be contiguous as requested, but "
+                            "buffering is not enabled");
+                    return 0;
+                }
+            }
+        }
+
         /*
          * If no alignment, byte swap, or casting is needed, and
          * the inner stride of this operand works for the whole
diff --git a/numpy/core/src/multiarray/new_iterator_pywrap.c b/numpy/core/src/multiarray/new_iterator_pywrap.c
index 68a741f9f..e448aac29 100644
--- a/numpy/core/src/multiarray/new_iterator_pywrap.c
+++ b/numpy/core/src/multiarray/new_iterator_pywrap.c
@@ -302,17 +302,23 @@ NpyIter_OpFlagsConverter(PyObject *op_flags_in,
                 if (strcmp(str, "allocate") == 0) {
                     flag = NPY_ITER_ALLOCATE;
                 }
+                if (strcmp(str, "aligned") == 0) {
+                    flag = NPY_ITER_ALIGNED;
+                }
                 break;
             case 'c':
                 if (strcmp(str, "copy") == 0) {
                     flag = NPY_ITER_COPY;
                 }
+                if (strcmp(str, "contig") == 0) {
+                    flag = NPY_ITER_CONTIG;
+                }
                 break;
             case 'n':
                 switch (str[1]) {
                     case 'b':
-                        if (strcmp(str, "nbo_aligned") == 0) {
-                            flag = NPY_ITER_NBO_ALIGNED;
+                        if (strcmp(str, "nbo") == 0) {
+                            flag = NPY_ITER_NBO;
                         }
                         break;
                     case 'o':
@@ -396,14 +402,14 @@ npyiter_convert_op_flags_array(PyObject *op_flags_in,
         if (f == NULL) {
             return 0;
         }
+        /* If the first item is a string, try as one set of flags */
+        if (iiter == 0 && (PyString_Check(f) || PyUnicode_Check(f))) {
+            Py_DECREF(f);
+            goto try_single_flags;
+        }
         if (NpyIter_OpFlagsConverter(f,
                         &op_flags_array[iiter]) != 1) {
             Py_DECREF(f);
-            /* If the first one doesn't work, try the whole thing as flags */
-            if (iiter == 0) {
-                PyErr_Clear();
-                goto try_single_flags;
-            }
             return 0;
         }
         
@@ -938,7 +944,7 @@ NpyIter_NestedIters(PyObject *NPY_UNUSED(self),
                 !(op_flags[iiter]&(NPY_ITER_COPY|
                                    NPY_ITER_UPDATEIFCOPY|
                                    NPY_ITER_ALLOCATE))) {
-            op_flags[iiter] &= ~NPY_ITER_NBO_ALIGNED;
+            op_flags[iiter] &= ~(NPY_ITER_NBO|NPY_ITER_ALIGNED|NPY_ITER_CONTIG);
             op_request_dtypes_inner[iiter] = op_request_dtypes[iiter];
             op_request_dtypes[iiter] = NULL;
         }
diff --git a/numpy/core/tests/test_new_iterator.py b/numpy/core/tests/test_new_iterator.py
index dc34b655a..06aa34560 100644
--- a/numpy/core/tests/test_new_iterator.py
+++ b/numpy/core/tests/test_new_iterator.py
@@ -668,8 +668,8 @@ def test_iter_flags_errors():
     # Can't iterate if size is zero
     assert_raises(ValueError, newiter, np.array([]))
 
-def test_iter_nbo_align():
-    # Check that byte order and alignment changes work
+def test_iter_nbo_align_contig():
+    # Check that byte order, alignment, and contig changes work
 
     # Byte order change by requesting a specific dtype
     a = np.arange(6, dtype='f4')
@@ -685,11 +685,11 @@ def test_iter_nbo_align():
     i = None
     assert_equal(au, [2]*6)
 
-    # Byte order change by requesting NBO_ALIGNED
+    # Byte order change by requesting NBO
     a = np.arange(6, dtype='f4')
     au = a.byteswap().newbyteorder()
     assert_(a.dtype.byteorder != au.dtype.byteorder)
-    i = newiter(au, [], [['readwrite','updateifcopy','nbo_aligned']], casting='equiv')
+    i = newiter(au, [], [['readwrite','updateifcopy','nbo']], casting='equiv')
     assert_equal(i.dtypes[0].byteorder, a.dtype.byteorder)
     assert_equal(i.operands[0].dtype.byteorder, a.dtype.byteorder)
     assert_equal(i.operands[0], a)
@@ -702,18 +702,31 @@ def test_iter_nbo_align():
     a.dtype = 'f4'
     a[:] = np.arange(6, dtype='f4')
     assert_(not a.flags.aligned)
-    # Without 'nbo_aligned', shouldn't copy
+    # Without 'aligned', shouldn't copy
     i = newiter(a, [], [['readonly']])
     assert_(not i.operands[0].flags.aligned)
     assert_equal(i.operands[0], a);
-    # With 'nbo_aligned', should make a copy
-    i = newiter(a, [], [['readwrite','updateifcopy','nbo_aligned']])
+    # With 'aligned', should make a copy
+    i = newiter(a, [], [['readwrite','updateifcopy','aligned']])
     assert_(i.operands[0].flags.aligned)
     assert_equal(i.operands[0], a);
     i.operands[0][:] = 3
     i = None
     assert_equal(a, [3]*6)
 
+    # Discontiguous input
+    a = arange(12)
+    # If it is contiguous, shouldn't copy
+    i = newiter(a[:6], [], [['readonly']])
+    assert_(i.operands[0].flags.contiguous)
+    assert_equal(i.operands[0], a[:6]);
+    # If it isn't contiguous, should buffer
+    i = newiter(a[::2], ['buffered','no_inner_iteration'],
+                        [['readonly','contig']],
+                        buffersize=10)
+    assert_(i[0].flags.contiguous)
+    assert_equal(i[0], a[::2])
+
 def test_iter_array_cast():
     # Check that arrays are cast as requested
 
@@ -1314,7 +1327,7 @@ def test_iter_buffering():
         for buffersize in (1,2,3,5,8,11,16,1024):
             vals = []
             i = newiter(a, ['buffered','no_inner_iteration'],
-                           [['readonly','nbo_aligned']],
+                           [['readonly','nbo','aligned']],
                            order='C',
                            casting='equiv',
                            buffersize=buffersize)
@@ -1330,7 +1343,7 @@ def test_iter_write_buffering():
     # F-order swapped array
     a = np.arange(24).reshape(2,3,4).T.newbyteorder().byteswap()
     i = newiter(a, ['buffered'],
-                   [['readwrite','nbo_aligned']],
+                   [['readwrite','nbo','aligned']],
                    casting='equiv',
                    order='C',
                    buffersize=16)
@@ -1370,7 +1383,7 @@ def test_iter_buffered_cast_simple():
 
     a = np.arange(10, dtype='f4')
     i = newiter(a, ['buffered','no_inner_iteration'],
-                   [['readwrite','nbo_aligned']],
+                   [['readwrite','nbo','aligned']],
                    casting='same_kind',
                    op_dtypes=[np.dtype('f8')],
                    buffersize=3)
@@ -1384,7 +1397,7 @@ def test_iter_buffered_cast_byteswapped():
 
     a = np.arange(10, dtype='f4').newbyteorder().byteswap()
     i = newiter(a, ['buffered','no_inner_iteration'],
-                   [['readwrite','nbo_aligned','same_kind_casts']],
+                   [['readwrite','nbo','aligned','same_kind_casts']],
                    op_dtypes=[np.dtype('f8').newbyteorder()],
                    buffersize=3)
     for v in i:
@@ -1394,7 +1407,7 @@ def test_iter_buffered_cast_byteswapped():
 
     a = np.arange(10, dtype='f8').newbyteorder().byteswap()
     i = newiter(a, ['buffered','no_inner_iteration'],
-                   [['readwrite','nbo_aligned','unsafe_casts']],
+                   [['readwrite','nbo','aligned','unsafe_casts']],
                    op_dtypes=[np.dtype('c8').newbyteorder()],
                    buffersize=3)
     for v in i:
@@ -1408,7 +1421,7 @@ def test_iter_buffered_cast_byteswapped():
     a = np.arange(10, dtype='c8').newbyteorder().byteswap()
     a += 2j
     i = newiter(a, ['buffered','no_inner_iteration'],
-                   [['readwrite','nbo_aligned']],
+                   [['readwrite','nbo','aligned']],
                    casting='same_kind',
                    op_dtypes=[np.dtype('c16')],
                    buffersize=3)
@@ -1419,7 +1432,7 @@ def test_iter_buffered_cast_byteswapped():
     a = np.arange(10, dtype='c8')
     a += 2j
     i = newiter(a, ['buffered','no_inner_iteration'],
-                   [['readwrite','nbo_aligned']],
+                   [['readwrite','nbo','aligned']],
                    casting='same_kind',
                    op_dtypes=[np.dtype('c16').newbyteorder()],
                    buffersize=3)
@@ -1430,7 +1443,7 @@ def test_iter_buffered_cast_byteswapped():
     a = np.arange(10, dtype=np.clongdouble).newbyteorder().byteswap()
     a += 2j
     i = newiter(a, ['buffered','no_inner_iteration'],
-                   [['readwrite','nbo_aligned']],
+                   [['readwrite','nbo','aligned']],
                    casting='same_kind',
                    op_dtypes=[np.dtype('c16')],
                    buffersize=3)
@@ -1440,7 +1453,7 @@ def test_iter_buffered_cast_byteswapped():
 
     a = np.arange(10, dtype=np.longdouble).newbyteorder().byteswap()
     i = newiter(a, ['buffered','no_inner_iteration'],
-                   [['readwrite','nbo_aligned']],
+                   [['readwrite','nbo','aligned']],
                    casting='same_kind',
                    op_dtypes=[np.dtype('f4')],
                    buffersize=7)