ENH: iter: Change the meaning of itersize to be the number of elements, not the number of iterations

3 files changed, 43 insertions, 37 deletions

diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src
index 1215882ca..dd5bc9cc3 100644
--- a/numpy/core/src/multiarray/new_iterator.c.src
+++ b/numpy/core/src/multiarray/new_iterator.c.src
@@ -18,21 +18,23 @@
 /* Internal iterator flags */
 
 /* The perm is the identity */
-#define NPY_ITFLAG_IDENTPERM   0x001
+#define NPY_ITFLAG_IDENTPERM    0x001
 /* The perm has negative entries (indicating flipped axes) */
-#define NPY_ITFLAG_NEGPERM     0x002
+#define NPY_ITFLAG_NEGPERM      0x002
 /* The iterator is tracking an index */
-#define NPY_ITFLAG_HASINDEX    0x004
+#define NPY_ITFLAG_HASINDEX     0x004
 /* The iterator is tracking coordinates */
-#define NPY_ITFLAG_HASCOORDS   0x008
+#define NPY_ITFLAG_HASCOORDS    0x008
 /* The iteration order was forced on construction */
-#define NPY_ITFLAG_FORCEDORDER 0x010
+#define NPY_ITFLAG_FORCEDORDER  0x010
 /* The inner loop is handled outside the iterator */
-#define NPY_ITFLAG_NOINNER     0x020
+#define NPY_ITFLAG_NOINNER      0x020
 /* The iterator is buffered */
-#define NPY_ITFLAG_BUFFER      0x040
+#define NPY_ITFLAG_BUFFER       0x040
 /* The iterator should grow the buffered inner loop when possible */
-#define NPY_ITFLAG_GROWINNER   0x080
+#define NPY_ITFLAG_GROWINNER    0x080
+/* There is just one iteration, can specialize iternext for that */
+#define NPY_ITFLAG_ONEITERATION 0x100
 
 /* Internal iterator per-operand iterator flags */
 
@@ -536,10 +538,15 @@ NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
         ndim = NIT_NDIM(iter);
     }
 
-    /* Now that the axes are finished, adjust ITERSIZE if necessary */
+    /*
+     * Now that the axes are finished, check whether we can apply
+     * the single iteration optimization to the iternext function.
+     */
     if ((itflags&NPY_ITFLAG_NOINNER) && !(itflags&NPY_ITFLAG_BUFFER)) {
         NpyIter_AxisData *axisdata = NIT_AXISDATA(iter);
-        NIT_ITERSIZE(iter) /= NAD_SHAPE(axisdata);
+        if (NIT_ITERSIZE(iter) == NAD_SHAPE(axisdata)) {
+            NIT_ITFLAGS(iter) |= NPY_ITFLAG_ONEITERATION;
+        }
     }
 
     if (itflags&NPY_ITFLAG_BUFFER) {
@@ -549,11 +556,6 @@ NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
             return NULL;
         }
 
-        /* BUFFERED + NOINNER may not have a predictable itersize */
-        if (itflags&NPY_ITFLAG_NOINNER) {
-            NIT_ITERSIZE(iter) = 0;
-        }
-
         /* Prepare the next buffers and set pos/size */
         npyiter_copy_to_buffers(iter);
     }
