ENH: iter: Add getting and setting of IterIndex, the iteration index

4 files changed, 258 insertions, 34 deletions

diff --git a/numpy/core/src/multiarray/new_iterator.c.src b/numpy/core/src/multiarray/new_iterator.c.src
index 74163e628..480e5bd1a 100644
--- a/numpy/core/src/multiarray/new_iterator.c.src
+++ b/numpy/core/src/multiarray/new_iterator.c.src
@@ -342,7 +342,7 @@ NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
     /* Calculate how many dimensions the iterator should have */
     ndim = npyiter_calculate_ndim(niter, op_in, oa_ndim);
 
-    /* If 'ndim' is zero, any outputs should scalars */
+    /* If 'ndim' is zero, any outputs should be scalars */
     if (ndim == 0) {
         output_scalars = 1;
         ndim = 1;
@@ -805,6 +805,13 @@ int NpyIter_GotoCoords(NpyIter *iter, npy_intp *coords)
         return NPY_FAIL;
     }
 
+    if (itflags&NPY_ITFLAG_NOINNER) {
+        PyErr_SetString(PyExc_ValueError,
+                "Cannot call GotoCoords on an iterator which "
+                "has the flag NO_INNER_ITERATION");
+        return NPY_FAIL;
+    }
+
     perm = NIT_PERM(iter);
     dataptr = NIT_RESETDATAPTR(iter);
     axisdata = NIT_AXISDATA(iter);
@@ -861,7 +868,8 @@ int NpyIter_GotoCoords(NpyIter *iter, npy_intp *coords)
     return NPY_SUCCEED;
 }
 
-/* If the iterator is tracking an index, sets the iterator
+/*
+ * If the iterator is tracking an index, sets the iterator
  * to the specified index.
  *
  * Returns NPY_SUCCEED on success, NPY_FAIL on failure.
@@ -891,6 +899,13 @@ int NpyIter_GotoIndex(NpyIter *iter, npy_intp index)
         return NPY_FAIL;
     }
 
+    if (itflags&NPY_ITFLAG_NOINNER) {
+        PyErr_SetString(PyExc_ValueError,
+                "Cannot call GotoIndex on an iterator which "
+                "has the flag NO_INNER_ITERATION");
+        return NPY_FAIL;
+    }
+
     if (index < 0 || index >= NIT_ITERSIZE(iter)) {
         PyErr_SetString(PyExc_IndexError,
                 "Iterator GotoIndex called with out-of-bounds "
@@ -941,6 +956,114 @@ int NpyIter_GotoIndex(NpyIter *iter, npy_intp index)
     return NPY_SUCCEED;
 }
 
+/*
+ * Sets the iterator position to the specified iterindex,
+ * which matches the iteration order of the iterator.
+ *
+ * Returns NPY_SUCCEED on success, NPY_FAIL on failure.
+ */
+int NpyIter_GotoIterIndex(NpyIter *iter, npy_intp iterindex)
+{
+    npy_uint32 itflags = NIT_ITFLAGS(iter);
+    npy_intp idim, ndim = NIT_NDIM(iter);
+    npy_intp niter = NIT_NITER(iter);
+
+    char **dataptr;
+    NpyIter_AxisData *axisdata;
+    npy_intp sizeof_axisdata;
+    npy_intp istrides, nstrides, i, shape;
+
+    if (itflags&NPY_ITFLAG_BUFFER) {
+        PyErr_SetString(PyExc_ValueError,
+                "Cannot call GotoIterIndex on an iterator which "
+                "is buffered");
+        return NPY_FAIL;
+    }
+
+    if (itflags&NPY_ITFLAG_NOINNER) {
+        PyErr_SetString(PyExc_ValueError,
+                "Cannot call GotoIterIndex on an iterator which "
+                "has the flag NO_INNER_ITERATION");
+        return NPY_FAIL;
+    }
+
+    if (iterindex < 0 || iterindex >= NIT_ITERSIZE(iter)) {
+        PyErr_SetString(PyExc_IndexError,
+                "Iterator GotoIterIndex called with out-of-bounds "
+                "index.");
+        return NPY_FAIL;
+    }
+
+    axisdata = NIT_AXISDATA(iter);
+    sizeof_axisdata = NIT_SIZEOF_AXISDATA(itflags, ndim, niter);
+    nstrides = NAD_NSTRIDES();
+
+    /*
+     * Set the coordinates, from the fastest-changing to the
+     * slowest-changing.
+     */
+    axisdata = NIT_AXISDATA(iter);
+    shape = NAD_SHAPE(axisdata);
+    i = iterindex;
+    iterindex /= shape;
+    NAD_COORD(axisdata) = i - iterindex * shape;
+    for (idim = 0; idim < ndim-1; ++idim) {
+        NIT_ADVANCE_AXISDATA(axisdata, 1);
+        shape = NAD_SHAPE(axisdata);
+        i = iterindex;
+        iterindex /= shape;
+        NAD_COORD(axisdata) = i - iterindex * shape;
+    }
+
+    dataptr = NIT_RESETDATAPTR(iter);
+
+    /*
+     * The successive pointers accumulate
+     * the offsets in the opposite order, starting from the
+     * original data pointers.
+     */
+    for (idim = 0; idim < ndim; ++idim, NIT_ADVANCE_AXISDATA(axisdata, -1)) {
+        npy_intp *strides;
+        char **ptrs;
+
+        strides = NAD_STRIDES(axisdata);
+        ptrs = NAD_PTRS(axisdata);
+
+        i = NAD_COORD(axisdata);
+
+        for (istrides = 0; istrides < nstrides; ++istrides) {
+            ptrs[istrides] = dataptr[istrides] + i*strides[istrides];
+        }
+
+        dataptr = ptrs;
+    }
+
+    return NPY_SUCCEED;
+}
+
+npy_intp NpyIter_GetIterIndex(NpyIter *iter)
+{
+    npy_uint32 itflags = NIT_ITFLAGS(iter);
+    npy_intp idim, ndim = NIT_NDIM(iter);
+    npy_intp niter = NIT_NITER(iter);
+
+    NpyIter_AxisData *axisdata;
+    npy_intp sizeof_axisdata;
+    npy_intp iterindex = 0;
+
+    sizeof_axisdata = NIT_SIZEOF_AXISDATA(itflags, ndim, niter);
+    axisdata = NIT_INDEX_AXISDATA(NIT_AXISDATA(iter), ndim-1);
+
+    for (idim = ndim-2; idim >= 0; --idim) {
+        iterindex += NAD_COORD(axisdata);
+        NIT_ADVANCE_AXISDATA(axisdata, -1);
+        iterindex *= NAD_SHAPE(axisdata);
+    }
+    iterindex += NAD_COORD(axisdata);
+
+    return iterindex;
+}
+
 /* SPECIALIZED iternext functions that handle the non-buffering part */
 
 /**begin repeat
diff --git a/numpy/core/src/multiarray/new_iterator.h b/numpy/core/src/multiarray/new_iterator.h
index df0250d6f..61218141b 100644
--- a/numpy/core/src/multiarray/new_iterator.h
+++ b/numpy/core/src/multiarray/new_iterator.h
@@ -27,57 +27,69 @@ typedef enum {
 } NPY_CASTING;
 
 
-/* Allocate a new iterator over one array object */
+/* Allocate a new iterator for one array object */
 NpyIter*
 NpyIter_New(PyArrayObject* op, npy_uint32 flags,
                   NPY_ORDER order, NPY_CASTING casting,
                   PyArray_Descr* dtype,
                   npy_intp a_ndim, npy_intp *axes, npy_intp buffersize);
 
