Fix constant and zero/one padding when stacking neighborhood iterator.

Stacking a neighborhood iterator on top of another one now works as expected.

2 files changed, 66 insertions, 2 deletions

diff --git a/numpy/core/include/numpy/_neighborhood_iterator_imp.h b/numpy/core/include/numpy/_neighborhood_iterator_imp.h
index bd1ebd36d..37a03ebd1 100644
--- a/numpy/core/include/numpy/_neighborhood_iterator_imp.h
+++ b/numpy/core/include/numpy/_neighborhood_iterator_imp.h
@@ -332,14 +332,39 @@ int PyArrayNeighborhoodIter_NextCircular(PyArrayNeighborhoodIterObject* iter)
     return 0;
 }
 
-static NPY_INLINE int PyArrayNeighborhoodIter_Next(PyArrayNeighborhoodIterObject* iter)
+#define _INF_SET_PTR(c) \
+    bd = niter->coordinates[c] + p->coordinates[c]; \
+    if (bd < p->bounds[c][0] || bd > p->bounds[c][1]) { \
+        return niter->constant; \
+    } \
+    _coordinates[c] = bd;
+
+static inline char*
+_neigh_iter_get_ptr_const(PyArrayNeighborhoodIterObject* niter)
+{
+    int i;
+    npy_intp bd;
+    PyArrayIterObject *p = niter->_internal_iter;
+    npy_intp _coordinates[NPY_MAXDIMS];
+
+    for(i = 0; i < niter->nd; ++i) {
+        _INF_SET_PTR(i)
+    }
+
+    return p->translate(p, _coordinates);
+}
+#undef _INF_SET_PTR
+
+static NPY_INLINE int
+PyArrayNeighborhoodIter_Next(PyArrayNeighborhoodIterObject* iter)
 {
     _PyArrayNeighborhoodIter_IncrCoord (iter);
     switch (iter->mode) {
         case NPY_NEIGHBORHOOD_ITER_ZERO_PADDING:
         case NPY_NEIGHBORHOOD_ITER_ONE_PADDING:
         case NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING:
-            _PyArrayNeighborhoodIter_SetPtrConstant(iter);
+            //iter->dataptr = _neigh_iter_get_ptr_const(iter);
+            iter->dataptr = iter->translate((PyArrayIterObject*)iter, iter->coordinates);
             break;
         case NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING:
             _PyArrayNeighborhoodIter_SetPtrMirror(iter);
@@ -410,6 +435,9 @@ PyArrayNeighborhoodIter_Reset(PyArrayNeighborhoodIterObject* iter)
     for (i = 0; i < iter->nd; ++i) {
         iter->coordinates[i] = iter->bounds[i][0];
     }
+    iter->dataptr = _neigh_iter_get_ptr_const(iter);
+
+#if 0
     switch (iter->mode) {
         case NPY_NEIGHBORHOOD_ITER_ZERO_PADDING:
         case NPY_NEIGHBORHOOD_ITER_ONE_PADDING:
@@ -423,6 +451,7 @@ PyArrayNeighborhoodIter_Reset(PyArrayNeighborhoodIterObject* iter)
             _PyArrayNeighborhoodIter_SetPtrCircular(iter);
             break;
     }
+#endif
 
     return 0;
 }
diff --git a/numpy/core/src/multiarray/iterators.c b/numpy/core/src/multiarray/iterators.c
index fc70d6826..7cdb199ab 100644
--- a/numpy/core/src/multiarray/iterators.c
+++ b/numpy/core/src/multiarray/iterators.c
@@ -1793,6 +1793,31 @@ static char* _set_constant(PyArrayNeighborhoodIterObject* iter,
     return ret;
 }
 
+#define _INF_SET_PTR(c) \
+    bd = coordinates[c] + p->coordinates[c]; \
+    if (bd < p->bounds[c][0] || bd > p->bounds[c][1]) { \
+        return niter->constant; \
+    } \
+    _coordinates[c] = bd;
+
+/* set the dataptr from its current coordinates */
+static char*
+get_ptr_constant(PyArrayIterObject* _iter, npy_intp *coordinates)
+{
+    int i;
+    npy_intp bd, _coordinates[NPY_MAXDIMS];
+    PyArrayNeighborhoodIterObject *niter = (PyArrayNeighborhoodIterObject*)_iter;
+    PyArrayIterObject *p = niter->_internal_iter;
+
+    for(i = 0; i < niter->nd; ++i) {
+        _INF_SET_PTR(i)
+    }
+
+    return p->translate(p, _coordinates);
+}
+#undef _INF_SET_PTR
+
+
 /*
  * fill and x->ao should have equivalent types 
  */
@@ -1832,10 +1857,12 @@ PyArray_NeighborhoodIterNew(PyArrayIterObject *x, intp *bounds,
         case NPY_NEIGHBORHOOD_ITER_ZERO_PADDING:
             ret->constant = PyArray_Zero(x->ao);
             ret->mode = mode;
+            ret->translate = &get_ptr_constant;
             break;
         case NPY_NEIGHBORHOOD_ITER_ONE_PADDING:
             ret->constant = PyArray_One(x->ao);
             ret->mode = mode;
+            ret->translate = &get_ptr_constant;
             break;
         case NPY_NEIGHBORHOOD_ITER_CONSTANT_PADDING:
             /* New reference in returned value of _set_constant if array
@@ -1846,12 +1873,20 @@ PyArray_NeighborhoodIterNew(PyArrayIterObject *x, intp *bounds,
                 goto clean_x;
             }
             ret->mode = mode;
+            ret->translate = &get_ptr_constant;
             break;
+#if 0
         case NPY_NEIGHBORHOOD_ITER_MIRROR_PADDING:
+            ret->mode = mode;
+            ret->constant = NULL;
+            ret->translate = get_coordinates_mirror;
+            break;
         case NPY_NEIGHBORHOOD_ITER_CIRCULAR_PADDING:
             ret->mode = mode;
             ret->constant = NULL;
+            ret->translate = get_coordinates_circular;
             break;
+#endif
         default:
             PyErr_SetString(PyExc_ValueError, "Unsupported padding mode");
             goto clean_x;