@@ -654,7 +656,7 @@ int NpyIter_RemoveCoords(NpyIter *iter)
 /* Removes the inner loop handling (adds NPY_ITER_NO_INNER_ITERATION) */
 int NpyIter_RemoveInnerLoop(NpyIter *iter)
 {
-    npy_uint32 itflags = NIT_ITFLAGS(iter);;
+    npy_uint32 itflags = NIT_ITFLAGS(iter);
     npy_intp ndim = NIT_NDIM(iter);
     npy_intp niter = NIT_NITER(iter);
 
@@ -670,15 +672,15 @@ int NpyIter_RemoveInnerLoop(NpyIter *iter)
         itflags |= NPY_ITFLAG_NOINNER;
         NIT_ITFLAGS(iter) = itflags;
         
-        /* Adjust ITERSIZE */
-        if (itflags&NPY_ITFLAG_BUFFER) {
-            /* BUFFERED + NOINNER may not have a predictable itersize */
-            NIT_ITERSIZE(iter) = 0;
-                                            
-        }
-        else {
+        /*
+         * Check whether we can apply the single iteration
+         * optimization to the iternext function.
+         */
+        if (!(itflags&NPY_ITFLAG_BUFFER)) {
             NpyIter_AxisData *axisdata = NIT_AXISDATA(iter);
-            NIT_ITERSIZE(iter) /= NAD_SHAPE(axisdata);
+            if (NIT_ITERSIZE(iter) == NAD_SHAPE(axisdata)) {
+                NIT_ITFLAGS(iter) |= NPY_ITFLAG_ONEITERATION;
+            }
         }
     }
 
@@ -1162,10 +1164,10 @@ NpyIter_IterNext_Fn NpyIter_GetIterNext(NpyIter *iter)
     npy_intp niter = NIT_NITER(iter);
 
     /*
-     * When there is just one element being iterated,
-     * the iternext function is very simple
+     * When there is just one iteration and buffering is disabled
+     * the iternext function is very simple.
      */
-    if (NIT_ITERSIZE(iter) == 1) {
+    if (itflags&NPY_ITFLAG_ONEITERATION) {
         return &npyiter_iternext_sizeone;
     }
 
diff --git a/numpy/core/src/multiarray/new_iterator.h b/numpy/core/src/multiarray/new_iterator.h
index e89fb2569..70811e662 100644
--- a/numpy/core/src/multiarray/new_iterator.h
+++ b/numpy/core/src/multiarray/new_iterator.h
@@ -74,7 +74,7 @@ NpyIter_GetCoords_Fn NpyIter_GetGetCoords(NpyIter *iter);
 npy_intp NpyIter_GetNDim(NpyIter *iter);
 /* Gets the number of objects being iterated */
 npy_intp NpyIter_GetNIter(NpyIter *iter);
-/* Gets the number of times the iterator iterates */
+/* Gets the number of elements being iterated */
 npy_intp NpyIter_GetIterSize(NpyIter *iter);
 /* Gets the broadcast shape (if coords are enabled) */
 int NpyIter_GetShape(NpyIter *iter, npy_intp *outshape);
diff --git a/numpy/core/tests/test_new_iterator.py b/numpy/core/tests/test_new_iterator.py
index 6aec368dd..b77cd32b4 100644
--- a/numpy/core/tests/test_new_iterator.py
+++ b/numpy/core/tests/test_new_iterator.py
@@ -384,7 +384,8 @@ def test_iter_no_inner_full_coalesce():
     # Check no_inner iterators which coalesce into a single inner loop
 
     for shape in [(5,), (3,4), (2,3,4), (2,3,4,3), (2,3,2,2,3)]:
-        a = arange(np.prod(shape))
+        size = np.prod(shape)
+        a = arange(size)
         # Test each combination of forward and backwards indexing
         for dirs in range(2**len(shape)):
             dirs_index = [slice(None)]*len(shape)
@@ -397,17 +398,17 @@ def test_iter_no_inner_full_coalesce():
             # C-order
             i = newiter(aview, ['no_inner_iteration'], [['readonly']])
             assert_equal(i.ndim, 1)
-            assert_equal(i.itersize, 1)
+            assert_equal(i[0].shape, (size,))
             # Fortran-order
             i = newiter(aview.T, ['no_inner_iteration'], [['readonly']])
             assert_equal(i.ndim, 1)
-            assert_equal(i.itersize, 1)
+            assert_equal(i[0].shape, (size,))
             # Other order
             if len(shape) > 2:
                 i = newiter(aview.swapaxes(0,1),
                                     ['no_inner_iteration'], [['readonly']])
                 assert_equal(i.ndim, 1)
-                assert_equal(i.itersize, 1)
+                assert_equal(i[0].shape, (size,))
 
 def test_iter_no_inner_dim_coalescing():
     # Check no_inner iterators whose dimensions may not coalesce completely
@@ -417,21 +418,21 @@ def test_iter_no_inner_dim_coalescing():
     a = arange(24).reshape(2,3,4)[:,:,:-1]
     i = newiter(a, ['no_inner_iteration'], [['readonly']])
     assert_equal(i.ndim, 2)
-    assert_equal(i.itersize, 6)
+    assert_equal(i[0].shape, (3,))
     a = arange(24).reshape(2,3,4)[:,:-1,:]
     i = newiter(a, ['no_inner_iteration'], [['readonly']])
     assert_equal(i.ndim, 2)
-    assert_equal(i.itersize, 2)
+    assert_equal(i[0].shape, (8,))
     a = arange(24).reshape(2,3,4)[:-1,:,:]
     i = newiter(a, ['no_inner_iteration'], [['readonly']])
     assert_equal(i.ndim, 1)
-    assert_equal(i.itersize, 1)
+    assert_equal(i[0].shape, (12,))
     
     # Even with lots of 1-sized dimensions, should still coalesce
     a = arange(24).reshape(1,1,2,1,1,3,1,1,4,1,1)
     i = newiter(a, ['no_inner_iteration'], [['readonly']])
     assert_equal(i.ndim, 1)
-    assert_equal(i.itersize, 1)
+    assert_equal(i[0].shape, (24,))
 
 def test_iter_dim_coalescing():
     # Check that the correct number of dimensions are coalesced
@@ -1110,9 +1111,12 @@ def test_iter_remove_coords_inner_loop():
     assert_equal(i.itviews[0].shape, (24,))
 
     # Removing the inner loop means there's just one iteration
+    i.reset()
     assert_equal(i.itersize, 24)
+    assert_equal(i[0].shape, tuple())
     i.remove_inner_loop()
-    assert_equal(i.itersize, 1)
+    assert_equal(i.itersize, 24)
+    assert_equal(i[0].shape, (24,))
     assert_equal(i.value, arange(24))
 
 def test_iter_buffering():