-/* Allocate a new iterator over multiple array objects */
+/* Allocate a new iterator for multiple array objects */
 NpyIter*
 NpyIter_MultiNew(npy_intp niter, PyArrayObject **op_in, npy_uint32 flags,
                  NPY_ORDER order, NPY_CASTING casting,
                  npy_uint32 *op_flags, PyArray_Descr **op_request_dtypes,
                  npy_intp oa_ndim, npy_intp **op_axes, npy_intp buffersize);
 
-/* Removes coords support from an iterator */
-int NpyIter_RemoveCoords(NpyIter *iter);
-/* Removes the inner loop handling (adds NPY_ITER_NO_INNER_ITERATION) */
-int NpyIter_RemoveInnerLoop(NpyIter *iter);
-
 /* Deallocate an iterator */
 int NpyIter_Deallocate(NpyIter* iter);
 
+/* Whether the iterator handles the inner loop */
+int NpyIter_HasInnerLoop(NpyIter *iter);
+/* Removes the inner loop handling (so HasInnerLoop returns false) */
+int NpyIter_RemoveInnerLoop(NpyIter *iter);
+/* Get the array of strides for the inner loop (when HasInnerLoop is false) */
+npy_intp *NpyIter_GetInnerStrideArray(NpyIter *iter);
+/* Get a pointer to the size of the inner loop  (when HasInnerLoop is false) */
+npy_intp* NpyIter_GetInnerLoopSizePtr(NpyIter *iter);
+
 /* Resets the iterator to its initial state */
 void NpyIter_Reset(NpyIter *iter);
 /* Resets the iterator to its initial state, with new base data pointers */
 void NpyIter_ResetBasePointers(NpyIter *iter, char **baseptrs);
