ENH: missingdata: Move getting the masked ufunc inner loop to a different function

This is so that backwards compatibility can be retained with denying
regular ufuncs from supporting dynamic auxiliary data, but allowing
it for the masked inner loops so that the default masked loop can store the
regular inner loop it proxies calls to.

4 files changed, 35 insertions, 39 deletions

diff --git a/numpy/core/include/numpy/ufuncobject.h b/numpy/core/include/numpy/ufuncobject.h
index 074fdd9eb..adf5c7dae 100644
--- a/numpy/core/include/numpy/ufuncobject.h
+++ b/numpy/core/include/numpy/ufuncobject.h
@@ -31,11 +31,14 @@ struct _tagPyUFuncObject;
  * This function should validate that the casting rule is being followed,
  * and fail if it is not.
  *
+ * For backwards compatibility, the regular type resolution function does not
+ * support auxiliary data with object semantics. The type resolution call
+ * which returns a masked generic function returns a standard NpyAuxData
+ * object, for which the NPY_AUXDATA_FREE and NPY_AUXDATA_CLONE macros
+ * work.
+ *
  * ufunc:             The ufunc object.
  * casting:           The 'casting' parameter provided to the ufunc.
- * masked:            If non-zero, must return a
- *                    PyUFuncGenericMaskedFunction instead of a regular
- *                    PyUFuncGenericFunction.
  * operands:          An array of length (ufunc->nin + ufunc->nout),
  *                    with the output parameters possibly NULL.
  * type_tup:          Either NULL, or the type_tup passed to the ufunc.
@@ -54,12 +57,19 @@ struct _tagPyUFuncObject;
 typedef int (PyUFunc_TypeResolutionFunc)(
                                 struct _tagPyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
                                 PyUFuncGenericFunction *out_innerloop,
                                 void **out_innerloopdata);
