BUG: missingdata: Fix src_itemsize in USE_MASKNA copy to buffer

This was causing some memory corruption that was difficult to track down.

7 files changed, 129 insertions, 22 deletions

diff --git a/numpy/core/src/multiarray/arrayobject.c b/numpy/core/src/multiarray/arrayobject.c
index 9189839b5..17cbd328d 100644
--- a/numpy/core/src/multiarray/arrayobject.c
+++ b/numpy/core/src/multiarray/arrayobject.c
@@ -454,6 +454,80 @@ dump_data(char **string, int *n, int *max_n, char *data, int nd,
 #undef CHECK_MEMORY
 }
 
+NPY_NO_EXPORT void
+PyArray_DebugPrint(PyArrayObject *obj)
+{
+    int i;
+    PyArrayObject_fieldaccess *fobj = (PyArrayObject_fieldaccess *)obj;
+
+    printf("-------------------------------------------------------\n");
+    printf(" Dump of NumPy ndarray at address %p\n", obj);
+    if (obj == NULL) {
+        printf(" It's NULL!\n");
+        printf("-------------------------------------------------------\n");
+        fflush(stdout);
+        return;
+    }
+    printf(" ndim   : %d\n", fobj->nd);
+    printf(" shape  :");
+    for (i = 0; i < fobj->nd; ++i) {
+        printf(" %d", (int)fobj->dimensions[i]);
+    }
+    printf("\n");
+
+    printf(" dtype  : ");
+    PyObject_Print((PyObject *)fobj->descr, stdout, 0);
+    printf("\n");
+    printf(" data   : %p\n", fobj->data);
+    printf(" strides:");
+    for (i = 0; i < fobj->nd; ++i) {
+        printf(" %d", (int)fobj->strides[i]);
+    }
+    printf("\n");
+
+    printf(" base   : %p\n", fobj->base);
+
+    printf(" flags :");
+    if (fobj->flags & NPY_ARRAY_C_CONTIGUOUS)
+        printf(" C_CONTIGUOUS");
+    if (fobj->flags & NPY_ARRAY_F_CONTIGUOUS)
+        printf(" F_CONTIGUOUS");
+    if (fobj->flags & NPY_ARRAY_OWNDATA)
+        printf(" OWNDATA");
+    if (fobj->flags & NPY_ARRAY_ALIGNED)
+        printf(" ALIGNED");
+    if (fobj->flags & NPY_ARRAY_WRITEABLE)
+        printf(" WRITEABLE");
+    if (fobj->flags & NPY_ARRAY_UPDATEIFCOPY)
+        printf(" UPDATEIFCOPY");
+    if (fobj->flags & NPY_ARRAY_MASKNA)
+        printf(" MASKNA");
+    if (fobj->flags & NPY_ARRAY_OWNMASKNA)
+        printf(" OWNMASKNA");
+    printf("\n");
+
+    if (fobj->flags & NPY_ARRAY_MASKNA) {
+        printf(" maskna dtype  : ");
+        PyObject_Print((PyObject *)fobj->maskna_dtype, stdout, 0);
+        printf("\n");
+        printf(" maskna data   : %p\n", fobj->maskna_data);
+        printf(" maskna strides:");
+        for (i = 0; i < fobj->nd; ++i) {
+            printf(" %d", (int)fobj->maskna_strides[i]);
+        }
+        printf("\n");
+    }
+
+    if (fobj->base != NULL && PyArray_Check(fobj->base)) {
+        printf("<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n");
+        printf("Dump of array's BASE:\n");
+        PyArray_DebugPrint((PyArrayObject *)fobj->base);
+        printf(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>\n");
+    }
+    printf("-------------------------------------------------------\n");
+    fflush(stdout);
+}
+
 static PyObject *
 array_repr_builtin(PyArrayObject *self, int repr)
 {
diff --git a/numpy/core/src/multiarray/arrayobject.h b/numpy/core/src/multiarray/arrayobject.h
index ec3361435..b4c4c086c 100644
--- a/numpy/core/src/multiarray/arrayobject.h
+++ b/numpy/core/src/multiarray/arrayobject.h
@@ -12,4 +12,7 @@ _strings_richcompare(PyArrayObject *self, PyArrayObject *other, int cmp_op,
 NPY_NO_EXPORT PyObject *
 array_richcompare(PyArrayObject *self, PyObject *other, int cmp_op);
 
+NPY_NO_EXPORT void
+PyArray_DebugPrint(PyArrayObject *obj);
+
 #endif
diff --git a/numpy/core/src/multiarray/convert_datatype.c b/numpy/core/src/multiarray/convert_datatype.c
index fa0d6b4bb..0d72807ae 100644
--- a/numpy/core/src/multiarray/convert_datatype.c
+++ b/numpy/core/src/multiarray/convert_datatype.c
@@ -707,14 +707,17 @@ PyArray_CanCastArrayTo(PyArrayObject *arr, PyArray_Descr *to,
 {
     PyArray_Descr *from = PyArray_DESCR(arr);
 
-    /* If it's not a scalar, use the standard rules */
-    if (PyArray_NDIM(arr) > 0) {
-        return PyArray_CanCastTypeTo(from, to, casting);
-    }
-    /* Otherwise, check the value */
-    else {
-        return can_cast_scalar_to(from, PyArray_DATA(arr), to, casting);
+    /* If it's a scalar, check the value */
+    if (PyArray_NDIM(arr) == 0 && !PyArray_HASFIELDS(arr)) {
+        /* Only check the value if it's not masked */
+        if (!PyArray_HASMASKNA(arr) ||
+                NpyMaskValue_IsExposed((npy_mask)*PyArray_MASKNA_DATA(arr))) {
+            return can_cast_scalar_to(from, PyArray_DATA(arr), to, casting);
+        }
     }
+
+    /* Otherwise, use the standard rules */
+    return PyArray_CanCastTypeTo(from, to, casting);
 }
 
 /*NUMPY_API
diff --git a/numpy/core/src/multiarray/nditer_api.c b/numpy/core/src/multiarray/nditer_api.c
index 0c925d8b9..0dccd6148 100644
--- a/numpy/core/src/multiarray/nditer_api.c
+++ b/numpy/core/src/multiarray/nditer_api.c
@@ -1591,6 +1591,7 @@ NpyIter_DebugPrint(NpyIter *iter)
     }
 
     printf("------- END ITERATOR DUMP -------\n");
+    fflush(stdout);
 
     PyGILState_Release(gilstate);
 }
@@ -2044,6 +2045,7 @@ npyiter_copy_to_buffers(NpyIter *iter, char **prev_dataptrs)
     npy_uint32 itflags = NIT_ITFLAGS(iter);
     int ndim = NIT_NDIM(iter);
     int iop, nop = NIT_NOP(iter);
+    int first_maskna_op = NIT_FIRST_MASKNA_OP(iter);
 
     char *op_itflags = NIT_OPITFLAGS(iter);
     NpyIter_BufferData *bufferdata = NIT_BUFFERDATA(iter);
@@ -2373,7 +2375,7 @@ npyiter_copy_to_buffers(NpyIter *iter, char **prev_dataptrs)
         }
 
         if (stransfer != NULL) {
-            npy_intp src_itemsize = PyArray_DESCR(operands[iop])->elsize;
+            npy_intp src_itemsize;
             npy_intp op_transfersize;
 
             npy_intp dst_stride, *src_strides, *src_coords, *src_shape;
@@ -2381,6 +2383,14 @@ npyiter_copy_to_buffers(NpyIter *iter, char **prev_dataptrs)
 
             npy_bool skip_transfer = 0;
 
+            /* Need to pick the right item size for the data vs mask */
+            if (iop < first_maskna_op) {
+                src_itemsize = PyArray_DTYPE(operands[iop])->elsize;
+            }
+            else {
+                src_itemsize = PyArray_MASKNA_DTYPE(operands[iop])->elsize;
+            }
+
             /* If stransfer wasn't set to NULL, buffering is required */
             any_buffered = 1;
 
diff --git a/numpy/core/src/multiarray/nditer_constr.c b/numpy/core/src/multiarray/nditer_constr.c
index a7f3ab30e..07efba65c 100644
--- a/numpy/core/src/multiarray/nditer_constr.c
+++ b/numpy/core/src/multiarray/nditer_constr.c
@@ -14,6 +14,8 @@
 #define NPY_ITERATOR_IMPLEMENTATION_CODE
 #include "nditer_impl.h"
 
+#include "arrayobject.h"
+
 /* Internal helper functions private to this file */
 static int
 npyiter_check_global_flags(npy_uint32 flags, npy_uint32* itflags);
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 09360a8d2..aafae83ed 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -2705,13 +2705,13 @@ conform_reduce_result(int ndim, npy_bool *axis_flags, PyArrayObject *out)
 /* Allocate 'result' or conform 'out' to 'self' (in 'result') */
 static PyArrayObject *
 allocate_or_conform_reduce_result(PyArrayObject *arr, PyArrayObject *out,
-                        npy_bool *axis_flags, PyArray_Descr * otype_dtype,
+                        npy_bool *axis_flags, PyArray_Descr *otype_dtype,
                         int addmask)
 {
     PyArrayObject *result;
 
     if (out == NULL) {
-        Py_INCREF(otype_dtype);
+        Py_XINCREF(otype_dtype);
         result = allocate_reduce_result(arr, axis_flags, otype_dtype);
 
         /* Allocate an NA mask if necessary */
@@ -3075,6 +3075,9 @@ PyUFunc_Reduce(PyUFuncObject *self, PyArrayObject *arr, PyArrayObject *out,
                     result_strides[idim] = PyArray_MASKNA_STRIDES(result)[idim];
                 }
             }
+            if (!PyArray_HASMASKNA(arr) || !PyArray_HASMASKNA(result))
+                printf ("hasmaskna %d %d\n", PyArray_HASMASKNA(arr),
+                                PyArray_HASMASKNA(result));
             if (PyArray_ReduceMaskNAArray(ndim, PyArray_DIMS(arr),
                         PyArray_MASKNA_DTYPE(arr),
                         PyArray_MASKNA_DATA(arr),
@@ -4183,12 +4186,9 @@ PyUFunc_GenericReduction(PyUFuncObject *self, PyObject *args,
     }
     /* Try to interpret axis as an integer */
     else {
-        /* TODO: PyNumber_Index would be good to use here */
         long axis = PyInt_AsLong(axes_in);
+        /* TODO: PyNumber_Index would be good to use here */
         if (axis == -1 && PyErr_Occurred()) {
-            PyErr_Clear();
-            PyErr_SetString(PyExc_ValueError,
-                    "invalid argument for 'axis'");
             Py_XDECREF(otype);
             Py_DECREF(mp);
             return NULL;
diff --git a/numpy/core/tests/test_maskna.py b/numpy/core/tests/test_maskna.py
index 3c188de21..0e5cc33bc 100644
--- a/numpy/core/tests/test_maskna.py
+++ b/numpy/core/tests/test_maskna.py
@@ -826,26 +826,41 @@ def test_array_maskna_sum_prod_methods():
     # ndarray.sum, ndarray.prod
     a = np.array([[2, np.NA, 10],
                   [4, 8, 7],
-                  [12, 4, np.NA]], maskna=True)
+                  [12, 4, np.NA],
+                  [3, 2, 5]], maskna=True)
 
     res = a.sum(axis=0)
     assert_equal(np.isna(res), [0,1,1])
-    assert_equal(res[~np.isna(res)], [18])
+    assert_equal(res[~np.isna(res)], [21])
 
     res = a.sum(axis=-1)
-    assert_equal(np.isna(res), [1,0,1])
-    assert_equal(res[~np.isna(res)], [19])
+    assert_equal(np.isna(res), [1,0,1,0])
+    assert_equal(res[~np.isna(res)], [19,10])
 
     res = a.prod(axis=0)
     assert_equal(np.isna(res), [0,1,1])
-    assert_equal(res[~np.isna(res)], [2*4*12])
+    assert_equal(res[~np.isna(res)], [2*4*12*3])
 
     res = a.prod(axis=-1)
-    assert_equal(np.isna(res), [1,0,1])
-    assert_equal(res[~np.isna(res)], [4*8*7])
+    assert_equal(np.isna(res), [1,0,1,0])
+    assert_equal(res[~np.isna(res)], [4*8*7,3*2*5])
+
+    # Check also with Fortran-order
+    a = np.array([[2, np.NA, 10],
+                  [4, 8, 7],
+                  [12, 4, np.NA],
+                  [3, 2, 5]], maskna=True, order='F')
+
+    res = a.sum(axis=0)
+    assert_equal(np.isna(res), [0,1,1])
+    assert_equal(res[~np.isna(res)], [21])
+
+    res = a.sum(axis=-1)
+    assert_equal(np.isna(res), [1,0,1,0])
+    assert_equal(res[~np.isna(res)], [19,10])
+
 
 def test_array_maskna_std_mean_methods():
-    return
     # ndarray.std, ndarray.mean
     a = np.array([[2, np.NA, 10],
                   [4, 8, 7],