-/* Sets the iterator to point at the coordinates in 'coords' */
-int NpyIter_GotoCoords(NpyIter *iter, npy_intp *coords);
-/* Sets the iterator to point at the given index */
-int NpyIter_GotoIndex(NpyIter *iter, npy_intp index);
-
-/* Whether the iterator handles the inner loop */
-int NpyIter_HasInnerLoop(NpyIter *iter);
-/* Whether the iterator is tracking coordinates */
-int NpyIter_HasCoords(NpyIter *iter);
-/* Whether the iterator is tracking an index */
-int NpyIter_HasIndex(NpyIter *iter);
-
-/* Compute a specialized iteration function for an iterator */
-NpyIter_IterNext_Fn NpyIter_GetIterNext(NpyIter *iter);
-/* Compute a specialized getcoords function for an iterator */
-NpyIter_GetCoords_Fn NpyIter_GetGetCoords(NpyIter *iter);
 
 /* Gets the number of dimensions being iterated */
 npy_intp NpyIter_GetNDim(NpyIter *iter);
 /* Gets the number of objects being iterated */
 npy_intp NpyIter_GetNIter(NpyIter *iter);
+
+/* Compute the specialized iteration function for an iterator */
+NpyIter_IterNext_Fn NpyIter_GetIterNext(NpyIter *iter);
+
 /* Gets the number of elements being iterated */
 npy_intp NpyIter_GetIterSize(NpyIter *iter);
+/* Gets the current iteration index */
+npy_intp NpyIter_GetIterIndex(NpyIter *iter);
+/* Sets the iterator to point at the given iteration index */
+int NpyIter_GotoIterIndex(NpyIter *iter, npy_intp iterindex);
+
+/* Whether the iterator is tracking coordinates */
+int NpyIter_HasCoords(NpyIter *iter);
 /* Gets the broadcast shape (if coords are enabled) */
 int NpyIter_GetShape(NpyIter *iter, npy_intp *outshape);
+/* Compute a specialized getcoords function for the iterator */
+NpyIter_GetCoords_Fn NpyIter_GetGetCoords(NpyIter *iter);
+/* Sets the iterator to point at the coordinates in 'coords' */
+int NpyIter_GotoCoords(NpyIter *iter, npy_intp *coords);
+/* Removes coords support from an iterator */
+int NpyIter_RemoveCoords(NpyIter *iter);
+
+/* Whether the iterator is tracking an index */
+int NpyIter_HasIndex(NpyIter *iter);
+/* Get a pointer to the index, if it is being tracked */
+npy_intp *NpyIter_GetIndexPtr(NpyIter *iter);
+/* Sets the iterator to point at the given index */
+int NpyIter_GotoIndex(NpyIter *iter, npy_intp index);
 
 /* Get the array of data pointers (1 per object being iterated) */
 char **NpyIter_GetDataPtrArray(NpyIter *iter);
@@ -87,19 +99,12 @@ PyArray_Descr **NpyIter_GetDescrArray(NpyIter *iter);
 PyArrayObject **NpyIter_GetObjectArray(NpyIter *iter);
 /* Returns a view to the i-th object with the iterator's internal axes */
 PyArrayObject *NpyIter_GetIterView(NpyIter *iter, npy_intp i);
-/* Get a pointer to the index, if it is being tracked */
-npy_intp *NpyIter_GetIndexPtr(NpyIter *iter);
 
 /* Gets an array of read flags (1 per object being iterated) */
 void NpyIter_GetReadFlags(NpyIter *iter, char *outreadflags);
 /* Gets an array of write flags (1 per object being iterated) */
 void NpyIter_GetWriteFlags(NpyIter *iter, char *outwriteflags);
 
-/* Get the array of strides for the inner loop */
-npy_intp *NpyIter_GetInnerStrideArray(NpyIter *iter);
-/* Get a pointer to the size of the inner loop */
-npy_intp* NpyIter_GetInnerLoopSizePtr(NpyIter *iter);
-
 /* For debugging */
 NPY_NO_EXPORT void NpyIter_DebugPrint(NpyIter *iter);
 
diff --git a/numpy/core/src/multiarray/new_iterator_pywrap.c b/numpy/core/src/multiarray/new_iterator_pywrap.c
index 063a4fe86..fc6ddc6ce 100644
--- a/numpy/core/src/multiarray/new_iterator_pywrap.c
+++ b/numpy/core/src/multiarray/new_iterator_pywrap.c
@@ -1288,8 +1288,6 @@ static PyObject *npyiter_itviews_get(NewNpyArrayIterObject *self)
 static PyObject *
 npyiter_next(NewNpyArrayIterObject *self)
 {
-    NewNpyArrayIterObject * tmp;
-
     if (self->iter == NULL || self->finished) {
         return NULL;
     }
@@ -1476,6 +1474,49 @@ static int npyiter_index_set(NewNpyArrayIterObject *self, PyObject *value)
     }
 }
 