+typedef int (PyUFunc_TypeResolutionMaskedFunc)(
+                                struct _tagPyUFuncObject *ufunc,
+                                NPY_CASTING casting,
+                                PyArrayObject **operands,
+                                PyObject *type_tup,
+                                PyArray_Descr **out_dtypes,
+                                PyUFuncGenericMaskedFunction *out_innerloop,
+                                NpyAuxData **out_innerloopdata);
 
 typedef struct _tagPyUFuncObject {
         PyObject_HEAD
@@ -129,6 +139,12 @@ typedef struct _tagPyUFuncObject {
          * have a different set of rules.
          */
         PyUFunc_TypeResolutionFunc *type_resolution_function;
+        /*
+         * A function which resolves the types and returns an inner loop.
+         * This is used by the regular ufunc when it requires using
+         * a mask to select which elements to compute.
+         */
+        PyUFunc_TypeResolutionMaskedFunc *type_resolution_masked_function;
 } PyUFuncObject;
 
 #include "arrayobject.h"
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 98ddbcf84..de541af42 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -1733,7 +1733,6 @@ ensure_dtype_nbo(PyArray_Descr *type)
 NPY_NO_EXPORT int
 PyUFunc_DefaultTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -1789,7 +1788,6 @@ PyUFunc_DefaultTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_SimpleBinaryComparisonTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -1817,7 +1815,7 @@ PyUFunc_SimpleBinaryComparisonTypeResolution(PyUFuncObject *ufunc,
     type_num2 = PyArray_DESCR(operands[1])->type_num;
     if (type_num1 >= NPY_NTYPES || type_num2 >= NPY_NTYPES ||
             type_num1 == NPY_OBJECT || type_num2 == NPY_OBJECT) {
-        return PyUFunc_DefaultTypeResolution(ufunc, casting, masked, operands,
+        return PyUFunc_DefaultTypeResolution(ufunc, casting, operands,
                 type_tup, out_dtypes, out_innerloop, out_innerloopdata);
     }
 
@@ -1836,7 +1834,7 @@ PyUFunc_SimpleBinaryComparisonTypeResolution(PyUFuncObject *ufunc,
          * default type resolution handle this one.
          */
         if (!PyTuple_Check(type_tup) || PyTuple_GET_SIZE(type_tup) != 1) {
-            return PyUFunc_DefaultTypeResolution(ufunc, casting, masked,
+            return PyUFunc_DefaultTypeResolution(ufunc, casting,
                     operands, type_tup, out_dtypes,
                     out_innerloop, out_innerloopdata);
         }
@@ -1916,7 +1914,6 @@ PyUFunc_SimpleBinaryComparisonTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_SimpleUnaryOperationTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -1942,7 +1939,7 @@ PyUFunc_SimpleUnaryOperationTypeResolution(PyUFuncObject *ufunc,
      */
     type_num1 = PyArray_DESCR(operands[0])->type_num;
     if (type_num1 >= NPY_NTYPES || type_num1 == NPY_OBJECT) {
-        return PyUFunc_DefaultTypeResolution(ufunc, casting, masked, operands,
+        return PyUFunc_DefaultTypeResolution(ufunc, casting, operands,
                 type_tup, out_dtypes, out_innerloop, out_innerloopdata);
     }
 
@@ -1961,7 +1958,7 @@ PyUFunc_SimpleUnaryOperationTypeResolution(PyUFuncObject *ufunc,
          * default type resolution handle this one.
          */
         if (!PyTuple_Check(type_tup) || PyTuple_GET_SIZE(type_tup) != 1) {
-            return PyUFunc_DefaultTypeResolution(ufunc, casting, masked,
+            return PyUFunc_DefaultTypeResolution(ufunc, casting,
                     operands, type_tup, out_dtypes,
                     out_innerloop, out_innerloopdata);
         }
@@ -2025,7 +2022,6 @@ PyUFunc_SimpleUnaryOperationTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_OnesLikeTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING NPY_UNUSED(casting),
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -2033,7 +2029,7 @@ PyUFunc_OnesLikeTypeResolution(PyUFuncObject *ufunc,
                                 void **out_innerloopdata)
 {
     return PyUFunc_SimpleUnaryOperationTypeResolution(ufunc,
-                        NPY_UNSAFE_CASTING, masked,
+                        NPY_UNSAFE_CASTING,
                         operands, type_tup, out_dtypes,
                         out_innerloop, out_innerloopdata);
 }
@@ -2054,7 +2050,6 @@ PyUFunc_OnesLikeTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_SimpleBinaryOperationTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -2082,7 +2077,7 @@ PyUFunc_SimpleBinaryOperationTypeResolution(PyUFuncObject *ufunc,
     type_num2 = PyArray_DESCR(operands[1])->type_num;
     if (type_num1 >= NPY_NTYPES || type_num2 >= NPY_NTYPES ||
             type_num1 == NPY_OBJECT || type_num2 == NPY_OBJECT) {
-        return PyUFunc_DefaultTypeResolution(ufunc, casting, masked, operands,
+        return PyUFunc_DefaultTypeResolution(ufunc, casting, operands,
                 type_tup, out_dtypes, out_innerloop, out_innerloopdata);
     }
 
@@ -2103,7 +2098,7 @@ PyUFunc_SimpleBinaryOperationTypeResolution(PyUFuncObject *ufunc,
          * default type resolution handle this one.
          */
         if (!PyTuple_Check(type_tup) || PyTuple_GET_SIZE(type_tup) != 1) {
-            return PyUFunc_DefaultTypeResolution(ufunc, casting, masked,
+            return PyUFunc_DefaultTypeResolution(ufunc, casting,
                     operands, type_tup, out_dtypes,
                     out_innerloop, out_innerloopdata);
         }
@@ -2171,7 +2166,6 @@ PyUFunc_SimpleBinaryOperationTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_AbsoluteTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -2180,12 +2174,12 @@ PyUFunc_AbsoluteTypeResolution(PyUFuncObject *ufunc,
 {
     /* Use the default for complex types, to find the loop producing float */
     if (PyTypeNum_ISCOMPLEX(PyArray_DESCR(operands[0])->type_num)) {
-        return PyUFunc_DefaultTypeResolution(ufunc, casting, masked, operands,
+        return PyUFunc_DefaultTypeResolution(ufunc, casting, operands,
                     type_tup, out_dtypes, out_innerloop, out_innerloopdata);
     }
     else {
         return PyUFunc_SimpleUnaryOperationTypeResolution(ufunc, casting,
-                    masked, operands, type_tup, out_dtypes,
+                    operands, type_tup, out_dtypes,
                     out_innerloop, out_innerloopdata);
     }
 }
@@ -2279,7 +2273,6 @@ timedelta_dtype_with_copied_meta(PyArray_Descr *dtype)
 NPY_NO_EXPORT int
 PyUFunc_AdditionTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -2298,7 +2291,7 @@ PyUFunc_AdditionTypeResolution(PyUFuncObject *ufunc,
 
     /* Use the default when datetime and timedelta are not involved */
     if (!PyTypeNum_ISDATETIME(type_num1) && !PyTypeNum_ISDATETIME(type_num2)) {
-        return PyUFunc_DefaultTypeResolution(ufunc, casting, masked, operands,
+        return PyUFunc_DefaultTypeResolution(ufunc, casting, operands,
                     type_tup, out_dtypes, out_innerloop, out_innerloopdata);
     }
 
@@ -2487,7 +2480,6 @@ type_reso_error: {
 NPY_NO_EXPORT int
 PyUFunc_SubtractionTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -2506,7 +2498,7 @@ PyUFunc_SubtractionTypeResolution(PyUFuncObject *ufunc,
 
     /* Use the default when datetime and timedelta are not involved */
     if (!PyTypeNum_ISDATETIME(type_num1) && !PyTypeNum_ISDATETIME(type_num2)) {
-        return PyUFunc_DefaultTypeResolution(ufunc, casting, masked, operands,
+        return PyUFunc_DefaultTypeResolution(ufunc, casting, operands,
                     type_tup, out_dtypes, out_innerloop, out_innerloopdata);
     }
 
@@ -2673,7 +2665,6 @@ type_reso_error: {
 NPY_NO_EXPORT int
 PyUFunc_MultiplicationTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -2692,7 +2683,7 @@ PyUFunc_MultiplicationTypeResolution(PyUFuncObject *ufunc,
 
     /* Use the default when datetime and timedelta are not involved */
     if (!PyTypeNum_ISDATETIME(type_num1) && !PyTypeNum_ISDATETIME(type_num2)) {
-        return PyUFunc_DefaultTypeResolution(ufunc, casting, masked, operands,
+        return PyUFunc_DefaultTypeResolution(ufunc, casting, operands,
                     type_tup, out_dtypes, out_innerloop, out_innerloopdata);
     }
 
@@ -2834,7 +2825,6 @@ type_reso_error: {
 NPY_NO_EXPORT int
 PyUFunc_DivisionTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -2853,7 +2843,7 @@ PyUFunc_DivisionTypeResolution(PyUFuncObject *ufunc,
 
     /* Use the default when datetime and timedelta are not involved */
     if (!PyTypeNum_ISDATETIME(type_num1) && !PyTypeNum_ISDATETIME(type_num2)) {
-        return PyUFunc_DefaultTypeResolution(ufunc, casting, masked, operands,
+        return PyUFunc_DefaultTypeResolution(ufunc, casting, operands,
                     type_tup, out_dtypes, out_innerloop, out_innerloopdata);
     }
 
@@ -3607,7 +3597,7 @@ PyUFunc_GeneralizedFunction(PyUFuncObject *self,
     NPY_UF_DBG_PRINT("Finding inner loop\n");
 
 
-    retval = self->type_resolution_function(self, casting, 0,
+    retval = self->type_resolution_function(self, casting,
                         op, type_tup, dtype, &innerloop, &innerloopdata);
     if (retval < 0) {
         goto fail;
@@ -3933,7 +3923,7 @@ PyUFunc_GenericFunction(PyUFuncObject *self,
 
     NPY_UF_DBG_PRINT("Finding inner loop\n");
 
-    retval = self->type_resolution_function(self, casting, 0,
+    retval = self->type_resolution_function(self, casting,
                         op, type_tup, dtype, &innerloop, &innerloopdata);
     if (retval < 0) {
         goto fail;
diff --git a/numpy/core/src/umath/ufunc_object.h b/numpy/core/src/umath/ufunc_object.h
index 7a851ed77..59754380c 100644
--- a/numpy/core/src/umath/ufunc_object.h
+++ b/numpy/core/src/umath/ufunc_object.h
@@ -10,7 +10,6 @@ ufunc_seterr(PyObject *NPY_UNUSED(dummy), PyObject *args);
 NPY_NO_EXPORT int
 PyUFunc_SimpleBinaryComparisonTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -20,7 +19,6 @@ PyUFunc_SimpleBinaryComparisonTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_SimpleUnaryOperationTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -30,7 +28,6 @@ PyUFunc_SimpleUnaryOperationTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_OnesLikeTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -40,7 +37,6 @@ PyUFunc_OnesLikeTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_SimpleBinaryOperationTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -50,7 +46,6 @@ PyUFunc_SimpleBinaryOperationTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_AbsoluteTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -60,7 +55,6 @@ PyUFunc_AbsoluteTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_AdditionTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -70,7 +64,6 @@ PyUFunc_AdditionTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_SubtractionTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -80,7 +73,6 @@ PyUFunc_SubtractionTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_MultiplicationTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
@@ -89,7 +81,6 @@ PyUFunc_MultiplicationTypeResolution(PyUFuncObject *ufunc,
 NPY_NO_EXPORT int
 PyUFunc_DivisionTypeResolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,
diff --git a/numpy/core/src/umath/umathmodule.c.src b/numpy/core/src/umath/umathmodule.c.src
index ff72187bb..82afd7fd8 100644
--- a/numpy/core/src/umath/umathmodule.c.src
+++ b/numpy/core/src/umath/umathmodule.c.src
@@ -47,7 +47,6 @@ static PyUFuncGenericFunction pyfunc_functions[] = {PyUFunc_On_Om};
 static int
 object_ufunc_type_resolution(PyUFuncObject *ufunc,
                                 NPY_CASTING casting,
-                                int masked,
                                 PyArrayObject **operands,
                                 PyObject *type_tup,
                                 PyArray_Descr **out_dtypes,