+static PyObject *npyiter_iterindex_get(NewNpyArrayIterObject *self)
+{
+    if (self->iter == NULL || self->finished) {
+        PyErr_SetString(PyExc_ValueError,
+                "Iterator is past the end");
+        return NULL;
+    }
+
+    return PyInt_FromLong(NpyIter_GetIterIndex(self->iter));
+}
+
+static int npyiter_iterindex_set(NewNpyArrayIterObject *self, PyObject *value)
+{
+    npy_intp iterindex;
+
+    if (self->iter == NULL) {
+        PyErr_SetString(PyExc_ValueError,
+                "Iterator was not constructed correctly");
+        return -1;
+    }
+
+    if (value == NULL) {
+        PyErr_SetString(PyExc_ValueError,
+                "Cannot delete iterindex");
+        return -1;
+    }
+
+    iterindex = PyInt_AsLong(value);
+    if (iterindex==-1 && PyErr_Occurred()) {
+        return -1;
+    }
+    if (NpyIter_GotoIterIndex(self->iter, iterindex) != NPY_SUCCEED) {
+        return -1;
+    }
+    self->started = 0;
+    self->finished = 0;
+
+    /* If there is nesting, the nested iterators should be reset */
+    npyiter_resetbasepointers(self);
+
+    return 0;
+}
+
 static PyObject *npyiter_hascoords_get(NewNpyArrayIterObject *self)
 {
     if (self->iter == NULL) {
@@ -1787,6 +1828,10 @@ static PyGetSetDef npyiter_getsets[] = {
         (getter)npyiter_index_get,
         (setter)npyiter_index_set,
         NULL, NULL},
+    {"iterindex",
+        (getter)npyiter_iterindex_get,
+        (setter)npyiter_iterindex_set,
+        NULL, NULL},
     {"operands",
         (getter)npyiter_operands_get,
         NULL, NULL, NULL},
diff --git a/numpy/core/tests/test_new_iterator.py b/numpy/core/tests/test_new_iterator.py
index 284b83ee4..6ded7cea0 100644
--- a/numpy/core/tests/test_new_iterator.py
+++ b/numpy/core/tests/test_new_iterator.py
@@ -19,6 +19,13 @@ def iter_indices(i):
         i.iternext()
     return ret
 
+def iter_iterindices(i):
+    ret = []
+    while not i.finished:
+        ret.append(i.iterindex)
+        i.iternext()
+    return ret
+
 def test_iter_refcount():
     # Make sure the iterator doesn't leak
 
@@ -36,6 +43,35 @@ def test_iter_refcount():
     assert_equal(sys.getrefcount(a), rc_a)
     assert_equal(sys.getrefcount(dt), rc_dt)
 
+def test_iter_iterindex():
+    # Make sure iterindex works
+
+    a = arange(6).reshape(2,3)
+    i = newiter(a)
+    assert_equal(iter_iterindices(i), range(6))
+    i.iterindex = 2
+    assert_equal(iter_iterindices(i), range(2,6))
+
+    i = newiter(a, order='F')
+    assert_equal(iter_iterindices(i), range(6))
+    i.iterindex = 3
+    assert_equal(iter_iterindices(i), range(3,6))
+
+    i = newiter(a[::-1], order='F')
+    assert_equal(iter_iterindices(i), range(6))
+    i.iterindex = 4
+    assert_equal(iter_iterindices(i), range(4,6))
+
+    i = newiter(a[::-1,::-1], order='C')
+    assert_equal(iter_iterindices(i), range(6))
+    i.iterindex = 5
+    assert_equal(iter_iterindices(i), range(5,6))
+
+    i = newiter(a[::1,::-1])
+    assert_equal(iter_iterindices(i), range(6))
+    i.iterindex = 1
+    assert_equal(iter_iterindices(i), range(1,6))
+
 def test_iter_best_order():
     # The iterator should always find the iteration order
     # with increasing memory addresses
@@ -611,6 +647,21 @@ def test_iter_flags_errors():
     assert_raises(ValueError, lambda i:i.shape, i)
     # Index available only with an index flag
     assert_raises(ValueError, lambda i:i.index, i)
+    # GotoCoords and GotoIndex incompatible with buffering or no_inner
+    def assign_coords(i):
+        i.coords = (0,)
+    def assign_index(i):
+        i.index = 0
+    def assign_iterindex(i):
+        i.iterindex = 0;
+    i = newiter(arange(6), ['no_inner_iteration'])
+    assert_raises(ValueError, assign_coords, i)
+    assert_raises(ValueError, assign_index, i)
+    assert_raises(ValueError, assign_iterindex, i)
+    i = newiter(arange(6), ['buffered'])
+    assert_raises(ValueError, assign_coords, i)
+    assert_raises(ValueError, assign_index, i)
+    assert_raises(ValueError, assign_iterindex, i)
 
 def test_iter_nbo_align():
     # Check that byte order and alignment changes work